Fix Gitaly duration timings of RefService RPCs

For many Gitaly RPCs, previously the `gitaly_duration_s` log timings
only accounted for the initial request/response time.  We now measure
the total time it takes to consume the response for the following RPCs:

1. FindAllBranches
2. FindAllRemoteBranches
3. FindAllBranchNames
4. FindAllTagNames
5. ListNewCommits
6. ListNewBlobs
7. FindLocalBranches
8. FindAllTags
9. ListBranchNamesContainingCommit
10. GetTagMessages

This came up in https://gitlab.com/gitlab-org/gitlab/-/issues/222247.

Part of https://gitlab.com/gitlab-org/gitlab/-/issues/30334

changelogs/unreleased/sh-fix-gitaly-ref-service-durations.yml

+---
+title: Fix Gitaly duration tracking of RefService RPCs
+merge_request: 34904
+author:
+type: other

lib/gitlab/gitaly_client/ref_service.rb

@@ -14,17 +14,17 @@ module Gitlab
 
       def branches
         request = Gitaly::FindAllBranchesRequest.new(repository: @gitaly_repo)
-        response = GitalyClient.call(@storage, :ref_service, :find_all_branches, request, timeout: GitalyClient.fast_timeout)
-
+        GitalyClient.streaming_call(@storage, :ref_service, :find_all_branches, request, timeout: GitalyClient.fast_timeout) do |response|
           consume_find_all_branches_response(response)
         end
+      end
 
       def remote_branches(remote_name)
         request = Gitaly::FindAllRemoteBranchesRequest.new(repository: @gitaly_repo, remote_name: remote_name)
-        response = GitalyClient.call(@repository.storage, :ref_service, :find_all_remote_branches, request, timeout: GitalyClient.medium_timeout)
-
+        GitalyClient.streaming_call(@storage, :ref_service, :find_all_remote_branches, request, timeout: GitalyClient.medium_timeout) do |response|
           consume_find_all_remote_branches_response(remote_name, response)
         end
+      end
 
       def merged_branches(branch_names = [])
         request = Gitaly::FindAllBranchesRequest.new(
@@ -32,10 +32,10 @@ module Gitlab
           merged_only: true,
           merged_branches: branch_names.map { |s| encode_binary(s) }
         )
-        response = GitalyClient.call(@storage, :ref_service, :find_all_branches, request, timeout: GitalyClient.fast_timeout)
-
+        GitalyClient.streaming_call(@storage, :ref_service, :find_all_branches, request, timeout: GitalyClient.fast_timeout) do |response|
           consume_find_all_branches_response(response)
         end
+      end
 
       def default_branch_name
         request = Gitaly::FindDefaultBranchNameRequest.new(repository: @gitaly_repo)
@@ -45,15 +45,17 @@ module Gitlab
 
       def branch_names
         request = Gitaly::FindAllBranchNamesRequest.new(repository: @gitaly_repo)
-        response = GitalyClient.call(@storage, :ref_service, :find_all_branch_names, request, timeout: GitalyClient.fast_timeout)
+        GitalyClient.streaming_call(@storage, :ref_service, :find_all_branch_names, request, timeout: GitalyClient.fast_timeout) do |response|
           consume_refs_response(response) { |name| Gitlab::Git.branch_name(name) }
         end
+      end
 
       def tag_names
         request = Gitaly::FindAllTagNamesRequest.new(repository: @gitaly_repo)
-        response = GitalyClient.call(@storage, :ref_service, :find_all_tag_names, request, timeout: GitalyClient.fast_timeout)
+        GitalyClient.streaming_call(@storage, :ref_service, :find_all_tag_names, request, timeout: GitalyClient.fast_timeout) do |response|
           consume_refs_response(response) { |name| Gitlab::Git.tag_name(name) }
         end
+      end
 
       def find_ref_name(commit_id, ref_prefix)
         request = Gitaly::FindRefNameRequest.new(
@@ -71,15 +73,15 @@ module Gitlab
           commit_id: newrev
         )
 
-        response = GitalyClient
-          .call(@storage, :ref_service, :list_new_commits, request, timeout: GitalyClient.medium_timeout)
-
         commits = []
+
+        GitalyClient.streaming_call(@storage, :ref_service, :list_new_commits, request, timeout: GitalyClient.medium_timeout) do |response|
           response.each do |msg|
             msg.commits.each do |c|
               commits << Gitlab::Git::Commit.new(@repository, c)
             end
           end
+        end
 
         commits
       end
@@ -98,15 +100,14 @@ module Gitlab
             GitalyClient.medium_timeout
           end
 
-        response = GitalyClient
-          .call(@storage, :ref_service, :list_new_blobs, request, timeout: timeout)
-
+        GitalyClient.streaming_call(@storage, :ref_service, :list_new_blobs, request, timeout: timeout) do |response|
           response.flat_map do |msg|
             # Returns an Array of Gitaly::NewBlobObject objects
             # Available methods are: #size, #oid and #path
             msg.new_blob_objects
           end
         end
+      end
 
       def count_tag_names
         tag_names.count
@@ -119,15 +120,17 @@ module Gitlab
       def local_branches(sort_by: nil)
         request = Gitaly::FindLocalBranchesRequest.new(repository: @gitaly_repo)
         request.sort_by = sort_by_param(sort_by) if sort_by
-        response = GitalyClient.call(@storage, :ref_service, :find_local_branches, request, timeout: GitalyClient.fast_timeout)
+        GitalyClient.streaming_call(@storage, :ref_service, :find_local_branches, request, timeout: GitalyClient.fast_timeout) do |response|
           consume_find_local_branches_response(response)
         end
+      end
 
       def tags
         request = Gitaly::FindAllTagsRequest.new(repository: @gitaly_repo)
-        response = GitalyClient.call(@storage, :ref_service, :find_all_tags, request, timeout: GitalyClient.medium_timeout)
+        GitalyClient.streaming_call(@storage, :ref_service, :find_all_tags, request, timeout: GitalyClient.medium_timeout) do |response|
           consume_tags_response(response)
         end
+      end
 
       def ref_exists?(ref_name)
         request = Gitaly::RefExistsRequest.new(repository: @gitaly_repo, ref: encode_binary(ref_name))
@@ -171,9 +174,9 @@ module Gitlab
           limit: limit
         )
 
-        stream = GitalyClient.call(@repository.storage, :ref_service, :list_tag_names_containing_commit, request, timeout: GitalyClient.medium_timeout)
-
-        consume_ref_contains_sha_response(stream, :tag_names)
+        GitalyClient.streaming_call(@storage, :ref_service, :list_tag_names_containing_commit, request, timeout: GitalyClient.medium_timeout) do |response|
+          consume_ref_contains_sha_response(response, :tag_names)
+        end
       end
 
       # Limit: 0 implies no limit, thus all tag names will be returned
@@ -184,23 +187,23 @@ module Gitlab
           limit: limit
         )
 
-        stream = GitalyClient.call(@repository.storage, :ref_service, :list_branch_names_containing_commit, request, timeout: GitalyClient.medium_timeout)
-
-        consume_ref_contains_sha_response(stream, :branch_names)
+        GitalyClient.streaming_call(@storage, :ref_service, :list_branch_names_containing_commit, request, timeout: GitalyClient.medium_timeout) do |response|
+          consume_ref_contains_sha_response(response, :branch_names)
+        end
       end
 
       def get_tag_messages(tag_ids)
         request = Gitaly::GetTagMessagesRequest.new(repository: @gitaly_repo, tag_ids: tag_ids)
-        response = GitalyClient.call(@repository.storage, :ref_service, :get_tag_messages, request, timeout: GitalyClient.fast_timeout)
-
         messages = Hash.new { |h, k| h[k] = +''.b }
         current_tag_id = nil
 
+        GitalyClient.streaming_call(@storage, :ref_service, :get_tag_messages, request, timeout: GitalyClient.fast_timeout) do |response|
           response.each do |rpc_message|
             current_tag_id = rpc_message.tag_id if rpc_message.tag_id.present?
 
             messages[current_tag_id] << rpc_message.message
           end
+        end
 
         messages
       end

spec/lib/gitlab/gitaly_client/ref_service_spec.rb

@@ -34,7 +34,7 @@ describe Gitlab::GitalyClient::RefService do
       subject
     end
 
-    it 'concantes and returns the response branches as Gitlab::Git::Branch objects' do
+    it 'concatenates and returns the response branches as Gitlab::Git::Branch objects' do
       target_commits = create_list(:gitaly_commit, 4)
       response_branches = target_commits.each_with_index.map do |gitaly_commit, i|
         Gitaly::Branch.new(name: "#{remote_name}/#{i}", target_commit: gitaly_commit)
@@ -59,6 +59,17 @@ describe Gitlab::GitalyClient::RefService do
     end
   end
 
+  describe '#merged_branches' do
+    it 'sends a find_all_branches message' do
+      expect_any_instance_of(Gitaly::RefService::Stub)
+        .to receive(:find_all_branches)
+        .with(gitaly_request_with_params(merged_only: true, merged_branches: ['test']), kind_of(Hash))
+        .and_return([])
+
+      client.merged_branches(%w(test))
+    end
+  end
+
   describe '#branch_names' do
     it 'sends a find_all_branch_names message' do
       expect_any_instance_of(Gitaly::RefService::Stub)
@@ -135,6 +146,38 @@ describe Gitlab::GitalyClient::RefService do
     end
   end
 
+  describe '#tags' do
+    it 'sends a find_all_tags message' do
+      expect_any_instance_of(Gitaly::RefService::Stub)
+        .to receive(:find_all_tags)
+        .and_return([])
+
+      client.tags
+    end
+  end
+
+  describe '#branch_names_contains_sha' do
+    it 'sends a list_branch_names_containing_commit message' do
+      expect_any_instance_of(Gitaly::RefService::Stub)
+        .to receive(:list_branch_names_containing_commit)
+        .with(gitaly_request_with_params(commit_id: '123', limit: 0), kind_of(Hash))
+        .and_return([])
+
+      client.branch_names_contains_sha('123')
+    end
+  end
+
+  describe '#get_tag_messages' do
+    it 'sends a get_tag_messages message' do
+      expect_any_instance_of(Gitaly::RefService::Stub)
+        .to receive(:get_tag_messages)
+        .with(gitaly_request_with_params(tag_ids: ['some_tag_id']), kind_of(Hash))
+        .and_return([])
+
+      client.get_tag_messages(['some_tag_id'])
+    end
+  end
+
   describe '#find_ref_name', :seed_helper do
     subject { client.find_ref_name(SeedRepo::Commit::ID, 'refs/heads/master') }