Merge branch 'pks-transactions-operations' into 'master'

Start using transactions for UserCreateBranch

See merge request gitlab-org/gitaly!2364

12 files changed, 97 insertions, 24 deletions

diff --git a/changelogs/unreleased/pks-transactions-operations.yml b/changelogs/unreleased/pks-transactions-operations.yml
new file mode 100644
index 000000000..43861fa2f
--- /dev/null
+++ b/changelogs/unreleased/pks-transactions-operations.yml
@@ -0,0 +1,5 @@
+---
+title: Start using transactions for UserCreateBranch
+merge_request: 2364
+author:
+type: added
diff --git a/internal/metadata/featureflag/feature_flags.go b/internal/metadata/featureflag/feature_flags.go
index 2b075bc91..1580ef556 100644
--- a/internal/metadata/featureflag/feature_flags.go
+++ b/internal/metadata/featureflag/feature_flags.go
@@ -14,12 +14,19 @@ var (
 	RemoteBranchesLsRemote = FeatureFlag{Name: "ruby_remote_branches_ls_remote", OnByDefault: true}
 	// GoFetchSourceBranch enables a go implementation of FetchSourceBranch
 	GoFetchSourceBranch = FeatureFlag{Name: "go_fetch_source_branch", OnByDefault: false}
-	// ReferenceTransactions will handle Git reference updates via the transaction service for strong consistency
-	ReferenceTransactions = FeatureFlag{Name: "reference_transactions", OnByDefault: false}
 	// DistributedReads allows praefect to redirect accessor operations to up-to-date secondaries
 	DistributedReads = FeatureFlag{Name: "distributed_reads", OnByDefault: false}
 	// GoPreReceiveHook will bypass the ruby pre-receive hook and use the go implementation
 	GoPreReceiveHook = FeatureFlag{Name: "go_prereceive_hook", OnByDefault: true}
+
+	// ReferenceTransactions will handle Git reference updates via the transaction service for strong consistency
+	ReferenceTransactions = FeatureFlag{Name: "reference_transactions", OnByDefault: false}
+	// ReferenceTransactionsOperationService will enable reference transactions for the OperationService
+	ReferenceTransactionsOperationService = FeatureFlag{Name: "reference_transactions_operation_service", OnByDefault: true}
+	// ReferenceTransactionsSmartHTTPService will enable reference transactions for the SmartHTTPService
+	ReferenceTransactionsSmartHTTPService = FeatureFlag{Name: "reference_transactions_smarthttp_service", OnByDefault: true}
+	// ReferenceTransactionsSSHService will enable reference transactions for the SSHService
+	ReferenceTransactionsSSHService = FeatureFlag{Name: "reference_transactions_ssh_service", OnByDefault: true}
 )
 
 const (
diff --git a/internal/praefect/coordinator.go b/internal/praefect/coordinator.go
index 38c16c282..274ea6419 100644
--- a/internal/praefect/coordinator.go
+++ b/internal/praefect/coordinator.go
@@ -157,9 +157,27 @@ func (c *Coordinator) accessorStreamParameters(ctx context.Context, call grpcCal
 	}, nil, nil, nil), nil
 }
 
-var transactionRPCs = map[string]struct{}{
-	"/gitaly.SmartHTTPService/PostReceivePack": {},
-	"/gitaly.SSHService/SSHReceivePack":        {},
+var transactionRPCs = map[string]featureflag.FeatureFlag{
+	"/gitaly.OperationService/UserCreateBranch": featureflag.ReferenceTransactionsOperationService,
+	"/gitaly.SSHService/SSHReceivePack":         featureflag.ReferenceTransactionsSSHService,
+	"/gitaly.SmartHTTPService/PostReceivePack":  featureflag.ReferenceTransactionsSmartHTTPService,
+}
+
+func shouldUseTransaction(ctx context.Context, method string) bool {
+	if !featureflag.IsEnabled(ctx, featureflag.ReferenceTransactions) {
+		return false
+	}
+
+	flag, ok := transactionRPCs[method]
+	if !ok {
+		return false
+	}
+
+	if !featureflag.IsEnabled(ctx, flag) {
+		return false
+	}
+
+	return true
 }
 
 func (c *Coordinator) mutatorStreamParameters(ctx context.Context, call grpcCall, targetRepo *gitalypb.Repository) (*proxy.StreamParameters, error) {
@@ -194,7 +212,7 @@ func (c *Coordinator) mutatorStreamParameters(ctx context.Context, call grpcCall
 
 	var secondaryDests []proxy.Destination
 
-	if _, ok := transactionRPCs[call.fullMethodName]; ok && featureflag.IsEnabled(ctx, featureflag.ReferenceTransactions) {
+	if shouldUseTransaction(ctx, call.fullMethodName) {
 		// Make sure to only let healthy nodes take part in transactions, otherwise we'll be
 		// completely blocked until they come back.
 		healthySecondaries := shard.GetHealthySecondaries()
diff --git a/internal/rubyserver/proxy.go b/internal/rubyserver/proxy.go
index 87052c900..f0b8e3f9b 100644
--- a/internal/rubyserver/proxy.go
+++ b/internal/rubyserver/proxy.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
+	praefect_metadata "gitlab.com/gitlab-org/gitaly/internal/praefect/metadata"
 	"gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 	"google.golang.org/grpc/metadata"
 )
@@ -62,7 +63,11 @@ func setHeaders(ctx context.Context, repo *gitalypb.Repository, mustExist bool)
 	)
 
 	// list of http/2 headers that will be forwarded as-is to gitaly-ruby
-	proxyHeaderWhitelist := []string{"gitaly-servers"}
+	proxyHeaderWhitelist := []string{
+		"gitaly-servers",
+		praefect_metadata.TransactionMetadataKey,
+		praefect_metadata.PraefectMetadataKey,
+	}
 
 	if inMD, ok := metadata.FromIncomingContext(ctx); ok {
 		// Automatically whitelist any Ruby-specific feature flag
diff --git a/internal/service/operations/branches_test.go b/internal/service/operations/branches_test.go
index e290e3948..cbc90472d 100644
--- a/internal/service/operations/branches_test.go
+++ b/internal/service/operations/branches_test.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/internal/git/log"
+	"gitlab.com/gitlab-org/gitaly/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
@@ -75,10 +76,17 @@ func TestSuccessfulCreateBranchRequest(t *testing.T) {
 }
 
 func TestSuccessfulGitHooksForUserCreateBranchRequest(t *testing.T) {
-	ctx, cancel := testhelper.Context()
-	defer cancel()
+	featureSets, err := testhelper.NewFeatureSets([]featureflag.FeatureFlag{featureflag.ReferenceTransactions})
+	require.NoError(t, err)
 
-	testSuccessfulGitHooksForUserCreateBranchRequest(t, ctx)
+	for _, featureSet := range featureSets {
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		ctx = featureSet.WithParent(ctx)
+
+		testSuccessfulGitHooksForUserCreateBranchRequest(t, ctx)
+	}
 }
 
 func testSuccessfulGitHooksForUserCreateBranchRequest(t *testing.T, ctx context.Context) {
diff --git a/ruby/lib/gitaly_server.rb b/ruby/lib/gitaly_server.rb
index b6a085f57..a3e68fbff 100644
--- a/ruby/lib/gitaly_server.rb
+++ b/ruby/lib/gitaly_server.rb
@@ -14,6 +14,8 @@ require_relative 'gitaly_server/remote_service.rb'
 require_relative 'gitaly_server/health_service.rb'
 require_relative 'gitaly_server/feature_flags.rb'
 
+require_relative 'praefect/transaction.rb'
+
 module GitalyServer
   STORAGE_PATH_HEADER = 'gitaly-storage-path'.freeze
   REPO_PATH_HEADER = 'gitaly-repo-path'.freeze
diff --git a/ruby/lib/gitaly_server/operations_service.rb b/ruby/lib/gitaly_server/operations_service.rb
index d4aebac0c..0f4db6c7d 100644
--- a/ruby/lib/gitaly_server/operations_service.rb
+++ b/ruby/lib/gitaly_server/operations_service.rb
@@ -49,10 +49,11 @@ module GitalyServer
       repo = Gitlab::Git::Repository.from_gitaly(request.repository, call)
       target = get_param!(request, :start_point)
       gitaly_user = get_param!(request, :user)
+      transaction = Praefect::Transaction.from_metadata(call.metadata)
 
       branch_name = request.branch_name
       user = Gitlab::Git::User.from_gitaly(gitaly_user)
-      created_branch = repo.add_branch(branch_name, user: user, target: target)
+      created_branch = repo.add_branch(branch_name, user: user, target: target, transaction: transaction)
       Gitaly::UserCreateBranchResponse.new unless created_branch
 
       rugged_commit = created_branch.dereferenced_target.rugged_commit
diff --git a/ruby/lib/gitlab/git/hook.rb b/ruby/lib/gitlab/git/hook.rb
index 78fa009fd..9fb54506e 100644
--- a/ruby/lib/gitlab/git/hook.rb
+++ b/ruby/lib/gitlab/git/hook.rb
@@ -28,13 +28,13 @@ module Gitlab
         File.exist?(path)
       end
 
-      def trigger(gl_id, gl_username, oldrev, newrev, ref, push_options: nil)
+      def trigger(gl_id, gl_username, oldrev, newrev, ref, push_options: nil, transaction: nil)
         return [true, nil] unless exists?
 
         Bundler.with_clean_env do
           case name
           when "pre-receive", "post-receive"
-            call_receive_hook(gl_id, gl_username, oldrev, newrev, ref, push_options)
+            call_receive_hook(gl_id, gl_username, oldrev, newrev, ref, push_options, transaction)
           when "update"
             call_update_hook(gl_id, gl_username, oldrev, newrev, ref)
           end
@@ -43,7 +43,7 @@ module Gitlab
 
       private
 
-      def call_receive_hook(gl_id, gl_username, oldrev, newrev, ref, push_options)
+      def call_receive_hook(gl_id, gl_username, oldrev, newrev, ref, push_options, transaction)
         changes = [oldrev, newrev, ref].join(" ")
 
         exit_status = false
@@ -51,6 +51,7 @@ module Gitlab
 
         vars = env_base_vars(gl_id, gl_username)
         vars.merge!(push_options.env_data) if push_options
+        vars.merge!(transaction.env_vars) if transaction
 
         options = {
           chdir: repo_path
diff --git a/ruby/lib/gitlab/git/hooks_service.rb b/ruby/lib/gitlab/git/hooks_service.rb
index 30501d388..5ad8e43b7 100644
--- a/ruby/lib/gitlab/git/hooks_service.rb
+++ b/ruby/lib/gitlab/git/hooks_service.rb
@@ -3,7 +3,7 @@ module Gitlab
     class HooksService
       attr_accessor :oldrev, :newrev, :ref
 
-      def execute(pusher, repository, oldrev, newrev, ref, push_options:)
+      def execute(pusher, repository, oldrev, newrev, ref, push_options:, transaction: nil)
         @repository  = repository
         @gl_id       = pusher.gl_id
         @gl_username = pusher.username
@@ -11,6 +11,7 @@ module Gitlab
         @newrev      = newrev
         @ref         = ref
         @push_options = push_options
+        @transaction = transaction
 
         %w[pre-receive update].each do |hook_name|
           status, message = run_hook(hook_name)
@@ -29,7 +30,7 @@ module Gitlab
 
       def run_hook(name)
         hook = Gitlab::Git::Hook.new(name, @repository)
-        hook.trigger(@gl_id, @gl_username, oldrev, newrev, ref, push_options: @push_options)
+        hook.trigger(@gl_id, @gl_username, oldrev, newrev, ref, push_options: @push_options, transaction: @transaction)
       end
     end
   end
diff --git a/ruby/lib/gitlab/git/operation_service.rb b/ruby/lib/gitlab/git/operation_service.rb
index 5be4d27e0..e6025b673 100644
--- a/ruby/lib/gitlab/git/operation_service.rb
+++ b/ruby/lib/gitlab/git/operation_service.rb
@@ -25,11 +25,11 @@ module Gitlab
         @repository = new_repository
       end
 
-      def add_branch(branch_name, newrev)
+      def add_branch(branch_name, newrev, transaction: nil)
         ref = Gitlab::Git::BRANCH_REF_PREFIX + branch_name
         oldrev = Gitlab::Git::BLANK_SHA
 
-        update_ref_in_hooks(ref, newrev, oldrev)
+        update_ref_in_hooks(ref, newrev, oldrev, transaction: transaction)
       end
 
       def rm_branch(branch)
@@ -172,20 +172,21 @@ module Gitlab
         end
       end
 
-      def update_ref_in_hooks(ref, newrev, oldrev, push_options: nil)
-        with_hooks(ref, newrev, oldrev, push_options: push_options) do
+      def update_ref_in_hooks(ref, newrev, oldrev, push_options: nil, transaction: nil)
+        with_hooks(ref, newrev, oldrev, push_options: push_options, transaction: transaction) do
           update_ref(ref, newrev, oldrev)
         end
       end
 
-      def with_hooks(ref, newrev, oldrev, push_options: nil)
+      def with_hooks(ref, newrev, oldrev, push_options: nil, transaction: nil)
         Gitlab::Git::HooksService.new.execute(
           user,
           repository,
           oldrev,
           newrev,
           ref,
-          push_options: push_options
+          push_options: push_options,
+          transaction: transaction
         ) do |service|
           yield(service)
         end
diff --git a/ruby/lib/gitlab/git/repository.rb b/ruby/lib/gitlab/git/repository.rb
index 28502287c..675f4b2e2 100644
--- a/ruby/lib/gitlab/git/repository.rb
+++ b/ruby/lib/gitlab/git/repository.rb
@@ -88,11 +88,11 @@ module Gitlab
         @feature_flags.enabled?(flag, on_by_default: on_by_default)
       end
 
-      def add_branch(branch_name, user:, target:)
+      def add_branch(branch_name, user:, target:, transaction: nil)
         target_object = Ref.dereference_object(lookup(target))
         raise InvalidRef, "target not found: #{target}" unless target_object
 
-        OperationService.new(user, self).add_branch(branch_name, target_object.oid)
+        OperationService.new(user, self).add_branch(branch_name, target_object.oid, transaction: transaction)
         find_branch(branch_name)
       rescue Rugged::ReferenceError => ex
         raise InvalidRef, ex
diff --git a/ruby/lib/praefect/transaction.rb b/ruby/lib/praefect/transaction.rb
new file mode 100644
index 000000000..3ddf6c5f6
--- /dev/null
+++ b/ruby/lib/praefect/transaction.rb
@@ -0,0 +1,24 @@
+module Praefect
+  class Transaction
+    PRAEFECT_SERVER_KEY = "praefect-server".freeze
+    PRAEFECT_SERVER_ENV = "PRAEFECT_SERVER".freeze
+    TRANSACTION_KEY = "transaction".freeze
+    TRANSACTION_ENV = "REFERENCE_TRANSACTION".freeze
+
+    def self.from_metadata(metadata)
+      new(metadata[PRAEFECT_SERVER_KEY], metadata[TRANSACTION_KEY])
+    end
+
+    def initialize(server, transaction)
+      @server = server
+      @transaction = transaction
+    end
+
+    def env_vars
+      {
+        TRANSACTION_ENV => @transaction,
+        PRAEFECT_SERVER_ENV => @server
+      }.reject { |_, v| v.nil? }
+    end
+  end
+end