Search: Counts endpoint respond with 408 for query timeouts

app/controllers/search_controller.rb

@@ -5,6 +5,8 @@ class SearchController < ApplicationController
   include SearchHelper
   include RedisTracking
 
+  RESCUE_FROM_TIMEOUT_ACTIONS = [:count, :show].freeze
+
   track_redis_hll_event :show, name: 'i_search_total'
 
   around_action :allow_gitaly_ref_name_caching
@@ -154,13 +156,22 @@ class SearchController < ApplicationController
   end
 
   def render_timeout(exception)
-    raise exception unless action_name.to_sym == :show
+    raise exception unless action_name.to_sym.in?(RESCUE_FROM_TIMEOUT_ACTIONS)
 
     log_exception(exception)
 
     @timeout = true
+
+    if count_action_name?
+      render json: {}, status: :request_timeout
+    else
       render status: :request_timeout
     end
+  end
+
+  def count_action_name?
+    action_name.to_sym == :count
+  end
 end
 
 SearchController.prepend_mod_with('SearchController')

spec/controllers/search_controller_spec.rb

@@ -53,6 +53,20 @@ RSpec.describe SearchController do
       end
     end
 
+    shared_examples_for 'support for active record query timeouts' do |action, params, method_to_stub, format|
+      before do
+        allow_next_instance_of(SearchService) do |service|
+          allow(service).to receive(method_to_stub).and_raise(ActiveRecord::QueryCanceled)
+        end
+      end
+
+      it 'renders a 408 when a timeout occurs' do
+        get action, params: params, format: format
+
+        expect(response).to have_gitlab_http_status(:request_timeout)
+      end
+    end
+
     describe 'GET #show' do
       it_behaves_like 'when the user cannot read cross project', :show, { search: 'hello' } do
         it 'still allows accessing the search page' do
@@ -63,6 +77,7 @@ RSpec.describe SearchController do
       end
 
       it_behaves_like 'with external authorization service enabled', :show, { search: 'hello' }
+      it_behaves_like 'support for active record query timeouts', :show, { search: 'hello' }, :search_objects, :html
 
       context 'uses the right partials depending on scope' do
         using RSpec::Parameterized::TableSyntax
@@ -230,6 +245,7 @@ RSpec.describe SearchController do
     describe 'GET #count' do
       it_behaves_like 'when the user cannot read cross project', :count, { search: 'hello', scope: 'projects' }
       it_behaves_like 'with external authorization service enabled', :count, { search: 'hello', scope: 'projects' }
+      it_behaves_like 'support for active record query timeouts', :count, { search: 'hello', scope: 'projects' }, :search_results, :json
 
       it 'returns the result count for the given term and scope' do
         create(:project, :public, name: 'hello world')