Schedule rebalancing if we timeout when finding gaps

This is a symptom that we need to rebalance a set of issues. We plug
into the gap-finding mechanism to diagnose this and schedule remediation

app/models/issue.rb

@@ -440,6 +440,22 @@ class Issue < ApplicationRecord
     key = Gitlab::Routing.url_helpers.realtime_changes_project_issue_path(project, self)
     Gitlab::EtagCaching::Store.new.touch(key)
   end
+
+  def find_next_gap_before
+    super
+  rescue ActiveRecord::QueryCanceled => e
+    # Symptom of running out of space - schedule rebalancing
+    IssueRebalancingWorker.perform_async(id)
+    raise e
+  end
+
+  def find_next_gap_after
+    super
+  rescue ActiveRecord::QueryCanceled => e
+    # Symptom of running out of space - schedule rebalancing
+    IssueRebalancingWorker.perform_async(id)
+    raise e
+  end
 end
 
 Issue.prepend_if_ee('EE::Issue')

spec/models/issue_spec.rb

@@ -1183,4 +1183,32 @@ RSpec.describe Issue do
       expect(context[:label_url_method]).to eq(:project_issues_url)
     end
   end
+
+  describe 'scheduling rebalancing' do
+    before do
+      allow(issue).to receive(:find_next_gap) { raise ActiveRecord::QueryCanceled }
+    end
+
+    let(:project) { build(:project_empty_repo) }
+    let(:issue) { build_stubbed(:issue, relative_position: 100, project: project) }
+
+    describe '#find_next_gap_before' do
+      it 'schedules rebalancing if we time-out when finding a gap' do
+        lhs = build_stubbed(:issue, relative_position: 99, project: project)
+        to_move = build(:issue, project: project)
+        expect(IssueRebalancingWorker).to receive(:perform_async).with(issue.id)
+
+        expect { to_move.move_between(lhs, issue) }.to raise_error(ActiveRecord::QueryCanceled)
+      end
+    end
+
+    describe '#find_next_gap_after' do
+      it 'schedules rebalancing if we time-out when finding a gap' do
+        allow(issue).to receive(:find_next_gap) { raise ActiveRecord::QueryCanceled }
+        expect(IssueRebalancingWorker).to receive(:perform_async).with(issue.id)
+
+        expect { issue.move_sequence_after }.to raise_error(ActiveRecord::QueryCanceled)
+      end
+    end
+  end
 end