From 10fb904d636ab49b87e2deb756f64fc26d2eacec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kim=20=22BKC=22=20Carlb=C3=A4cker?=
 <kim.carlbacker@gmail.com>
Date: Wed, 17 Jan 2018 06:43:30 +0100
Subject: [PATCH] Migrate Git::Repository#write_ref to Gitaly

---
 app/models/repository.rb                      |  2 ++
 lib/gitlab/git/repository.rb                  | 20 +++++++++++++++----
 .../gitaly_client/repository_service.rb       | 16 +++++++++++++++
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/app/models/repository.rb b/app/models/repository.rb
index 5b06dc5a39b..2146e5a74f2 100644
--- a/app/models/repository.rb
+++ b/app/models/repository.rb
@@ -255,6 +255,8 @@ class Repository
 
     # This will still fail if the file is corrupted (e.g. 0 bytes)
     raw_repository.write_ref(keep_around_ref_name(sha), sha, shell: false)
+  rescue Gitlab::Git::CommandError => ex
+    Rails.logger.error "Unable to create keep-around reference for repository #{path}: #{ex}"
   end
 
   def kept_around?(sha)
diff --git a/lib/gitlab/git/repository.rb b/lib/gitlab/git/repository.rb
index 3a7930154e5..e65719e7f92 100644
--- a/lib/gitlab/git/repository.rb
+++ b/lib/gitlab/git/repository.rb
@@ -1105,10 +1105,14 @@ module Gitlab
       end
 
       def write_ref(ref_path, ref, old_ref: nil, shell: true)
-        if shell
-          shell_write_ref(ref_path, ref, old_ref)
-        else
-          rugged_write_ref(ref_path, ref)
+        ref_path = "#{Gitlab::Git::BRANCH_REF_PREFIX}#{ref_path}" unless ref_path.start_with?("refs/") || ref_path == "HEAD"
+
+        gitaly_migrate(:write_ref) do |is_enabled|
+          if is_enabled
+            gitaly_repository_client.write_ref(ref_path, ref, old_ref, shell)
+          else
+            local_write_ref(ref_path, ref, old_ref: old_ref, shell: shell)
+          end
         end
       end
 
@@ -1419,6 +1423,14 @@ module Gitlab
 
       private
 
+      def local_write_ref(ref_path, ref, old_ref: nil, shell: true)
+        if shell
+          shell_write_ref(ref_path, ref, old_ref)
+        else
+          rugged_write_ref(ref_path, ref)
+        end
+      end
+
       def shell_write_ref(ref_path, ref, old_ref)
         raise ArgumentError, "invalid ref_path #{ref_path.inspect}" if ref_path.include?(' ')
         raise ArgumentError, "invalid ref #{ref.inspect}" if ref.include?("\x00")
diff --git a/lib/gitlab/gitaly_client/repository_service.rb b/lib/gitlab/gitaly_client/repository_service.rb
index b0dbaf11598..7adf32af209 100644
--- a/lib/gitlab/gitaly_client/repository_service.rb
+++ b/lib/gitlab/gitaly_client/repository_service.rb
@@ -203,6 +203,22 @@ module Gitlab
           timeout: GitalyClient.default_timeout
         )
       end
+
+      def write_ref(ref_path, ref, old_ref, shell)
+        request = Gitaly::WriteRefRequest.new(
+          repository: @gitaly_repo,
+          ref: ref_path.b,
+          revision: ref.b,
+          shell: shell
+        )
+        request.old_revision = old_ref.b unless old_ref.nil?
+
+        response = GitalyClient.call(@storage, :repository_service, :write_ref, request)
+
+        raise Gitlab::Git::CommandError, encode!(response.error) if response.error.present?
+
+        true
+      end
     end
   end
 end