From d64b8844a545b4663d50ef83305a63d67afa332d Mon Sep 17 00:00:00 2001
From: Abhishek Pandey <abhishek@alcion.ai>
Date: Fri, 21 Jul 2023 00:05:33 -0700
Subject: [PATCH] Fix race cond

---
 src/internal/m365/onedrive/collection.go | 15 ++++++---------
 src/pkg/extensions/mock_extensions.go    | 16 ++++++++++++++--
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/internal/m365/onedrive/collection.go b/src/internal/m365/onedrive/collection.go
index d4561b534..fea3cc809 100644
--- a/src/internal/m365/onedrive/collection.go
+++ b/src/internal/m365/onedrive/collection.go
@@ -411,9 +411,6 @@ type driveStats struct {
 	itemsFound int64
 }
 
-// global counter
-var ctr int32
-
 // populateItems iterates through items added to the collection
 // and uses the collection `itemReader` to read the item
 func (oc *Collection) populateItems(ctx context.Context, errs *fault.Bus) {
@@ -455,13 +452,13 @@ func (oc *Collection) populateItems(ctx context.Context, errs *fault.Bus) {
 			defer wg.Done()
 			defer func() { <-semaphoreCh }()
 
-			itf := []extensions.CreateItemExtensioner{
-				&extensions.MockItemExtensionFactory{
-					ItemNumber: ctr,
-				},
-			}
+			var itf []extensions.CreateItemExtensioner
 
-			atomic.AddInt32(&ctr, 1)
+			if item.GetFile() != nil {
+				itf = []extensions.CreateItemExtensioner{
+					extensions.GetMockExtensions(),
+				}
+			}
 
 			// Read the item
 			oc.populateDriveItem(
diff --git a/src/pkg/extensions/mock_extensions.go b/src/pkg/extensions/mock_extensions.go
index 24975b7a8..2f390abcd 100644
--- a/src/pkg/extensions/mock_extensions.go
+++ b/src/pkg/extensions/mock_extensions.go
@@ -33,8 +33,7 @@ type MockExtension struct {
 }
 
 func (me *MockExtension) Read(p []byte) (int, error) {
-	// randomly fail with 50% prob
-	if me.ItemNumber > 3 {
+	if me.ItemNumber > 1 {
 		return 0, clues.New("mock read error")
 	}
 
@@ -103,6 +102,8 @@ func (m *MockItemExtensionFactory) CreateItemExtension(
 		return nil, clues.New("factory error")
 	}
 
+	logger.Ctx(ctx).Infow("mock extension created", "itemnumber", m.ItemNumber)
+
 	return &MockExtension{
 		Ctx:         ctx,
 		InnerRc:     rc,
@@ -113,3 +114,14 @@ func (m *MockItemExtensionFactory) CreateItemExtension(
 		ItemNumber:  m.ItemNumber,
 	}, nil
 }
+
+var ctr int32
+
+func GetMockExtensions() CreateItemExtensioner {
+	atomic.AddInt32(&ctr, 1)
+	mf := &MockItemExtensionFactory{
+		ItemNumber: ctr,
+	}
+
+	return mf
+}