Partition fork with source repository in CreateFork

Gitaly's TransactionManager requires all object pool members to be
in the same partition. Currently we're ensuring doing that generally
by extracting the additional repository from the request and ensuring
the target repository of the RPC gets partitioned with it. Additional
repository is used in the ObjectPoolService's RPCs to tag the other
repository being accessed in the pool related operations, and it may
be either the object pool or one of the member repositories depending
on the RPC.

When a repository is first accessed, we also check whether it has an
existing alternate link on the disk, and place the repository in the
same partition with its alternate if so.

These two methods are enough to ensure in general the pools and their
members get partitioned together. One execption to this is CreateFork.
The created fork should be placed in the same partition as the origin
repository as they'll both eventually be connected to the same pool.
CreateFork does not tag the source repository as an additional
repository so the general handling does not apply for it. Tagging it
as the additional repository won't work with Praefect as Praefect
rewrites the paths of additional repositories. The additional repository
is fetched through the API so it needs to have its original relative
path intact.

This commit introduces special handling for CreateFork. The transaction
middleware checks whether the request is a CreateForkRequest. If so,
the source repository is extracted and the newly created fork gets
partitioned with it.

Along with the behavior changes, we add a test that exercises the
entire fork creation flow as typically done. Turns out Gitaly didn't
have a test covering the scenario at all.

1 files changed, 25 insertions, 8 deletions

diff --git a/internal/gitaly/storage/storagemgr/middleware.go b/internal/gitaly/storage/storagemgr/middleware.go
index 290545def..a72bbe88a 100644
--- a/internal/gitaly/storage/storagemgr/middleware.go
+++ b/internal/gitaly/storage/storagemgr/middleware.go
@@ -26,6 +26,9 @@ var (
 	// ErrQuarantineWithoutSnapshotRelativePath is returned when a request is configured with a quarantine but the snapshot's
 	// relative path was not sent in a header.
 	ErrQuarantineWithoutSnapshotRelativePath = errors.New("quarantined request did not contain snapshot relative path")
+	// ErrRepositoriesInDifferentStorages is returned when trying to access two repositories in different storages in the
+	// same transaction.
+	ErrRepositoriesInDifferentStorages = structerr.NewInvalidArgument("additional and target repositories are in different storages")
 )
 
 // NonTransactionalRPCs are the RPCs that do not support transactions.
@@ -247,22 +250,36 @@ func transactionalizeRequest(ctx context.Context, logger log.Logger, txRegistry
 		return transactionalizedRequest{}, err
 	}
 
-	var alternateRelativePath string
-	if additionalRepo, err := methodInfo.AdditionalRepo(req); err != nil {
+	// Object pools need to be placed in the same partition as their members. Below we figure out which repository,
+	// if any, the target repository of the RPC must be partitioned with. We figure this out using two strategies:
+	//
+	// The general case is handled by extracting the additional repository from the RPC, and paritioning the target
+	// repository of the RPC with the additional repository. Many of the ObjectPoolService's RPCs operate on two
+	// repositories. Depending on the RPC, the additional repository is either the object pool itself or a member
+	// of the pool.
+	//
+	// CreateFork is special cased. The fork must partitioned with the source repository in order to successfully
+	// link it with the object pool later. The source repository is not tagged as additional repository in the
+	// CreateForkRequest. If the request is CreateForkRequest, we extract the source repository and partition the
+	// fork with it.
+	var additionalRepo *gitalypb.Repository
+	if additionalRepo, err = methodInfo.AdditionalRepo(req); err != nil {
 		if !errors.Is(err, protoregistry.ErrRepositoryFieldNotFound) {
 			return transactionalizedRequest{}, fmt.Errorf("extract additional repository: %w", err)
 		}
 
 		// There was no additional repository.
-	} else {
-		if additionalRepo.StorageName != targetRepo.StorageName {
-			return transactionalizedRequest{}, structerr.NewInvalidArgument("additional and target repositories are in different storages")
-		}
+	}
+
+	if req, ok := req.(*gitalypb.CreateForkRequest); ok {
+		additionalRepo = req.GetSourceRepository()
+	}
 
-		alternateRelativePath = additionalRepo.RelativePath
+	if additionalRepo != nil && additionalRepo.StorageName != targetRepo.StorageName {
+		return transactionalizedRequest{}, ErrRepositoriesInDifferentStorages
 	}
 
-	tx, err := mgr.Begin(ctx, targetRepo.StorageName, targetRepo.RelativePath, alternateRelativePath, methodInfo.Operation == protoregistry.OpAccessor)
+	tx, err := mgr.Begin(ctx, targetRepo.StorageName, targetRepo.RelativePath, additionalRepo.GetRelativePath(), methodInfo.Operation == protoregistry.OpAccessor)
 	if err != nil {
 		return transactionalizedRequest{}, fmt.Errorf("begin transaction: %w", err)
 	}