Merge branch 'smh-fix-ptn-mgr-flake' into 'master'

Fix flaky PartitionManager test

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

Merged-by: Justin Tobler <jtobler@gitlab.com>
Approved-by: Justin Tobler <jtobler@gitlab.com>
Co-authored-by: Sami Hiltunen <shiltunen@gitlab.com>

2 files changed, 34 insertions, 24 deletions

diff --git a/internal/gitaly/storage/storagemgr/partition_manager.go b/internal/gitaly/storage/storagemgr/partition_manager.go
index 13f8c70a9..4349d5a8d 100644
--- a/internal/gitaly/storage/storagemgr/partition_manager.go
+++ b/internal/gitaly/storage/storagemgr/partition_manager.go
@@ -26,13 +26,22 @@ import (
 // ErrPartitionManagerClosed is returned when the PartitionManager stops processing transactions.
 var ErrPartitionManagerClosed = errors.New("partition manager closed")
 
+// transactionManager is the interface of TransactionManager as used by PartitionManager. See the
+// TransactionManager's documentation for more details.
+type transactionManager interface {
+	Begin(context.Context, TransactionOptions) (*Transaction, error)
+	Run() error
+	Close()
+	isClosing() bool
+}
+
 type transactionManagerFactory func(
 	partitionID partitionID,
 	storageMgr *storageManager,
 	cmdFactory git.CommandFactory,
 	housekeepingManager housekeeping.Manager,
 	relativePath, absoluteStateDir, stagingDir string,
-) *TransactionManager
+) transactionManager
 
 // PartitionManager is responsible for managing the lifecycle of each TransactionManager.
 type PartitionManager struct {
@@ -154,7 +163,7 @@ type partition struct {
 	// TransactionManager has closed and it is safe to start another one.
 	transactionManagerClosed chan struct{}
 	// transactionManager manages all transactions for the partition.
-	transactionManager *TransactionManager
+	transactionManager transactionManager
 	// pendingTransactionCount holds the current number of in flight transactions being processed by the manager.
 	pendingTransactionCount uint
 }
@@ -250,7 +259,7 @@ func NewPartitionManager(
 			cmdFactory git.CommandFactory,
 			housekeepingManager housekeeping.Manager,
 			relativePath, absoluteStateDir, stagingDir string,
-		) *TransactionManager {
+		) transactionManager {
 			return NewTransactionManager(
 				partitionID,
 				logger,
diff --git a/internal/gitaly/storage/storagemgr/partition_manager_test.go b/internal/gitaly/storage/storagemgr/partition_manager_test.go
index 77cab7d36..a80907445 100644
--- a/internal/gitaly/storage/storagemgr/partition_manager_test.go
+++ b/internal/gitaly/storage/storagemgr/partition_manager_test.go
@@ -26,6 +26,12 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
 )
 
+// stoppedTransactionManager is a wrapper type that prevents the transactionManager from
+// running. This is useful in tests that test certain order of operations.
+type stoppedTransactionManager struct{ transactionManager }
+
+func (stoppedTransactionManager) Run() error { return nil }
+
 func TestPartitionManager(t *testing.T) {
 	t.Parallel()
 
@@ -373,26 +379,21 @@ func TestPartitionManager(t *testing.T) {
 						commandFactory git.CommandFactory,
 						housekeepingManager housekeeping.Manager,
 						relativePath, absoluteStateDir, stagingDir string,
-					) *TransactionManager {
-						txMgr := NewTransactionManager(
-							partitionID,
-							storageMgr.logger,
-							storageMgr.database,
-							storageMgr.path,
-							relativePath,
-							absoluteStateDir,
-							stagingDir,
-							commandFactory,
-							housekeepingManager,
-							storageMgr.repoFactory,
-						)
-
-						// Fake a preexisting apply notification. This ensures that we would
-						// block indefinitely waiting for the notifcation and thus allows us to
-						// assert that we can indeed cancel this via the context.
-						txMgr.snapshotLocks[0] = &snapshotLock{}
-
-						return txMgr
+					) transactionManager {
+						return stoppedTransactionManager{
+							transactionManager: NewTransactionManager(
+								partitionID,
+								storageMgr.logger,
+								storageMgr.database,
+								storageMgr.path,
+								relativePath,
+								absoluteStateDir,
+								stagingDir,
+								commandFactory,
+								housekeepingManager,
+								storageMgr.repoFactory,
+							),
+						}
 					},
 				}
 			},
@@ -426,7 +427,7 @@ func TestPartitionManager(t *testing.T) {
 						commandFactory git.CommandFactory,
 						housekeepingManager housekeeping.Manager,
 						relativePath, absoluteStateDir, stagingDir string,
-					) *TransactionManager {
+					) transactionManager {
 						txMgr := NewTransactionManager(
 							partitionID,
 							logger,