From d86751d842f9547af0b7db17eacb89e940a106f7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alejandro=20Rodr=C3=ADguez?= <alejorro70@gmail.com>
Date: Wed, 17 Jan 2018 20:25:16 -0300
Subject: [PATCH] Incorporate Gitaly's RemoteService.UpdateRemoteMirror RPC

---
 GITALY_SERVER_VERSION                         |  2 +-
 lib/gitlab/git/remote_mirror.rb               | 24 +++++++++++++----
 lib/gitlab/gitaly_client/remote_service.rb    | 27 +++++++++++++++++++
 .../gitaly_client/remote_service_spec.rb      | 14 ++++++++++
 4 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/GITALY_SERVER_VERSION b/GITALY_SERVER_VERSION
index 534b316aef6..7375dee5f49 100644
--- a/GITALY_SERVER_VERSION
+++ b/GITALY_SERVER_VERSION
@@ -1 +1 @@
-0.70.0
+0.72.0
diff --git a/lib/gitlab/git/remote_mirror.rb b/lib/gitlab/git/remote_mirror.rb
index 38e9d2a8554..ebe46722890 100644
--- a/lib/gitlab/git/remote_mirror.rb
+++ b/lib/gitlab/git/remote_mirror.rb
@@ -6,7 +6,23 @@ module Gitlab
         @ref_name = ref_name
       end
 
-      def update(only_branches_matching: [], only_tags_matching: [])
+      def update(only_branches_matching: [])
+        @repository.gitaly_migrate(:remote_update_remote_mirror) do |is_enabled|
+          if is_enabled
+            gitaly_update(only_branches_matching)
+          else
+            rugged_update(only_branches_matching)
+          end
+        end
+      end
+
+      private
+
+      def gitaly_update(only_branches_matching)
+        @repository.gitaly_remote_client.update_remote_mirror(@ref_name, only_branches_matching)
+      end
+
+      def rugged_update(only_branches_matching)
         local_branches = refs_obj(@repository.local_branches, only_refs_matching: only_branches_matching)
         remote_branches = refs_obj(@repository.remote_branches(@ref_name), only_refs_matching: only_branches_matching)
 
@@ -15,8 +31,8 @@ module Gitlab
 
         delete_refs(local_branches, remote_branches)
 
-        local_tags = refs_obj(@repository.tags, only_refs_matching: only_tags_matching)
-        remote_tags = refs_obj(@repository.remote_tags(@ref_name), only_refs_matching: only_tags_matching)
+        local_tags = refs_obj(@repository.tags)
+        remote_tags = refs_obj(@repository.remote_tags(@ref_name))
 
         updated_tags = changed_refs(local_tags, remote_tags)
         @repository.push_remote_branches(@ref_name, updated_tags.keys) if updated_tags.present?
@@ -24,8 +40,6 @@ module Gitlab
         delete_refs(local_tags, remote_tags)
       end
 
-      private
-
       def refs_obj(refs, only_refs_matching: [])
         refs.each_with_object({}) do |ref, refs|
           next if only_refs_matching.present? && !only_refs_matching.include?(ref.name)
diff --git a/lib/gitlab/gitaly_client/remote_service.rb b/lib/gitlab/gitaly_client/remote_service.rb
index e58f641d69a..58c356edfd1 100644
--- a/lib/gitlab/gitaly_client/remote_service.rb
+++ b/lib/gitlab/gitaly_client/remote_service.rb
@@ -1,6 +1,8 @@
 module Gitlab
   module GitalyClient
     class RemoteService
+      MAX_MSG_SIZE = 128.kilobytes.freeze
+
       def initialize(repository)
         @repository = repository
         @gitaly_repo = repository.gitaly_repository
@@ -38,6 +40,31 @@ module Gitlab
 
         response.result
       end
+
+      def update_remote_mirror(ref_name, only_branches_matching)
+        req_enum = Enumerator.new do |y|
+          y.yield Gitaly::UpdateRemoteMirrorRequest.new(
+            repository: @gitaly_repo,
+            ref_name: ref_name
+          )
+
+          current_size = 0
+
+          slices = only_branches_matching.slice_before do |branch_name|
+            current_size += branch_name.bytesize
+
+            next false if current_size < MAX_MSG_SIZE
+
+            current_size = 0
+          end
+
+          slices.each do |slice|
+            y.yield Gitaly::UpdateRemoteMirrorRequest.new(only_branches_matching: slice)
+          end
+        end
+
+        GitalyClient.call(@storage, :remote_service, :update_remote_mirror, req_enum)
+      end
     end
   end
 end
diff --git a/spec/lib/gitlab/gitaly_client/remote_service_spec.rb b/spec/lib/gitlab/gitaly_client/remote_service_spec.rb
index 9d540446532..872377c93d8 100644
--- a/spec/lib/gitlab/gitaly_client/remote_service_spec.rb
+++ b/spec/lib/gitlab/gitaly_client/remote_service_spec.rb
@@ -44,4 +44,18 @@ describe Gitlab::GitalyClient::RemoteService do
       expect(client.fetch_internal_remote(remote_repository)).to be(true)
     end
   end
+
+  describe '#update_remote_mirror' do
+    let(:ref_name) { 'remote_mirror_1' }
+    let(:only_branches_matching) { ['my-branch', 'master'] }
+
+    it 'sends an update_remote_mirror message' do
+      expect_any_instance_of(Gitaly::RemoteService::Stub)
+        .to receive(:update_remote_mirror)
+        .with(kind_of(Enumerator), kind_of(Hash))
+        .and_return(double(:update_remote_mirror_response))
+
+      client.update_remote_mirror(ref_name, only_branches_matching)
+    end
+  end
 end