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, 41 insertions, 2 deletions

diff --git a/internal/gitaly/storage/storagemgr/middleware_test.go b/internal/gitaly/storage/storagemgr/middleware_test.go
index 794222874..9196d10b0 100644
--- a/internal/gitaly/storage/storagemgr/middleware_test.go
+++ b/internal/gitaly/storage/storagemgr/middleware_test.go
@@ -36,6 +36,7 @@ type mockRepositoryService struct {
 	removeRepositoryFunc   func(context.Context, *gitalypb.RemoveRepositoryRequest) (*gitalypb.RemoveRepositoryResponse, error)
 	setCustomHooksFunc     func(gitalypb.RepositoryService_SetCustomHooksServer) error
 	getCustomHooksFunc     func(*gitalypb.GetCustomHooksRequest, gitalypb.RepositoryService_GetCustomHooksServer) error
+	createForkFunc         func(context.Context, *gitalypb.CreateForkRequest) (*gitalypb.CreateForkResponse, error)
 	gitalypb.UnimplementedRepositoryServiceServer
 }
 
@@ -59,6 +60,10 @@ func (m mockRepositoryService) GetCustomHooks(req *gitalypb.GetCustomHooksReques
 	return m.getCustomHooksFunc(req, stream)
 }
 
+func (m mockRepositoryService) CreateFork(ctx context.Context, req *gitalypb.CreateForkRequest) (*gitalypb.CreateForkResponse, error) {
+	return m.createForkFunc(ctx, req)
+}
+
 type mockHealthService struct {
 	checkFunc func(context.Context, *grpc_health_v1.HealthCheckRequest) (*grpc_health_v1.HealthCheckResponse, error)
 	grpc_health_v1.UnimplementedHealthServer
@@ -396,6 +401,35 @@ messages and behavior by erroring out the requests before they even hit this int
 			},
 			expectHandlerInvoked: true,
 		},
+		{
+			desc: "successful CreateFork request",
+			performRequest: func(t *testing.T, ctx context.Context, cc *grpc.ClientConn) {
+				resp, err := gitalypb.NewRepositoryServiceClient(cc).CreateFork(ctx, &gitalypb.CreateForkRequest{
+					Repository:       validAdditionalRepository(),
+					SourceRepository: validRepository(),
+				})
+				require.NoError(t, err)
+				testhelper.ProtoEqual(t, &gitalypb.CreateForkResponse{}, resp)
+			},
+			assertAdditionalRepository: func(t *testing.T, ctx context.Context, actual *gitalypb.Repository) {
+				testhelper.ProtoEqual(t, validRepository(), actual)
+			},
+			expectHandlerInvoked: true,
+		},
+		{
+			desc: "CreateFork fails due to repositories in different storages",
+			performRequest: func(t *testing.T, ctx context.Context, cc *grpc.ClientConn) {
+				sourceRepository := validRepository()
+				sourceRepository.StorageName = "different_storage"
+
+				resp, err := gitalypb.NewRepositoryServiceClient(cc).CreateFork(ctx, &gitalypb.CreateForkRequest{
+					Repository:       validAdditionalRepository(),
+					SourceRepository: sourceRepository,
+				})
+				require.Equal(t, status.Error(codes.InvalidArgument, storagemgr.ErrRepositoriesInDifferentStorages.Error()), err)
+				require.Nil(t, resp)
+			},
+		},
 	} {
 		t.Run(tc.desc, func(t *testing.T) {
 			cfg := testcfg.Build(t)
@@ -414,7 +448,7 @@ messages and behavior by erroring out the requests before they even hit this int
 			handlerInvoked := false
 			var transactionID storage.TransactionID
 
-			assertHandler := func(ctx context.Context, isMutator bool, repo *gitalypb.Repository) {
+			assertHandler := func(ctx context.Context, shouldBeQuarantined bool, repo *gitalypb.Repository) {
 				handlerInvoked = true
 
 				// The repositories should be equal except for the relative path which
@@ -427,7 +461,7 @@ messages and behavior by erroring out the requests before they even hit this int
 				expectedRepo.RelativePath = ""
 				actualRepo.RelativePath = ""
 
-				if isMutator {
+				if shouldBeQuarantined {
 					// Mutators should have quarantine directory configured.
 					assert.NotEmpty(t, actualRepo.GitObjectDirectory)
 					actualRepo.GitObjectDirectory = ""
@@ -476,6 +510,11 @@ messages and behavior by erroring out the requests before they even hit this int
 					},
 				})
 				gitalypb.RegisterRepositoryServiceServer(server, mockRepositoryService{
+					createForkFunc: func(ctx context.Context, req *gitalypb.CreateForkRequest) (*gitalypb.CreateForkResponse, error) {
+						assertHandler(ctx, false, req.GetRepository())
+						tc.assertAdditionalRepository(t, ctx, req.GetSourceRepository())
+						return &gitalypb.CreateForkResponse{}, tc.handlerError
+					},
 					objectFormatFunc: func(ctx context.Context, req *gitalypb.ObjectFormatRequest) (*gitalypb.ObjectFormatResponse, error) {
 						assertHandler(ctx, false, req.GetRepository())
 						return &gitalypb.ObjectFormatResponse{}, tc.handlerError