Add FetchSourceBranch Gitaly call

4 files changed, 45 insertions, 8 deletions

diff --git a/lib/gitlab/git/repository.rb b/lib/gitlab/git/repository.rb
index 3cb9b254e6e..dcca20c75ef 100644
--- a/lib/gitlab/git/repository.rb
+++ b/lib/gitlab/git/repository.rb
@@ -1058,12 +1058,11 @@ module Gitlab
       end
 
       def fetch_source_branch!(source_repository, source_branch, local_ref)
-        with_repo_branch_commit(source_repository, source_branch) do |commit|
-          if commit
-            write_ref(local_ref, commit.sha)
-            true
+        Gitlab::GitalyClient.migrate(:fetch_source_branch) do |is_enabled|
+          if is_enabled
+            gitaly_repository_client.fetch_source_branch(source_repository, source_branch, local_ref)
           else
-            false
+            rugged_fetch_source_branch(source_repository, source_branch, local_ref)
           end
         end
       end
@@ -1216,6 +1215,17 @@ module Gitlab
 
       private
 
+      def rugged_fetch_source_branch(source_repository, source_branch, local_ref)
+        with_repo_branch_commit(source_repository, source_branch) do |commit|
+          if commit
+            write_ref(local_ref, commit.sha)
+            true
+          else
+            false
+          end
+        end
+      end
+
       # Gitaly note: JV: Trying to get rid of the 'filter' option so we can implement this with 'git'.
       def branches_filter(filter: nil, sort_by: nil)
         # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37464
diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index 0b35a787e07..572f4c892f6 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -75,6 +75,10 @@ module Gitlab
       address
     end
 
+    def self.address_metadata(storage)
+      Base64.strict_encode64(JSON.dump({ storage => { 'address' => address(storage), 'token' => token(storage) } }))
+    end
+
     # All Gitaly RPC call sites should use GitalyClient.call. This method
     # makes sure that per-request authentication headers are set.
     #
@@ -89,18 +93,19 @@ module Gitlab
     #   kwargs.merge(deadline: Time.now + 10)
     # end
     #
-    def self.call(storage, service, rpc, request)
+    def self.call(storage, service, rpc, request, remote_storage: nil)
       start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
       enforce_gitaly_request_limits(:call)
 
-      kwargs = request_kwargs(storage)
+      kwargs = request_kwargs(storage, remote_storage: remote_storage)
       kwargs = yield(kwargs) if block_given?
+
       stub(service, storage).__send__(rpc, request, kwargs) # rubocop:disable GitlabSecurity/PublicSend
     ensure
       self.query_time += Process.clock_gettime(Process::CLOCK_MONOTONIC) - start
     end
 
-    def self.request_kwargs(storage)
+    def self.request_kwargs(storage, remote_storage: nil)
       encoded_token = Base64.strict_encode64(token(storage).to_s)
       metadata = {
         'authorization' => "Bearer #{encoded_token}",
@@ -110,6 +115,7 @@ module Gitlab
       feature_stack = Thread.current[:gitaly_feature_stack]
       feature = feature_stack && feature_stack[0]
       metadata['call_site'] = feature.to_s if feature
+      metadata['gitaly-servers'] = address_metadata(remote_storage) if remote_storage
 
       { metadata: metadata }
     end
diff --git a/lib/gitlab/gitaly_client/repository_service.rb b/lib/gitlab/gitaly_client/repository_service.rb
index cef692d3c2a..70cb16bd810 100644
--- a/lib/gitlab/gitaly_client/repository_service.rb
+++ b/lib/gitlab/gitaly_client/repository_service.rb
@@ -65,6 +65,25 @@ module Gitlab
 
         response.value
       end
+
+      def fetch_source_branch(source_repository, source_branch, local_ref)
+        request = Gitaly::FetchSourceBranchRequest.new(
+          repository: @gitaly_repo,
+          source_repository: source_repository.gitaly_repository,
+          source_branch: source_branch.b,
+          target_ref: local_ref.b
+        )
+
+        response = GitalyClient.call(
+          @storage,
+          :repository_service,
+          :fetch_source_branch,
+          request,
+          remote_storage: source_repository.storage
+        )
+
+        response.result
+      end
     end
   end
 end
diff --git a/lib/tasks/gitlab/gitaly.rake b/lib/tasks/gitlab/gitaly.rake
index f2002d7a426..4d880c05f99 100644
--- a/lib/tasks/gitlab/gitaly.rake
+++ b/lib/tasks/gitlab/gitaly.rake
@@ -78,6 +78,8 @@ namespace :gitlab do
       config[:auth] = { token: 'secret' } if Rails.env.test?
       config[:'gitaly-ruby'] = { dir: File.join(Dir.pwd, 'ruby') } if gitaly_ruby
       config[:'gitlab-shell'] = { dir: Gitlab.config.gitlab_shell.path }
+      config[:bin_dir] = Gitlab.config.gitaly.client_path
+
       TOML.dump(config)
     end