From ec54a1ad3b5476c1907e3f35237ad28418074bfa Mon Sep 17 00:00:00 2001
From: Ankita Luthra <lankita@google.com>
Date: Mon, 29 Jul 2024 14:07:17 +0530
Subject: [PATCH] Minor changes for stat cache for HNS (#2229)

* refactor if condition in lookup folder for stat cache

* adds folder in stat cache when get folder is success

* adds comment
---
 internal/cache/metadata/stat_cache.go           | 16 ++++++++--------
 internal/storage/caching/fast_stat_bucket.go    | 10 +++++++++-
 .../storage/caching/fast_stat_bucket_test.go    | 17 +++++++++++++++++
 3 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/internal/cache/metadata/stat_cache.go b/internal/cache/metadata/stat_cache.go
index 000eef5ad2..029440ca5e 100644
--- a/internal/cache/metadata/stat_cache.go
+++ b/internal/cache/metadata/stat_cache.go
@@ -248,16 +248,16 @@ func (sc *statCacheBucketView) LookUpFolder(
 	// Look up in the LRU cache.
 	hit, entry := sc.sharedCacheLookup(folderName, now)
 
-	if hit {
-		// Adds scenario to check folder as well even if object with same name is already present
-		if entry.f == nil {
-			return false, nil
-		}
-
-		return true, entry.f
+	if !hit {
+		return false, nil
+	}
+	// Adds scenario to check folder as well even if object with same name is already present
+	// TODO: should be removed once integration for lookup is completed
+	if entry.f == nil {
+		return false, nil
 	}
 
-	return false, nil
+	return true, entry.f
 }
 
 func (sc *statCacheBucketView) sharedCacheLookup(key string, now time.Time) (bool, *entry) {
diff --git a/internal/storage/caching/fast_stat_bucket.go b/internal/storage/caching/fast_stat_bucket.go
index 01264a3c94..6a8b55a81d 100644
--- a/internal/storage/caching/fast_stat_bucket.go
+++ b/internal/storage/caching/fast_stat_bucket.go
@@ -335,7 +335,15 @@ func (b *fastStatBucket) GetFolder(
 	}
 
 	// Fetch the Folder
-	return b.wrapped.GetFolder(ctx, prefix)
+	folder, error := b.wrapped.GetFolder(ctx, prefix)
+
+	if error != nil {
+		return nil, error
+	}
+
+	// Record the new folder.
+	b.cache.InsertFolder(folder, b.clock.Now().Add(b.ttl))
+	return folder, nil
 }
 
 func (b *fastStatBucket) CreateFolder(ctx context.Context, folderName string) (f *gcs.Folder, err error) {
diff --git a/internal/storage/caching/fast_stat_bucket_test.go b/internal/storage/caching/fast_stat_bucket_test.go
index fe070b340b..caa7b79057 100644
--- a/internal/storage/caching/fast_stat_bucket_test.go
+++ b/internal/storage/caching/fast_stat_bucket_test.go
@@ -761,6 +761,8 @@ func (t *StatObjectTest) TestShouldCallGetFolderWhenEntryIsNotPresent() {
 
 	ExpectCall(t.cache, "LookUpFolder")(name, Any()).
 		WillOnce(Return(false, nil))
+	ExpectCall(t.cache, "InsertFolder")(folder, Any()).
+		WillOnce(Return())
 	ExpectCall(t.wrapped, "GetFolder")(Any(), name).
 		WillOnce(Return(folder, nil))
 
@@ -770,6 +772,21 @@ func (t *StatObjectTest) TestShouldCallGetFolderWhenEntryIsNotPresent() {
 	ExpectThat(result, Pointee(DeepEquals(*folder)))
 }
 
+func (t *StatObjectTest) TestShouldReturnNilWhenErrorIsReturnedFromGetFolder() {
+	const name = "some-name"
+	error := errors.New("connection error")
+
+	ExpectCall(t.cache, "LookUpFolder")(name, Any()).
+		WillOnce(Return(false, nil))
+	ExpectCall(t.wrapped, "GetFolder")(Any(), name).
+		WillOnce(Return(nil, error))
+
+	folder, result := t.bucket.GetFolder(context.TODO(), name)
+
+	AssertEq(nil, folder)
+	AssertEq(error, result)
+}
+
 func (t *StatObjectTest) TestRenameFolder() {
 	const name = "some-name"
 	const newName = "new-name"