From 78ea82c3e2073d7e81c59f7fbc09587892ba6729 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <godfat@godfat.org>
Date: Mon, 2 Apr 2018 17:40:27 +0800
Subject: [PATCH] Re-enable allowing n+1 Gitaly calls for cold cache

Whenever cache is enabled and cold, branches index still makes
n+1 calls. Make sure we catch this.
---
 .../projects/branches_controller.rb            | 10 +++++++---
 .../projects/branches_controller_spec.rb       | 18 +++++++++++++++++-
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/app/controllers/projects/branches_controller.rb b/app/controllers/projects/branches_controller.rb
index 176679f0849..b7b36f770f5 100644
--- a/app/controllers/projects/branches_controller.rb
+++ b/app/controllers/projects/branches_controller.rb
@@ -22,9 +22,13 @@ class Projects::BranchesController < Projects::ApplicationController
 
         @refs_pipelines = @project.pipelines.latest_successful_for_refs(@branches.map(&:name))
         @merged_branch_names = repository.merged_branch_names(@branches.map(&:name))
-        @max_commits = @branches.reduce(0) do |memo, branch|
-          diverging_commit_counts = repository.diverging_commit_counts(branch)
-          [memo, diverging_commit_counts[:behind], diverging_commit_counts[:ahead]].max
+
+        # n+1: https://gitlab.com/gitlab-org/gitaly/issues/992
+        Gitlab::GitalyClient.allow_n_plus_1_calls do
+          @max_commits = @branches.reduce(0) do |memo, branch|
+            diverging_commit_counts = repository.diverging_commit_counts(branch)
+            [memo, diverging_commit_counts[:behind], diverging_commit_counts[:ahead]].max
+          end
         end
 
         render
diff --git a/spec/controllers/projects/branches_controller_spec.rb b/spec/controllers/projects/branches_controller_spec.rb
index 3b9e06cb5ad..16fb377b002 100644
--- a/spec/controllers/projects/branches_controller_spec.rb
+++ b/spec/controllers/projects/branches_controller_spec.rb
@@ -398,6 +398,22 @@ describe Projects::BranchesController do
       end
     end
 
+    # We need :request_store because Gitaly only counts the queries whenever
+    # `RequestStore.active?` in GitalyClient.enforce_gitaly_request_limits
+    # And the main goal of this test is making sure TooManyInvocationsError
+    # was not raised whenever the cache is enabled yet cold.
+    context 'when cache is enabled yet cold', :request_store do
+      it 'return with a status 200' do
+        get :index,
+            namespace_id: project.namespace,
+            project_id: project,
+            state: 'all',
+            format: :html
+
+        expect(response).to have_gitlab_http_status(200)
+      end
+    end
+
     context 'when branch contains an invalid UTF-8 sequence' do
       before do
         project.repository.create_branch("wrong-\xE5-utf8-sequence")
@@ -414,7 +430,7 @@ describe Projects::BranchesController do
       end
     end
 
-    context 'when depreated sort/search/page parameters are specified' do
+    context 'when deprecated sort/search/page parameters are specified' do
       it 'returns with a status 301 when sort specified' do
         get :index,
             namespace_id: project.namespace,