Record initial values of updated references

The reference transaction hook has an issue where for both force updates
and reference deletions the old OID is zero OID. It's thus not possible
to distinguish them without resolving the references value in the 'prepare'
stage of the reference transaction hook. If we resolve the value in the
'prepare' phase, we also don't want to stage an update for a reference
before the 'commit' phase as the reference may not actually get updated.
This commit introduces RecordInitialReferenceValues method on the transaction
that can be used to record a reference's initial value as seen in the
'prepare' phase prior to any updates without actually staging the
reference for an update. This recorded initial values can then be used as
the old OID when ultimately staging the references for an update through
UpdateReferences.

2 files changed, 320 insertions, 1 deletions

diff --git a/internal/gitaly/storage/storagemgr/transaction_manager.go b/internal/gitaly/storage/storagemgr/transaction_manager.go
index 23ed201d3..5de67ab82 100644
--- a/internal/gitaly/storage/storagemgr/transaction_manager.go
+++ b/internal/gitaly/storage/storagemgr/transaction_manager.go
@@ -179,6 +179,7 @@ type Transaction struct {
 	snapshot Snapshot
 
 	skipVerificationFailures bool
+	initialReferenceValues   map[git.ReferenceName]git.ObjectID
 	referenceUpdates         ReferenceUpdates
 	defaultBranchUpdate      *DefaultBranchUpdate
 	customHooksUpdate        *CustomHooksUpdate
@@ -400,6 +401,49 @@ func (txn *Transaction) SkipVerificationFailures() {
 	txn.skipVerificationFailures = true
 }
 
+// RecordInitialReferenceValues records the initial values of the reference if they haven't yet been recorded. If oid is
+// not a zero OID, it's used as the initial value. If oid is a zero value, the reference's actual value is resolved.
+//
+// The reference's first recorded value is used as its old OID in the final committed update. RecordInitialReferenceValues
+// can be used to record the value without staging an update in the transaction. This is useful for example generally recording
+// the initial value in the 'prepare' phase of the reference transaction hook before any changes are made without staging
+// any updates before the 'committed' phase is reached.
+func (txn *Transaction) RecordInitialReferenceValues(ctx context.Context, initialValues map[git.ReferenceName]git.ObjectID) error {
+	if txn.initialReferenceValues == nil {
+		txn.initialReferenceValues = make(map[git.ReferenceName]git.ObjectID, len(initialValues))
+	}
+
+	for reference, oid := range initialValues {
+		if _, ok := txn.initialReferenceValues[reference]; ok {
+			// If the reference's starting value has already been recorded, we don't have to record it again.
+			continue
+		}
+
+		objectHash, err := txn.stagingRepository.ObjectHash(ctx)
+		if err != nil {
+			return fmt.Errorf("object hash: %w", err)
+		}
+
+		if objectHash.IsZeroOID(oid) {
+			// If this is a zero OID, resolve the value to see if this is a force update or the
+			// reference doesn't exist.
+			if current, err := txn.stagingRepository.ResolveRevision(ctx, reference.Revision()); err != nil {
+				if !errors.Is(err, git.ErrReferenceNotFound) {
+					return fmt.Errorf("resolve revision: %w", err)
+				}
+
+				// The reference doesn't exist, leave the value as zero oid.
+			} else {
+				oid = current
+			}
+		}
+
+		txn.initialReferenceValues[reference] = oid
+	}
+
+	return nil
+}
+
 // 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
@@ -413,6 +457,10 @@ func (txn *Transaction) UpdateReferences(updates ReferenceUpdates) {
 
 	for reference, update := range updates {
 		oldOID := update.OldOID
+		if initialValue, ok := txn.initialReferenceValues[reference]; ok {
+			oldOID = initialValue
+		}
+
 		if previousUpdate, ok := txn.referenceUpdates[reference]; ok {
 			oldOID = previousUpdate.OldOID
 		}
diff --git a/internal/gitaly/storage/storagemgr/transaction_manager_test.go b/internal/gitaly/storage/storagemgr/transaction_manager_test.go
index 7d3e59855..7575861c6 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
 	}
 
+	// RecordInitialReferenceValues calls RecordInitialReferenceValues on a transaction.
+	type RecordInitialReferenceValues struct {
+		// TransactionID identifies the transaction to prepare the reference updates on.
+		TransactionID int
+		// InitialValues are the initial values to record.
+		InitialValues map[git.ReferenceName]git.ObjectID
+	}
+
 	// UpdateReferences calls UpdateReferences on a transaction.
 	type UpdateReferences struct {
 		// TransactionID identifies the transaction to update references on.
@@ -1415,7 +1423,7 @@ func TestTransactionManager(t *testing.T) {
 			},
 		},
 		{
-			desc: "update reference multiple times fails due to wrong intial value",
+			desc: "update reference multiple times fails due to wrong initial value",
 			steps: steps{
 				StartManager{},
 				Begin{},
@@ -1441,6 +1449,264 @@ func TestTransactionManager(t *testing.T) {
 			},
 		},
 		{
+			desc: "recording initial value of a reference stages no updates",
+			steps: steps{
+				StartManager{},
+				Begin{},
+				RecordInitialReferenceValues{
+					InitialValues: map[git.ReferenceName]git.ObjectID{
+						"refs/heads/main": setup.Commits.First.OID,
+					},
+				},
+				Commit{},
+			},
+			expectedState: StateAssertion{
+				Database: DatabaseState{
+					string(keyAppliedLogIndex(relativePath)): LogIndex(1).toProto(),
+				},
+			},
+		},
+		{
+			desc: "update reference with non-existent initial value",
+			steps: steps{
+				StartManager{},
+				Begin{},
+				RecordInitialReferenceValues{
+					InitialValues: map[git.ReferenceName]git.ObjectID{
+						"refs/heads/main": setup.ObjectHash.ZeroOID,
+					},
+				},
+				UpdateReferences{
+					// The old oid is ignored as the references old value was already recorded.
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {NewOID: setup.Commits.First.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.First.OID.String()}},
+					},
+				},
+			},
+		},
+		{
+			desc: "update reference with the zero oid initial value",
+			steps: steps{
+				StartManager{},
+				Begin{
+					TransactionID: 1,
+				},
+				Commit{
+					TransactionID: 1,
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {OldOID: setup.ObjectHash.ZeroOID, NewOID: setup.Commits.First.OID},
+					},
+				},
+				Begin{
+					TransactionID: 2,
+					ExpectedSnapshot: Snapshot{
+						ReadIndex: 1,
+					},
+				},
+				RecordInitialReferenceValues{
+					TransactionID: 2,
+					InitialValues: map[git.ReferenceName]git.ObjectID{
+						"refs/heads/main": setup.ObjectHash.ZeroOID,
+					},
+				},
+				UpdateReferences{
+					TransactionID: 2,
+					// The old oid is ignored as the references old value was already recorded.
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {NewOID: setup.Commits.Second.OID},
+					},
+				},
+				Commit{
+					TransactionID: 2,
+				},
+			},
+			expectedState: StateAssertion{
+				Database: DatabaseState{
+					string(keyAppliedLogIndex(relativePath)): LogIndex(2).toProto(),
+				},
+				Repositories: RepositoryStates{
+					relativePath: {
+						References: []git.Reference{{Name: "refs/heads/main", Target: setup.Commits.Second.OID.String()}},
+					},
+				},
+			},
+		},
+		{
+			desc: "update reference with the correct initial value",
+			steps: steps{
+				StartManager{},
+				Begin{
+					TransactionID: 1,
+				},
+				Commit{
+					TransactionID: 1,
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {OldOID: setup.ObjectHash.ZeroOID, NewOID: setup.Commits.First.OID},
+					},
+				},
+				Begin{
+					TransactionID: 2,
+					ExpectedSnapshot: Snapshot{
+						ReadIndex: 1,
+					},
+				},
+				RecordInitialReferenceValues{
+					TransactionID: 2,
+					InitialValues: map[git.ReferenceName]git.ObjectID{
+						"refs/heads/main": setup.Commits.First.OID,
+					},
+				},
+				UpdateReferences{
+					TransactionID: 2,
+					// The old oid is ignored as the references old value was already recorded.
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {NewOID: setup.Commits.Second.OID},
+					},
+				},
+				Commit{
+					TransactionID: 2,
+				},
+			},
+			expectedState: StateAssertion{
+				Database: DatabaseState{
+					string(keyAppliedLogIndex(relativePath)): LogIndex(2).toProto(),
+				},
+				Repositories: RepositoryStates{
+					relativePath: {
+						References: []git.Reference{{Name: "refs/heads/main", Target: setup.Commits.Second.OID.String()}},
+					},
+				},
+			},
+		},
+		{
+			desc: "update reference with the incorrect initial value",
+			steps: steps{
+				StartManager{},
+				Begin{
+					TransactionID: 1,
+				},
+				Commit{
+					TransactionID: 1,
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {OldOID: setup.ObjectHash.ZeroOID, NewOID: setup.Commits.First.OID},
+					},
+				},
+				Begin{
+					TransactionID: 2,
+					ExpectedSnapshot: Snapshot{
+						ReadIndex: 1,
+					},
+				},
+				RecordInitialReferenceValues{
+					TransactionID: 2,
+					InitialValues: map[git.ReferenceName]git.ObjectID{
+						"refs/heads/main": setup.Commits.Third.OID,
+					},
+				},
+				UpdateReferences{
+					TransactionID: 2,
+					// The old oid is ignored as the references old value was already recorded.
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {NewOID: setup.Commits.Second.OID},
+					},
+				},
+				Commit{
+					TransactionID: 2,
+					ExpectedError: ReferenceVerificationError{
+						ReferenceName: "refs/heads/main",
+						ExpectedOID:   setup.Commits.Third.OID,
+						ActualOID:     setup.Commits.First.OID,
+					},
+				},
+			},
+			expectedState: StateAssertion{
+				Database: DatabaseState{
+					string(keyAppliedLogIndex(relativePath)): LogIndex(1).toProto(),
+				},
+				Repositories: RepositoryStates{
+					relativePath: {
+						References: []git.Reference{{Name: "refs/heads/main", Target: setup.Commits.First.OID.String()}},
+					},
+				},
+			},
+		},
+		{
+			desc: "initial value is only recorded on the first time",
+			steps: steps{
+				StartManager{},
+				Begin{},
+				RecordInitialReferenceValues{
+					InitialValues: map[git.ReferenceName]git.ObjectID{
+						"refs/heads/main": setup.ObjectHash.ZeroOID,
+					},
+				},
+				RecordInitialReferenceValues{
+					InitialValues: map[git.ReferenceName]git.ObjectID{
+						"refs/heads/main":     setup.Commits.Third.OID,
+						"refs/heads/branch-2": setup.ObjectHash.ZeroOID,
+					},
+				},
+				Commit{
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main":     {NewOID: setup.Commits.First.OID},
+						"refs/heads/branch-2": {NewOID: setup.Commits.First.OID},
+					},
+				},
+			},
+			expectedState: StateAssertion{
+				Database: DatabaseState{
+					string(keyAppliedLogIndex(relativePath)): LogIndex(1).toProto(),
+				},
+				Repositories: RepositoryStates{
+					relativePath: {
+						References: []git.Reference{
+							{Name: "refs/heads/branch-2", Target: setup.Commits.First.OID.String()},
+							{Name: "refs/heads/main", Target: setup.Commits.First.OID.String()},
+						},
+					},
+				},
+			},
+		},
+		{
+			desc: "initial value is set on the first update",
+			steps: steps{
+				StartManager{},
+				Begin{},
+				UpdateReferences{
+					ReferenceUpdates: ReferenceUpdates{
+						"refs/heads/main": {OldOID: setup.ObjectHash.ZeroOID, NewOID: setup.Commits.First.OID},
+					},
+				},
+				RecordInitialReferenceValues{
+					InitialValues: map[git.ReferenceName]git.ObjectID{
+						"refs/heads/main": setup.Commits.Third.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.First.OID.String()}},
+					},
+				},
+			},
+		},
+		{
 			desc: "set custom hooks successfully",
 			steps: steps{
 				StartManager{},
@@ -3958,6 +4224,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 RecordInitialReferenceValues:
+					require.Contains(t, openTransactions, step.TransactionID, "test error: record initial reference value on transaction before beginning it")
+
+					transaction := openTransactions[step.TransactionID]
+					require.NoError(t, transaction.RecordInitialReferenceValues(ctx, step.InitialValues))
 				case UpdateReferences:
 					require.Contains(t, openTransactions, step.TransactionID, "test error: reference updates aborted on committed before beginning it")