From 038c40667479983513c3cf3bb114fd08c6494edf Mon Sep 17 00:00:00 2001
From: aminst <aminst79@gmail.com>
Date: Mon, 29 Jan 2024 23:26:34 -0500
Subject: [PATCH] Fix storage WriteBucket bug

---
 pkg/storage/storage.go      | 10 +++++++---
 pkg/storage/storage_test.go | 20 ++++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/pkg/storage/storage.go b/pkg/storage/storage.go
index 71fe0f7..70228e0 100644
--- a/pkg/storage/storage.go
+++ b/pkg/storage/storage.go
@@ -230,6 +230,11 @@ func (s *StorageHandler) BatchWriteBucket(storageID int, readBucketBlocksList ma
 	ctx := context.Background()
 	dataResults := make(map[int]*redis.BoolCmd)
 	metadataResults := make(map[int]*redis.BoolCmd)
+	writtenBlocks = make(map[string]string)
+
+	log.Debug().Msgf("buckets from readBucketBlocksList: %v", readBucketBlocksList)
+	log.Debug().Msgf("shardNodeBlocks: %v", shardNodeBlocks)
+
 	for bucketID, readBucketBlocks := range readBucketBlocksList {
 		values := make([]string, s.Z+s.S)
 		metadatas := make([]string, s.Z+s.S)
@@ -239,13 +244,12 @@ func (s *StorageHandler) BatchWriteBucket(storageID int, readBucketBlocksList ma
 			realIndex[k] = k
 		}
 		shuffleArray(realIndex)
-		writtenBlocks = make(map[string]string)
 		i := 0
 		for key, value := range readBucketBlocks {
 			if strings.HasPrefix(key, "dummy") {
 				continue
 			}
-			if len(writtenBlocks) < s.Z {
+			if i < s.Z {
 				writtenBlocks[key] = value
 				values[realIndex[i]], err = Encrypt(value, s.key)
 				if err != nil {
@@ -262,7 +266,7 @@ func (s *StorageHandler) BatchWriteBucket(storageID int, readBucketBlocksList ma
 			if strings.HasPrefix(key, "dummy") {
 				continue
 			}
-			if len(writtenBlocks) < s.Z {
+			if i < s.Z {
 				writtenBlocks[key] = value
 				values[realIndex[i]], err = Encrypt(value, s.key)
 				if err != nil {
diff --git a/pkg/storage/storage_test.go b/pkg/storage/storage_test.go
index bd30a02..0897881 100644
--- a/pkg/storage/storage_test.go
+++ b/pkg/storage/storage_test.go
@@ -148,6 +148,26 @@ func TestBatchGetBlockOffset(t *testing.T) {
 	}
 }
 
+func TestBatchReadBucketReturnsBlocksInAllBuckets(t *testing.T) {
+	s := NewStorageHandler(4, 1, 9, 1, []config.RedisEndpoint{{ID: 0, IP: "localhost", Port: 6379}})
+	s.InitDatabase()
+	toWriteBlocks := map[int]map[string]string{1: {"usr1": "value1"}, 2: {"usr2": "value2"}, 3: {"usr3": "value3"}, 4: {"usr4": "value4"}, 5: {"usr5": "value5"}}
+	s.BatchWriteBucket(0, toWriteBlocks, map[string]string{})
+	blocks, err := s.BatchReadBucket([]int{1, 2, 3, 4, 5}, 0)
+	if err != nil {
+		t.Errorf("error reading bucket")
+	}
+	expectedReadBuckets := toWriteBlocks
+	log.Debug().Msgf("blocks: %v", expectedReadBuckets)
+	for bucketID, blockToVal := range expectedReadBuckets {
+		for block, val := range blockToVal {
+			if blocks[bucketID][block] != val {
+				t.Errorf("expected %s, but got %s", val, blocks[bucketID][block])
+			}
+		}
+	}
+}
+
 func TestRandom(t *testing.T) {
 	s := NewStorageHandler(4, 1, 9, 1, []config.RedisEndpoint{{ID: 0, IP: "localhost", Port: 6379}})
 	s.InitDatabase()