btrc-hub/backend/app/services/shared_resource_id_allocator.rb

require 'digest'


class SharedResourceIdAllocator
  SPACES = {
      'tags' => ['tags', 'external_tags'],
      'wiki' => ['tag_descriptions', 'tag_groups'],
      'wiki_history' => ['tag_description_versions', 'tag_group_versions'] }.freeze
  HISTORY_IDS = {
      'tags' => [
          ['tag_versions', 'tag_id'],
          ['nico_tag_versions', 'tag_id'],
          ['tag_description_versions', 'tag_id']],
      'wiki' => [['tag_group_versions', 'tag_group_id']] }.freeze
  LOCK_TIMEOUT = 10
  CONNECTION_LOCKS = ObjectSpace::WeakMap.new

  class LockCleanup
    def initialize &finish
      @finish = finish
    end

    def trigger_transactional_callbacks? = false
    def before_committed! = nil
    def committed!(**) = @finish.call(self)
    def rolledback!(**) = @finish.call(self)
  end
  private_constant :LockCleanup

  def self.assign! record
    space, tables = SPACES.find { |_name, names| names.include?(record.class.table_name) }
    raise ArgumentError, 'Unknown shared ID resource' unless tables

    connection = record.class.connection
    lock_space!(record.class, space)

    # 明示-tic な Id. 指定は許可する．
    return if record.id.present?

    next_id = record.class.uncached do
      historical_max = HISTORY_IDS.fetch(space, []).map { |table, column|
        quoted_table = connection.quote_table_name(table)
        quoted_column = connection.quote_column_name(column)
        connection.select_value("SELECT MAX(#{ quoted_column }) FROM #{ quoted_table }").to_i
      }.max.to_i

      [next_ids(connection, tables).max, historical_max + 1].max
    end

    record.id = next_id
  end

  def self.lock_space! record_class, space
    raise ArgumentError, 'Unknown shared ID resource' unless SPACES.key?(space)

    connection = record_class.connection
    transaction = connection.current_transaction
    unless transaction.open?
      raise ActiveRecord::ActiveRecordError, 'Shared ID requires a transaction'
    end

    database = connection.pool.db_config.database
    name = "btrc:ids:#{ Digest::SHA256.hexdigest(database)[0, 16] }:#{ space }"
    held_locks = CONNECTION_LOCKS[connection] ||= {}
    return if held_locks.key?(name)

    quoted_name = connection.quote(name)
    acquired = record_class.uncached do
      connection.select_value("SELECT GET_LOCK(#{ quoted_name }, #{ LOCK_TIMEOUT })")
    end
    unless acquired.to_i == 1
      raise ActiveRecord::LockWaitTimeout, "Shared ID lock unavailable: #{ space }"
    end
    held_locks[name] = true

    finish = lambda do |cleanup|
      next unless held_locks.key?(name)

      # Savepoint callbacks may run before the outer transaction ends.
      # Keep one session lock and transfer its cleanup to the remaining parent.
      parent = connection.current_transaction
      if parent.open?
        connection.add_transaction_record(cleanup)
        next
      end
      result = record_class.uncached do
        connection.select_value("SELECT RELEASE_LOCK(#{ quoted_name })")
      end
      unless result.to_i == 1
        raise ActiveRecord::ActiveRecordError, "Shared ID lock release failed: #{ space }"
      end
      held_locks.delete(name)
    rescue StandardError
      # Never return a session with an untracked advisory lock to the pool.
      begin
        connection.disconnect!
      ensure
        held_locks.clear
      end
      raise
    end
    begin
      # Rails finalises remaining records in ensure if another callback raises.
      connection.add_transaction_record(LockCleanup.new(&finish))
    rescue StandardError
      begin
        connection.disconnect!
      ensure
        held_locks.clear
      end
      raise
    end
  end

  def self.next_ids connection, tables
    expiry = connection.select_value('SELECT @@SESSION.information_schema_stats_expiry').to_i
    begin
      connection.execute('SET SESSION information_schema_stats_expiry = 0')
      quoted_tables = tables.map { connection.quote(_1) }.join(', ')
      rows = connection.select_all(<<~SQL).to_a
        SELECT TABLE_NAME, ENGINE, AUTO_INCREMENT
        FROM information_schema.tables
        WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME IN (#{ quoted_tables })
      SQL
      persistent = rows.all? do |row|
        row['ENGINE'] == 'InnoDB' && row['AUTO_INCREMENT'].to_i.positive?
      end
      unless rows.size == tables.size && persistent
        raise ActiveRecord::ActiveRecordError, 'Persistent InnoDB counters are required'
      end
      rows.map { _1.fetch('AUTO_INCREMENT').to_i }
    ensure
      connection.execute("SET SESSION information_schema_stats_expiry = #{ expiry }")
    end
  end
  private_class_method :next_ids
end