Fix race while update branch returned from ref manager branch update (treeverse#9360)

nopcoder · web-flow · commit bd1f9033d916 · 2025-07-29T07:17:07.000+03:00
* Fix race while update branch returned from ref manager branch update

Fetching branch information using batch executor will return a shared
result. This result is then used to update the branch information which
can lead to race condition. This commit fixes the race by cloning the
branch information before returning it.

* test race condition
diff --git a/pkg/graveler/graveler.go b/pkg/graveler/graveler.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"slices"
 	"strconv"
 	"strings"
 	"time"
@@ -503,6 +504,18 @@ type Branch struct {
 	Hidden                   bool
 }
 
+// Clone creates a deep copy of the Branch
+func (b *Branch) Clone() *Branch {
+	clone := &Branch{
+		CommitID:                 b.CommitID,
+		StagingToken:             b.StagingToken,
+		SealedTokens:             slices.Clone(b.SealedTokens),
+		CompactedBaseMetaRangeID: b.CompactedBaseMetaRangeID,
+		Hidden:                   b.Hidden,
+	}
+	return clone
+}
+
 // BranchRecord holds BranchID with the associated Branch data
 type BranchRecord struct {
 	BranchID BranchID `db:"id"`
diff --git a/pkg/graveler/ref/manager.go b/pkg/graveler/ref/manager.go
@@ -456,7 +456,6 @@ func (m *Manager) BranchUpdate(ctx context.Context, repository *graveler.Reposit
 		release, err := m.branchOwnership.Own(ctx, requestID, string(branchID))
 		if err != nil {
 			logging.FromContext(ctx).
-				WithFields(logging.Fields{}).
 				WithError(err).
 				Warn("Failed to get ownership on branch; continuing but may be slow")
 		} else {
@@ -467,6 +466,10 @@ func (m *Manager) BranchUpdate(ctx context.Context, repository *graveler.Reposit
 	if err != nil {
 		return err
 	}
+
+	// clone the branch information to avoid mutating the shared result returned by batch executor
+	b = b.Clone()
+
 	newBranch, err := f(b)
 	// return on error or nothing to update
 	if err != nil || newBranch == nil {
diff --git a/pkg/graveler/ref/manager_test.go b/pkg/graveler/ref/manager_test.go
@@ -10,6 +10,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -23,8 +24,11 @@ import (
 	"github.com/treeverse/lakefs/pkg/graveler/ref"
 	"github.com/treeverse/lakefs/pkg/ident"
 	"github.com/treeverse/lakefs/pkg/kv"
+	"github.com/treeverse/lakefs/pkg/kv/kvtest"
 	"github.com/treeverse/lakefs/pkg/kv/mock"
+	"github.com/treeverse/lakefs/pkg/logging"
 	"github.com/treeverse/lakefs/pkg/testutil"
+	"go.uber.org/ratelimit"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -547,6 +551,67 @@ func TestManager_BranchUpdate(t *testing.T) {
 	}
 }
 
+func TestManager_BranchUpdateRaceCondition(t *testing.T) {
+	ctx := context.Background()
+	kvStore := kvtest.GetStore(ctx, t)
+	executor := batch.NewExecutor(logging.Dummy())
+	cfg := ref.ManagerConfig{
+		Executor:              executor,
+		KVStore:               kvStore,
+		KVStoreLimited:        kv.NewStoreLimiter(kvStore, ratelimit.NewUnlimited()),
+		AddressProvider:       ident.NewHexAddressProvider(),
+		RepositoryCacheConfig: testRepoCacheConfig,
+		CommitCacheConfig:     testCommitCacheConfig,
+		MaxBatchDelay:         10 * time.Millisecond,
+	}
+	cancelCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	go executor.Run(cancelCtx)
+	r := ref.NewRefManager(cfg, NewStorageConfigMock(config.SingleBlockstoreID))
+
+	const (
+		repoID   = "repo1"
+		branchID = "race-branch"
+	)
+	repository, err := r.CreateRepository(context.Background(), repoID, graveler.Repository{
+		StorageID:        "sid",
+		StorageNamespace: "s3://test-bucket",
+		CreationDate:     time.Now(),
+		DefaultBranchID:  "main",
+	})
+	testutil.Must(t, err)
+
+	// Set initial branch state
+	testutil.Must(t, r.SetBranch(ctx, repository, branchID, graveler.Branch{
+		CommitID:     "test-commit-id",
+		StagingToken: "staging1",
+	}))
+
+	// Use goroutines to create a proper race condition
+	var wg sync.WaitGroup
+
+	// Start goroutines that will modify the branch concurrently
+	const concurrentGoroutines = 5
+	wg.Add(concurrentGoroutines)
+	for i := range concurrentGoroutines {
+		go func() {
+			defer wg.Done()
+
+			_ = r.BranchUpdate(ctx, repository, branchID, func(branch *graveler.Branch) (*graveler.Branch, error) {
+				// Modify the branch information
+				branch.StagingToken = graveler.StagingToken("staging" + strconv.Itoa(i))
+				branch.SealedTokens = slices.Repeat([]graveler.StagingToken{branch.StagingToken}, i)
+				return branch, nil
+			})
+		}()
+	}
+	wg.Wait()
+
+	branch, err := r.GetBranch(ctx, repository, branchID)
+	require.NoError(t, err)
+	require.NotNil(t, branch)
+}
+
 func TestManager_DeleteBranch(t *testing.T) {
 	r, _ := testRefManager(t)
 	ctx := context.Background()
