Merge branch 'pks-storagemgr-prioritize-cancellation-over-commit' into 'master'

storagemgr: Fix flaky test when committing with cancelled context

Closes #5542

See merge request https://gitlab.com/gitlab-org/gitaly/-/merge_requests/6299

Merged-by: Quang-Minh Nguyen <qmnguyen@gitlab.com>
Approved-by: Sami Hiltunen <shiltunen@gitlab.com>
Approved-by: Quang-Minh Nguyen <qmnguyen@gitlab.com>
Reviewed-by: Sami Hiltunen <shiltunen@gitlab.com>
Co-authored-by: Patrick Steinhardt <psteinhardt@gitlab.com>

2 files changed, 65 insertions, 3 deletions

diff --git a/internal/gitaly/storage/storagemgr/partition_manager.go b/internal/gitaly/storage/storagemgr/partition_manager.go
index b44a72621..e30e5cc23 100644
--- a/internal/gitaly/storage/storagemgr/partition_manager.go
+++ b/internal/gitaly/storage/storagemgr/partition_manager.go
@@ -26,6 +26,13 @@ import (
 // ErrPartitionManagerClosed is returned when the PartitionManager stops processing transactions.
 var ErrPartitionManagerClosed = errors.New("partition manager closed")
 
+type transactionManagerFactory func(
+	storageMgr *storageManager,
+	cmdFactory git.CommandFactory,
+	housekeepingManager housekeeping.Manager,
+	relativePath, absoluteStateDir, stagingDir string,
+) *TransactionManager
+
 // PartitionManager is responsible for managing the lifecycle of each TransactionManager.
 type PartitionManager struct {
 	// storages are the storages configured in this Gitaly server. The map is keyed by the storage name.
@@ -34,6 +41,9 @@ type PartitionManager struct {
 	commandFactory git.CommandFactory
 	// housekeepingManager access to the housekeeping.Manager.
 	housekeepingManager housekeeping.Manager
+	// transactionManagerFactory is a factory to create TransactionManagers. This shouldn't ever be changed
+	// during normal operation, but can be used to adjust the transaction manager's behaviour in tests.
+	transactionManagerFactory transactionManagerFactory
 }
 
 // storageManager represents a single storage.
@@ -221,7 +231,28 @@ func NewPartitionManager(
 		}
 	}
 
-	return &PartitionManager{storages: storages, commandFactory: cmdFactory, housekeepingManager: housekeepingManager}, nil
+	return &PartitionManager{
+		storages:            storages,
+		commandFactory:      cmdFactory,
+		housekeepingManager: housekeepingManager,
+		transactionManagerFactory: func(
+			storageMgr *storageManager,
+			cmdFactory git.CommandFactory,
+			housekeepingManager housekeeping.Manager,
+			relativePath, absoluteStateDir, stagingDir string,
+		) *TransactionManager {
+			return NewTransactionManager(
+				storageMgr.database,
+				storageMgr.path,
+				relativePath,
+				absoluteStateDir,
+				stagingDir,
+				cmdFactory,
+				housekeepingManager,
+				storageMgr.repoFactory,
+			)
+		},
+	}, nil
 }
 
 func stagingDirectoryPath(storagePath string) string {
@@ -283,7 +314,7 @@ func (pm *PartitionManager) Begin(ctx context.Context, repo storage.Repository,
 				return nil, fmt.Errorf("create staging directory: %w", err)
 			}
 
-			mgr := NewTransactionManager(storageMgr.database, storageMgr.path, relativePath, absoluteStateDir, stagingDir, pm.commandFactory, pm.housekeepingManager, storageMgr.repoFactory)
+			mgr := pm.transactionManagerFactory(storageMgr, pm.commandFactory, pm.housekeepingManager, relativePath, absoluteStateDir, stagingDir)
 
 			ptn.transactionManager = mgr
 
diff --git a/internal/gitaly/storage/storagemgr/partition_manager_test.go b/internal/gitaly/storage/storagemgr/partition_manager_test.go
index dc0a596bf..2e5101944 100644
--- a/internal/gitaly/storage/storagemgr/partition_manager_test.go
+++ b/internal/gitaly/storage/storagemgr/partition_manager_test.go
@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/catfile"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/housekeeping"
@@ -157,7 +158,8 @@ func TestPartitionManager(t *testing.T) {
 	}
 
 	type setupData struct {
-		steps steps
+		steps                     steps
+		transactionManagerFactory transactionManagerFactory
 	}
 
 	for _, tc := range []struct {
@@ -388,6 +390,31 @@ func TestPartitionManager(t *testing.T) {
 							expectedError: context.Canceled,
 						},
 					},
+					transactionManagerFactory: func(
+						storageMgr *storageManager,
+						commandFactory git.CommandFactory,
+						housekeepingManager housekeeping.Manager,
+						relativePath, absoluteStateDir, stagingDir string,
+					) *TransactionManager {
+						txMgr := NewTransactionManager(
+							storageMgr.database,
+							storageMgr.path,
+							relativePath,
+							absoluteStateDir,
+							stagingDir,
+							commandFactory,
+							housekeepingManager,
+							storageMgr.repoFactory,
+						)
+
+						// Unset the admission queue. This has the effect that we will block
+						// indefinitely when trying to submit the transaction to it in the
+						// commit step. Like this, we can racelessly verify that the context
+						// cancellation does indeed abort the commit.
+						txMgr.admissionQueue = nil
+
+						return txMgr
+					},
 				}
 			},
 		},
@@ -619,6 +646,10 @@ func TestPartitionManager(t *testing.T) {
 			partitionManager, err := NewPartitionManager(cfg.Storages, cmdFactory, housekeepingManager, localRepoFactory, testhelper.SharedLogger(t))
 			require.NoError(t, err)
 
+			if setup.transactionManagerFactory != nil {
+				partitionManager.transactionManagerFactory = setup.transactionManagerFactory
+			}
+
 			defer func() {
 				partitionManager.Close()
 				for _, storage := range cfg.Storages {