fix: adding lock while discovering partitions, Closes #412 (#413)

Signed-off-by: Yashash H L <yashashhl25@gmail.com>
numaproj · Dec 7, 2022 · 676ea1c · 676ea1c
1 parent c439a6a
commit 676ea1c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/pkg/pbq/store/memory/stores.go b/pkg/pbq/store/memory/stores.go
@@ -66,6 +66,8 @@ func (ms *memoryStores) CreateStore(ctx context.Context, partitionID partition.I
 }
 
 func (ms *memoryStores) DiscoverPartitions(ctx context.Context) ([]partition.ID, error) {
+	ms.RLock()
+	defer ms.RUnlock()
 	if ms.discoverFunc == nil {
 		partitionsIds := make([]partition.ID, 0)
 		for key := range ms.partitions {

diff --git a/pkg/reduce/reduce_test.go b/pkg/reduce/reduce_test.go
@@ -700,7 +700,7 @@ func TestDataForward_WithContextClose(t *testing.T) {
 	// wait for the partitions to be created
 	for {
 		partitionsList := pbqManager.ListPartitions()
-		if len(partitionsList) > 0 {
+		if len(partitionsList) > 1 {
 			childCancel()
 			break
 		}
@@ -717,7 +717,7 @@ func TestDataForward_WithContextClose(t *testing.T) {
 	for {
 		discoveredPartitions, _ = storeProvider.DiscoverPartitions(ctx)
 
-		if len(discoveredPartitions) > 0 {
+		if len(discoveredPartitions) > 1 {
 			break
 		}
 		select {