From 571f1dda17203d162dac65270a6b68366de1f101 Mon Sep 17 00:00:00 2001
From: "Jacob Vosmaer (GitLab)" <jacob@gitlab.com>
Date: Wed, 22 Nov 2017 10:19:42 +0000
Subject: [PATCH] Add FetchSourceBranch Gitaly call

---
 GITALY_SERVER_VERSION                         |  2 +-
 lib/gitlab/git/repository.rb                  | 20 ++++++++++++++-----
 lib/gitlab/gitaly_client.rb                   | 12 ++++++++---
 .../gitaly_client/repository_service.rb       | 19 ++++++++++++++++++
 lib/tasks/gitlab/gitaly.rake                  |  2 ++
 spec/requests/api/merge_requests_spec.rb      |  2 +-
 spec/requests/api/v3/merge_requests_spec.rb   |  2 +-
 spec/tasks/gitlab/gitaly_rake_spec.rb         |  1 +
 8 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/GITALY_SERVER_VERSION b/GITALY_SERVER_VERSION
index 524456c7767..316ba4bd9e6 100644
--- a/GITALY_SERVER_VERSION
+++ b/GITALY_SERVER_VERSION
@@ -1 +1 @@
-0.54.0
+0.55.0
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
 
diff --git a/spec/requests/api/merge_requests_spec.rb b/spec/requests/api/merge_requests_spec.rb
index a928ba79a4d..35c6b3bb2fb 100644
--- a/spec/requests/api/merge_requests_spec.rb
+++ b/spec/requests/api/merge_requests_spec.rb
@@ -628,7 +628,7 @@ describe API::MergeRequests do
 
     context 'forked projects' do
       let!(:user2) { create(:user) }
-      let!(:forked_project) { fork_project(project, user2) }
+      let!(:forked_project) { fork_project(project, user2, repository: true) }
       let!(:unrelated_project) { create(:project,  namespace: create(:user).namespace, creator_id: user2.id) }
 
       before do
diff --git a/spec/requests/api/v3/merge_requests_spec.rb b/spec/requests/api/v3/merge_requests_spec.rb
index 91897e5ee01..2e2b9449429 100644
--- a/spec/requests/api/v3/merge_requests_spec.rb
+++ b/spec/requests/api/v3/merge_requests_spec.rb
@@ -314,7 +314,7 @@ describe API::MergeRequests do
 
     context 'forked projects' do
       let!(:user2) { create(:user) }
-      let!(:forked_project) { fork_project(project, user2) }
+      let!(:forked_project) { fork_project(project, user2, repository: true) }
       let!(:unrelated_project) { create(:project,  namespace: create(:user).namespace, creator_id: user2.id) }
 
       before do
diff --git a/spec/tasks/gitlab/gitaly_rake_spec.rb b/spec/tasks/gitlab/gitaly_rake_spec.rb
index a51374e2645..6aba86fdc3c 100644
--- a/spec/tasks/gitlab/gitaly_rake_spec.rb
+++ b/spec/tasks/gitlab/gitaly_rake_spec.rb
@@ -112,6 +112,7 @@ describe 'gitlab:gitaly namespace rake task' do
         expected_output = <<~TOML
           # Gitaly storage configuration generated from #{Gitlab.config.source} on #{Time.current.to_s(:long)}
           # This is in TOML format suitable for use in Gitaly's config.toml file.
+          bin_dir = "tmp/tests/gitaly"
           socket_path = "/path/to/my.socket"
           [gitlab-shell]
           dir = "#{Gitlab.config.gitlab_shell.path}"
-- 
2.30.9