Support updating a reference multiple times within a transaction

Transaction manager currently only supports setting the reference
updates of a transaction once. This makes it difficult to gather the
reference updates from the reference transaction hook as there may
be multiple invocations of the hook within a single transaction. With
the current way things are, only the updates from the last invocation
of the hook would be committed.

This commit improves upon this by supporting multiple updates to the
references withing a single transaction. This enables then recording
the updates from the hooks which may be invoked multiple times. As
each transaction will only commit the reference updates at the end,
the updates withing a transaction are compacted from the start state
to the end state. So updates like 'oid-1 -> oid-2 -> oid-3' are
ultimately comitted as 'oid-1 -> oid-3' as the intermediate states
within a transaction don't make a difference when it comes to the
final result.

2 files changed, 88 insertions, 3 deletions

diff --git a/internal/gitaly/storage/storagemgr/transaction_manager.go b/internal/gitaly/storage/storagemgr/transaction_manager.go
index d08491c5f..23ed201d3 100644
--- a/internal/gitaly/storage/storagemgr/transaction_manager.go
+++ b/internal/gitaly/storage/storagemgr/transaction_manager.go
@@ -400,10 +400,28 @@ func (txn *Transaction) SkipVerificationFailures() {
 	txn.skipVerificationFailures = true
 }
 
-// UpdateReferences updates the given references as part of the transaction. If UpdateReferences is called
-// multiple times, only the changes from the latest invocation take place.
+// UpdateReferences updates the given references as part of the transaction.
+//
+// If a reference is updated multiple times during a transaction, its first recorded old OID is kept
+// and the new OID is updated. This means updates like 'oid-1 -> oid-2 -> oid-3' will ultimately be
+// committed as 'oid-1 -> oid-3'. The intermediate states are not relevant when committing the write
+// to the actual repository.
 func (txn *Transaction) UpdateReferences(updates ReferenceUpdates) {
-	txn.referenceUpdates = updates
+	if txn.referenceUpdates == nil {
+		txn.referenceUpdates = ReferenceUpdates{}
+	}
+
+	for reference, update := range updates {
+		oldOID := update.OldOID
+		if previousUpdate, ok := txn.referenceUpdates[reference]; ok {
+			oldOID = previousUpdate.OldOID
+		}
+
+		txn.referenceUpdates[reference] = ReferenceUpdate{
+			OldOID: oldOID,
+			NewOID: update.NewOID,
+		}
+	}
 }
 
 // DeleteRepository deletes the repository when the transaction is committed.
diff --git a/internal/gitaly/storage/storagemgr/transaction_manager_test.go b/internal/gitaly/storage/storagemgr/transaction_manager_test.go
index 42e90c188..7d3e59855 100644
--- a/internal/gitaly/storage/storagemgr/transaction_manager_test.go
+++ b/internal/gitaly/storage/storagemgr/transaction_manager_test.go
@@ -352,6 +352,14 @@ func TestTransactionManager(t *testing.T) {
 		ExpectedError error
 	}
 
+	// UpdateReferences calls UpdateReferences on a transaction.
+	type UpdateReferences struct {
+		// TransactionID identifies the transaction to update references on.
+		TransactionID int
+		// ReferenceUpdates are the reference updates to make.
+		ReferenceUpdates ReferenceUpdates
+	}
+
 	// Rollback calls Rollback on a transaction.
 	type Rollback struct {
 		// TransactionID identifies the transaction to rollback.
@@ -1379,6 +1387,60 @@ func TestTransactionManager(t *testing.T) {
 			},
 		},
 		{
+			desc: "update reference multiple times successfully in a transaction",
+			steps: steps{
+				StartManager{},
+				Begin{},
+				UpdateReferences{
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {OldOID: setup.ObjectHash.ZeroOID, NewOID: setup.Commits.First.OID},
+					},
+				},
+				UpdateReferences{
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {OldOID: setup.Commits.First.OID, NewOID: setup.Commits.Second.OID},
+					},
+				},
+				Commit{},
+			},
+			expectedState: StateAssertion{
+				Database: DatabaseState{
+					string(keyAppliedLogIndex(relativePath)): LogIndex(1).toProto(),
+				},
+				Repositories: RepositoryStates{
+					relativePath: {
+						References: []git.Reference{{Name: "refs/heads/main", Target: setup.Commits.Second.OID.String()}},
+					},
+				},
+			},
+		},
+		{
+			desc: "update reference multiple times fails due to wrong intial value",
+			steps: steps{
+				StartManager{},
+				Begin{},
+				UpdateReferences{
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {OldOID: setup.Commits.First.OID, NewOID: setup.Commits.Second.OID},
+					},
+				},
+				UpdateReferences{
+					ReferenceUpdates: ReferenceUpdates{
+						// The old oid should be ignored since there's already a recorded initial value for the
+						// reference.
+						"refs/heads/main": {NewOID: setup.Commits.Third.OID},
+					},
+				},
+				Commit{
+					ExpectedError: ReferenceVerificationError{
+						ReferenceName: "refs/heads/main",
+						ExpectedOID:   setup.Commits.First.OID,
+						ActualOID:     setup.ObjectHash.ZeroOID,
+					},
+				},
+			},
+		},
+		{
 			desc: "set custom hooks successfully",
 			steps: steps{
 				StartManager{},
@@ -3896,6 +3958,11 @@ func TestTransactionManager(t *testing.T) {
 					// determine that the deletion has actually been admitted, and is waiting for application to ensure the commit order is always
 					// as expected by the test.
 					<-transaction.admitted
+				case UpdateReferences:
+					require.Contains(t, openTransactions, step.TransactionID, "test error: reference updates aborted on committed before beginning it")
+
+					transaction := openTransactions[step.TransactionID]
+					transaction.UpdateReferences(step.ReferenceUpdates)
 				case Rollback:
 					require.Contains(t, openTransactions, step.TransactionID, "test error: transaction rollbacked before beginning it")
 					require.Equal(t, step.ExpectedError, openTransactions[step.TransactionID].Rollback())