makara run in new context

application.rb

module XXXProjectName
  class Application < Rails::Application
    # Settings in config/environments/* take precedence over those specified here.
    # Application configuration should go into files in config/initializers
    # -- all .rb files in that directory are automatically loaded.
    config.to_prepare do
      # Load application's model / class decorators
      Dir.glob(File.join(File.dirname(__FILE__), "../app/**/*_decorator*.rb")) do |c|
        Rails.configuration.cache_classes ? require(c) : load(c)
      end

    end
    
    # ... other configurations
  end
end

base_decorator.rb

ActiveRecord::Base.class_eval do
  class << self
    def query_in_new_context
      connection = ActiveRecord::Base.connection
      result = nil

      if makara_enabled?(connection)
        # We were previously reading from master, meaning master context is true.  This resets
        # the context so we can read from slave, but will allow it to get re-stuck to master if needed.
        connection.with_new_context do
          result = yield if block_given?
        end
      else
        Rails.logger.warn "*** Makara connection adapter not enabled.  Falling back to normal behavior. ***"
        result = yield if block_given?
      end

      result
    end
    alias_method :query_slave, :query_in_new_context

    def query_master
      force_master_config = Thread.current.thread_variable_get(:_makara_force_master)
      Thread.current.thread_variable_set(:_makara_force_master, true)

      query_in_new_context { yield } if block_given?
    ensure
      Thread.current.thread_variable_set(:_makara_force_master, force_master_config)
    end

    def makara_enabled?(connection)
      connection.is_a?(ActiveRecord::ConnectionAdapters::DistributedMysqlAdapter)
    end
  end
end

database.yml

development:
  <<: *defaults
  adapter: distributed_mysql
  database: database_name
  makara:
    # the following are default values
    blacklist_duration: 5
    master_ttl: 5
    master_strategy: failover
    sticky: true
    connections:
      - role: master
        host: 10.211.55.5
        username: username1
        password: password1
      - role: slave
        username: username2
        password: password2

distributed_mysql_adapter.rb

require 'active_record/connection_adapters/makara_mysql2_adapter'

module ActiveRecord
  module ConnectionHandling
    def distributed_mysql_connection(config)
      ActiveRecord::ConnectionAdapters::DistributedMysqlAdapter.new(config)
    end
  end
end

module ActiveRecord
  module ConnectionAdapters
    class DistributedMysqlAdapter < ActiveRecord::ConnectionAdapters::MakaraMysql2Adapter
      attr_reader :previous_context, :current_context

      def initialize(config)
        super

        @extended_sticky = @config_parser.makara_config[:extended_sticky]
        @extended_sticky = true if @extended_sticky.nil?
      end

      # 扩展 makara 内置的 needs_master? 方法，支持基于用户会话的主库粘连，粘连时间根据配置决定
      def needs_master?(method_name, args)
        return true if force_master?
        # 允许禁用 sticky 特性
        return super unless @sticky && @extended_sticky && current_user_id.present?

        # 捕捉到写操作
        sql = args.first.to_s
        if method_name.to_s.eql?("execute") && "COMMIT".eql?(sql)
          ttl_log("write! This will be expired in #{@ttl} seconds.")
          Rails.cache.write(user_cache_key, true, expires_in: @ttl)

          return true
        end

        need_master = super
        if !need_master && in_twindow? && sql_slave_matchers.any?{|m| sql =~ m }
          ttl_log("#{method_name}, #{sql}")
          true
        else
          need_master
        end
      end

      def with_new_context
        @current_context  = Makara::Context.get_current
        @previous_context = Makara::Context.get_previous
        @previous_extended_sticky = @extended_sticky

        Makara::Context.set_current(Makara::Context.generate)
        Makara::Context.set_previous(Makara::Context.generate)
        @extended_sticky = false

        yield if block_given?
      ensure
        Makara::Context.set_current(current_context)
        Makara::Context.set_previous(previous_context)
        @extended_sticky = @previous_extended_sticky
        @previous_extended_sticky = nil
      end

      private
      def force_master?
        Thread.current.thread_variable_get(:_makara_force_master)
      end

      def ttl_logger
        @ttl_logger ||= Logger.new("log/makara_read_ttl.log")
      end

      def ttl_log(content, type = :info)
        ttl_logger.send(:info, "T-#{current_request_id} U-#{current_user_id}: #{content}")
      end

      def user_cache_key
        "_makara_#{current_user_id}"
      end

      def in_twindow?
        Rails.cache.read(user_cache_key).present?
      end

      def current_user_id
        Thread.current.thread_variable_get(:user_id)
      end

      def current_request_id
        Thread.current.thread_variable_get(:request_id)
      end
    end
  end
end