Add latest changes from gitlab-org/gitlab@master

2 files changed, 39 insertions, 14 deletions

diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index b0f29d22ad4..9e3af00e00d 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -142,18 +142,39 @@ module Gitlab
     #   kwargs.merge(deadline: Time.now + 10)
     # end
     #
-    def self.call(storage, service, rpc, request, remote_storage: nil, timeout: default_timeout)
-      start = Gitlab::Metrics::System.monotonic_time
-      request_hash = request.is_a?(Google::Protobuf::MessageExts) ? request.to_h : {}
+    def self.call(storage, service, rpc, request, remote_storage: nil, timeout: default_timeout, &block)
+      self.measure_timings(service, rpc, request) do
+        self.execute(storage, service, rpc, request, remote_storage: remote_storage, timeout: timeout, &block)
+      end
+    end
 
+    # This method is like GitalyClient.call but should be used with
+    # Gitaly streaming RPCs. It measures how long the the RPC took to
+    # produce the full response, not just the initial response.
+    def self.streaming_call(storage, service, rpc, request, remote_storage: nil, timeout: default_timeout)
+      self.measure_timings(service, rpc, request) do
+        response = self.execute(storage, service, rpc, request, remote_storage: remote_storage, timeout: timeout)
+
+        yield(response)
+      end
+    end
+
+    def self.execute(storage, service, rpc, request, remote_storage:, timeout:)
       enforce_gitaly_request_limits(:call)
 
       kwargs = request_kwargs(storage, timeout: timeout.to_f, remote_storage: remote_storage)
       kwargs = yield(kwargs) if block_given?
 
       stub(service, storage).__send__(rpc, request, kwargs) # rubocop:disable GitlabSecurity/PublicSend
+    end
+
+    def self.measure_timings(service, rpc, request)
+      start = Gitlab::Metrics::System.monotonic_time
+
+      yield
     ensure
       duration = Gitlab::Metrics::System.monotonic_time - start
+      request_hash = request.is_a?(Google::Protobuf::MessageExts) ? request.to_h : {}
 
       # Keep track, separately, for the performance bar
       self.query_time += duration
diff --git a/lib/gitlab/gitaly_client/commit_service.rb b/lib/gitlab/gitaly_client/commit_service.rb
index dca55091be6..15318bc817a 100644
--- a/lib/gitlab/gitaly_client/commit_service.rb
+++ b/lib/gitlab/gitaly_client/commit_service.rb
@@ -200,8 +200,9 @@ module Gitlab
           to: to
         )
 
-        response = GitalyClient.call(@repository.storage, :commit_service, :commits_between, request, timeout: GitalyClient.medium_timeout)
-        consume_commits_response(response)
+        GitalyClient.streaming_call(@repository.storage, :commit_service, :commits_between, request, timeout: GitalyClient.medium_timeout) do |response|
+          consume_commits_response(response)
+        end
       end
 
       def diff_stats(left_commit_sha, right_commit_sha)
@@ -224,8 +225,9 @@ module Gitlab
         )
         request.order = opts[:order].upcase if opts[:order].present?
 
-        response = GitalyClient.call(@repository.storage, :commit_service, :find_all_commits, request, timeout: GitalyClient.medium_timeout)
-        consume_commits_response(response)
+        GitalyClient.streaming_call(@repository.storage, :commit_service, :find_all_commits, request, timeout: GitalyClient.medium_timeout) do |response|
+          consume_commits_response(response)
+        end
       end
 
       def list_commits_by_oid(oids)
@@ -233,8 +235,9 @@ module Gitlab
 
         request = Gitaly::ListCommitsByOidRequest.new(repository: @gitaly_repo, oid: oids)
 
-        response = GitalyClient.call(@repository.storage, :commit_service, :list_commits_by_oid, request, timeout: GitalyClient.medium_timeout)
-        consume_commits_response(response)
+        GitalyClient.streaming_call(@repository.storage, :commit_service, :list_commits_by_oid, request, timeout: GitalyClient.medium_timeout) do |response|
+          consume_commits_response(response)
+        end
       rescue GRPC::NotFound # If no repository is found, happens mainly during testing
         []
       end
@@ -249,8 +252,9 @@ module Gitlab
           offset: offset.to_i
         )
 
-        response = GitalyClient.call(@repository.storage, :commit_service, :commits_by_message, request, timeout: GitalyClient.medium_timeout)
-        consume_commits_response(response)
+        GitalyClient.streaming_call(@repository.storage, :commit_service, :commits_by_message, request, timeout: GitalyClient.medium_timeout) do |response|
+          consume_commits_response(response)
+        end
       end
 
       def languages(ref = nil)
@@ -323,9 +327,9 @@ module Gitlab
 
         request.paths = encode_repeated(Array(options[:path])) if options[:path].present?
 
-        response = GitalyClient.call(@repository.storage, :commit_service, :find_commits, request, timeout: GitalyClient.medium_timeout)
-
-        consume_commits_response(response)
+        GitalyClient.streaming_call(@repository.storage, :commit_service, :find_commits, request, timeout: GitalyClient.medium_timeout) do |response|
+          consume_commits_response(response)
+        end
       end
 
       def filter_shas_with_signatures(shas)