Add rake task to rebuild non-unique indexes

Add a new rake task to seamlessly rebuild a non-unique index,
which can be used as a replacement for some of the functionality
provided by pg_repack.

lib/gitlab/database/concurrent_reindex.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class ConcurrentReindex
+      include Gitlab::Utils::StrongMemoize
+      include MigrationHelpers
+
+      ReindexError = Class.new(StandardError)
+
+      PG_IDENTIFIER_LENGTH = 63
+      TEMPORARY_INDEX_PREFIX = 'tmp_reindex_'
+      REPLACED_INDEX_PREFIX = 'old_reindex_'
+
+      attr_reader :index_name, :logger
+
+      def initialize(index_name, logger:)
+        @index_name = index_name
+        @logger = logger
+      end
+
+      def execute
+        raise ReindexError, "index #{index_name} does not exist" unless index_exists?
+
+        raise ReindexError, 'UNIQUE indexes are currently not supported' if index_unique?
+
+        logger.debug("dropping dangling index from previous run: #{replacement_index_name}")
+        remove_replacement_index
+
+        begin
+          create_replacement_index
+
+          unless replacement_index_valid?
+            message = 'replacement index was created as INVALID'
+            logger.error("#{message}, cleaning up")
+            raise ReindexError, "failed to reindex #{index_name}: #{message}"
+          end
+
+          swap_replacement_index
+        rescue Gitlab::Database::WithLockRetries::AttemptsExhaustedError => e
+          logger.error('failed to obtain the required database locks to swap the indexes, cleaning up')
+          raise ReindexError, e.message
+        rescue ActiveRecord::ActiveRecordError, PG::Error => e
+          logger.error("database error while attempting reindex of #{index_name}: #{e.message}")
+          raise ReindexError, e.message
+        ensure
+          logger.info("dropping unneeded replacement index: #{replacement_index_name}")
+          remove_replacement_index
+        end
+      end
+
+      private
+
+      def connection
+        @connection ||= ActiveRecord::Base.connection
+      end
+
+      def replacement_index_name
+        @replacement_index_name ||= constrained_index_name(TEMPORARY_INDEX_PREFIX)
+      end
+
+      def index
+        strong_memoize(:index) do
+          find_index(index_name)
+        end
+      end
+
+      def index_exists?
+        !index.nil?
+      end
+
+      def index_unique?
+        index.indisunique
+      end
+
+      def constrained_index_name(prefix)
+        "#{prefix}#{index_name}".slice(0, PG_IDENTIFIER_LENGTH)
+      end
+
+      def create_replacement_index
+        create_replacement_index_statement = index.indexdef
+          .sub(/CREATE INDEX/, 'CREATE INDEX CONCURRENTLY')
+          .sub(/#{index_name}/, replacement_index_name)
+
+        logger.info("creating replacement index #{replacement_index_name}")
+        logger.debug("replacement index definition: #{create_replacement_index_statement}")
+
+        disable_statement_timeout do
+          connection.execute(create_replacement_index_statement)
+        end
+      end
+
+      def replacement_index_valid?
+        find_index(replacement_index_name).indisvalid
+      end
+
+      def find_index(index_name)
+        record = connection.select_one(<<~SQL)
+          SELECT
+            pg_index.indisunique,
+            pg_index.indisvalid,
+            pg_indexes.indexdef
+          FROM pg_index
+          INNER JOIN pg_class ON pg_class.oid = pg_index.indexrelid
+          INNER JOIN pg_namespace ON pg_class.relnamespace = pg_namespace.oid
+          INNER JOIN pg_indexes ON pg_class.relname = pg_indexes.indexname
+          WHERE pg_namespace.nspname = 'public'
+          AND pg_class.relname = #{connection.quote(index_name)}
+        SQL
+
+        OpenStruct.new(record) if record
+      end
+
+      def swap_replacement_index
+        replaced_index_name = constrained_index_name(REPLACED_INDEX_PREFIX)
+
+        logger.info("swapping replacement index #{replacement_index_name} with #{index_name}")
+
+        with_lock_retries do
+          rename_index(index_name, replaced_index_name)
+          rename_index(replacement_index_name, index_name)
+          rename_index(replaced_index_name, replacement_index_name)
+        end
+      end
+
+      def rename_index(old_index_name, new_index_name)
+        connection.execute("ALTER INDEX #{old_index_name} RENAME TO #{new_index_name}")
+      end
+
+      def remove_replacement_index
+        disable_statement_timeout do
+          connection.execute("DROP INDEX CONCURRENTLY IF EXISTS #{replacement_index_name}")
+        end
+      end
+
+      def with_lock_retries(&block)
+        arguments = { klass: self.class, logger: logger }
+
+        Gitlab::Database::WithLockRetries.new(arguments).run(raise_on_exhaustion: true, &block)
+      end
+    end
+  end
+end

lib/tasks/gitlab/db.rake

@@ -166,5 +166,12 @@ namespace :gitlab do
     Rake::Task['db:test:load'].enhance do
       Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
     end
+
+    desc 'reindex a regular (non-unique) index without downtime to eliminate bloat'
+    task :reindex, [:index_name] => :environment do |_, args|
+      raise ArgumentError, 'must give the index name to reindex' unless args[:index_name]
+
+      Gitlab::Database::ConcurrentReindex.new(args[:index_name], logger: Logger.new(STDOUT)).execute
+    end
   end
 end

spec/lib/gitlab/database/concurrent_reindex_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::ConcurrentReindex, '#execute' do
+  subject { described_class.new(index_name, logger: logger) }
+
+  let(:table_name) { '_test_reindex_table' }
+  let(:column_name) { '_test_column' }
+  let(:index_name) { '_test_reindex_index' }
+  let(:logger) { double('logger', debug: nil, info: nil, error: nil ) }
+  let(:connection) { ActiveRecord::Base.connection }
+
+  before do
+    connection.execute(<<~SQL)
+      CREATE TABLE #{table_name} (
+        id serial NOT NULL PRIMARY KEY,
+        #{column_name} integer NOT NULL);
+
+      CREATE INDEX #{index_name} ON #{table_name} (#{column_name});
+    SQL
+  end
+
+  context 'when the index does not exist' do
+    before do
+      connection.execute(<<~SQL)
+        DROP INDEX #{index_name}
+      SQL
+    end
+
+    it 'raises an error' do
+      expect { subject.execute }.to raise_error(described_class::ReindexError, /does not exist/)
+    end
+  end
+
+  context 'when the index is unique' do
+    before do
+      connection.execute(<<~SQL)
+        DROP INDEX #{index_name};
+        CREATE UNIQUE INDEX #{index_name} ON #{table_name} (#{column_name})
+      SQL
+    end
+
+    it 'raises an error' do
+      expect do
+        subject.execute
+      end.to raise_error(described_class::ReindexError, /UNIQUE indexes are currently not supported/)
+    end
+  end
+
+  context 'replacing the original index with a rebuilt copy' do
+    let(:replacement_name) { 'tmp_reindex__test_reindex_index' }
+    let(:replaced_name) { 'old_reindex__test_reindex_index' }
+
+    let(:create_index) { "CREATE INDEX CONCURRENTLY #{replacement_name} ON public.#{table_name} USING btree (#{column_name})" }
+    let(:drop_index) { "DROP INDEX CONCURRENTLY IF EXISTS #{replacement_name}" }
+
+    let!(:original_index) { find_index_create_statement }
+
+    before do
+      allow(subject).to receive(:connection).and_return(connection)
+      allow(subject).to receive(:disable_statement_timeout).and_yield
+    end
+
+    it 'replaces the existing index with an identical index' do
+      expect(subject).to receive(:disable_statement_timeout).exactly(3).times.and_yield
+
+      expect_to_execute_concurrently_in_order(drop_index)
+      expect_to_execute_concurrently_in_order(create_index)
+
+      expect_next_instance_of(::Gitlab::Database::WithLockRetries) do |instance|
+        expect(instance).to receive(:run).with(raise_on_exhaustion: true).and_yield
+      end
+
+      expect_to_execute_in_order("ALTER INDEX #{index_name} RENAME TO #{replaced_name}")
+      expect_to_execute_in_order("ALTER INDEX #{replacement_name} RENAME TO #{index_name}")
+      expect_to_execute_in_order("ALTER INDEX #{replaced_name} RENAME TO #{replacement_name}")
+
+      expect_to_execute_concurrently_in_order(drop_index)
+
+      subject.execute
+
+      check_index_exists
+    end
+
+    context 'when a dangling index is left from a previous run' do
+      before do
+        connection.execute("CREATE INDEX #{replacement_name} ON #{table_name} (#{column_name})")
+      end
+
+      it 'replaces the existing index with an identical index' do
+        expect(subject).to receive(:disable_statement_timeout).exactly(3).times.and_yield
+
+        expect_to_execute_concurrently_in_order(drop_index)
+        expect_to_execute_concurrently_in_order(create_index)
+
+        expect_next_instance_of(::Gitlab::Database::WithLockRetries) do |instance|
+          expect(instance).to receive(:run).with(raise_on_exhaustion: true).and_yield
+        end
+
+        expect_to_execute_in_order("ALTER INDEX #{index_name} RENAME TO #{replaced_name}")
+        expect_to_execute_in_order("ALTER INDEX #{replacement_name} RENAME TO #{index_name}")
+        expect_to_execute_in_order("ALTER INDEX #{replaced_name} RENAME TO #{replacement_name}")
+
+        expect_to_execute_concurrently_in_order(drop_index)
+
+        subject.execute
+
+        check_index_exists
+      end
+    end
+
+    context 'when it fails to create the replacement index' do
+      it 'safely cleans up and signals the error' do
+        expect_to_execute_concurrently_in_order(drop_index)
+
+        expect(connection).to receive(:execute).with(create_index).ordered
+          .and_raise(ActiveRecord::ConnectionTimeoutError, 'connect timeout')
+
+        expect_to_execute_concurrently_in_order(drop_index)
+
+        expect { subject.execute }.to raise_error(described_class::ReindexError, /connect timeout/)
+
+        check_index_exists
+      end
+    end
+
+    context 'when the replacement index is not valid' do
+      it 'safely cleans up and signals the error' do
+        expect_to_execute_concurrently_in_order(drop_index)
+        expect_to_execute_concurrently_in_order(create_index)
+
+        expect(subject).to receive(:replacement_index_valid?).and_return(false)
+
+        expect_to_execute_concurrently_in_order(drop_index)
+
+        expect { subject.execute }.to raise_error(described_class::ReindexError, /replacement index was created as INVALID/)
+
+        check_index_exists
+      end
+    end
+
+    context 'when a database error occurs while swapping the indexes' do
+      it 'safely cleans up and signals the error' do
+        expect_to_execute_concurrently_in_order(drop_index)
+        expect_to_execute_concurrently_in_order(create_index)
+
+        expect_next_instance_of(::Gitlab::Database::WithLockRetries) do |instance|
+          expect(instance).to receive(:run).with(raise_on_exhaustion: true).and_yield
+        end
+
+        expect(connection).to receive(:execute).ordered
+          .with("ALTER INDEX #{index_name} RENAME TO #{replaced_name}")
+          .and_raise(ActiveRecord::ConnectionTimeoutError, 'connect timeout')
+
+        expect_to_execute_concurrently_in_order(drop_index)
+
+        expect { subject.execute }.to raise_error(described_class::ReindexError, /connect timeout/)
+
+        check_index_exists
+      end
+    end
+
+    context 'when with_lock_retries fails to acquire the lock' do
+      it 'safely cleans up and signals the error' do
+        expect_to_execute_concurrently_in_order(drop_index)
+        expect_to_execute_concurrently_in_order(create_index)
+
+        expect_next_instance_of(::Gitlab::Database::WithLockRetries) do |instance|
+          expect(instance).to receive(:run).with(raise_on_exhaustion: true)
+            .and_raise(::Gitlab::Database::WithLockRetries::AttemptsExhaustedError, 'exhausted')
+        end
+
+        expect_to_execute_concurrently_in_order(drop_index)
+
+        expect { subject.execute }.to raise_error(described_class::ReindexError, /exhausted/)
+
+        check_index_exists
+      end
+    end
+  end
+
+  def expect_to_execute_concurrently_in_order(sql)
+    # Indexes cannot be created CONCURRENTLY in a transaction. Since the tests are wrapped in transactions,
+    # verify the original call but pass through the non-concurrent form.
+    expect(connection).to receive(:execute).with(sql).ordered.and_wrap_original do |method, sql|
+      method.call(sql.sub(/CONCURRENTLY/, ''))
+    end
+  end
+
+  def expect_to_execute_in_order(sql)
+    expect(connection).to receive(:execute).with(sql).ordered.and_call_original
+  end
+
+  def find_index_create_statement
+    ActiveRecord::Base.connection.select_value(<<~SQL)
+      SELECT indexdef
+      FROM pg_indexes
+      WHERE schemaname = 'public'
+      AND indexname = #{ActiveRecord::Base.connection.quote(index_name)}
+    SQL
+  end
+
+  def check_index_exists
+    expect(find_index_create_statement).to eq(original_index)
+  end
+end

spec/tasks/gitlab/db_rake_spec.rb

@@ -164,9 +164,31 @@ RSpec.describe 'gitlab:db namespace rake task' do
     end
   end
 
-  def run_rake_task(task_name)
+  describe 'reindex' do
+    context 'when no index_name is given' do
+      it 'raises an error' do
+        expect do
+          run_rake_task('gitlab:db:reindex', '')
+        end.to raise_error(ArgumentError, /must give the index name/)
+      end
+    end
+
+    it 'calls the index rebuilder with the proper arguments' do
+      reindex = double('rebuilder')
+
+      expect(Gitlab::Database::ConcurrentReindex).to receive(:new)
+        .with('some_index_name', logger: instance_of(Logger))
+        .and_return(reindex)
+
+      expect(reindex).to receive(:execute)
+
+      run_rake_task('gitlab:db:reindex', '[some_index_name]')
+    end
+  end
+
+  def run_rake_task(task_name, arguments = '')
     Rake::Task[task_name].reenable
-    Rake.application.invoke_task task_name
+    Rake.application.invoke_task("#{task_name}#{arguments}")
   end
 
   def expect_multiple_executions_of_task(test_task_name, task_to_invoke, count: 2)