diff options
Diffstat (limited to 'activerecord/lib/active_record/middleware/database_selector')
-rw-r--r-- | activerecord/lib/active_record/middleware/database_selector/resolver.rb | 88 | ||||
-rw-r--r-- | activerecord/lib/active_record/middleware/database_selector/resolver/session.rb | 45 |
2 files changed, 133 insertions, 0 deletions
diff --git a/activerecord/lib/active_record/middleware/database_selector/resolver.rb b/activerecord/lib/active_record/middleware/database_selector/resolver.rb new file mode 100644 index 0000000000..1c056e2156 --- /dev/null +++ b/activerecord/lib/active_record/middleware/database_selector/resolver.rb @@ -0,0 +1,88 @@ +# frozen_string_literal: true + +require "active_record/middleware/database_selector/resolver/session" + +module ActiveRecord + module Middleware + class DatabaseSelector + # The Resolver class is used by the DatabaseSelector middleware to + # determine which database the request should use. + # + # To change the behavior of the Resolver class in your application, + # create a custom resolver class that inherts from + # DatabaseSelector::Resolver and implements the methods that need to + # be changed. + # + # By default the Resolver class will send read traffic to the replica + # if it's been 2 seconds since the last write. + class Resolver # :nodoc: + SEND_TO_REPLICA_DELAY = 2.seconds + + def self.call(resolver) + new(resolver) + end + + def initialize(resolver) + @resolver = resolver + @instrumenter = ActiveSupport::Notifications.instrumenter + end + + attr_reader :resolver, :instrumenter + + def read(&blk) + if read_from_primary? + read_from_primary(&blk) + else + read_from_replica(&blk) + end + end + + def write(&blk) + write_to_primary(&blk) + end + + private + + def read_from_primary(&blk) + ActiveRecord::Base.connection.while_preventing_writes do + ActiveRecord::Base.connected_to(role: :writing) do + instrumenter.instrument("database_selector.active_record.read_from_primary") do + yield + end + end + end + end + + def read_from_replica(&blk) + ActiveRecord::Base.connected_to(role: :reading) do + instrumenter.instrument("database_selector.active_record.read_from_replica") do + yield + end + end + end + + def write_to_primary(&blk) + ActiveRecord::Base.connected_to(role: :writing) do + instrumenter.instrument("database_selector.active_record.wrote_to_primary") do + resolver.update_last_write_timestamp + yield + end + end + end + + def read_from_primary? + !time_since_last_write_ok? + end + + def send_to_replica_delay + (ActiveRecord::Base.database_selector && ActiveRecord::Base.database_selector[:delay]) || + SEND_TO_REPLICA_DELAY + end + + def time_since_last_write_ok? + Time.now - resolver.last_write_timestamp >= send_to_replica_delay + end + end + end + end +end diff --git a/activerecord/lib/active_record/middleware/database_selector/resolver/session.rb b/activerecord/lib/active_record/middleware/database_selector/resolver/session.rb new file mode 100644 index 0000000000..33e0af5ee4 --- /dev/null +++ b/activerecord/lib/active_record/middleware/database_selector/resolver/session.rb @@ -0,0 +1,45 @@ +# frozen_string_literal: true + +module ActiveRecord + module Middleware + class DatabaseSelector + class Resolver + # The session class is used by the DatabaseSelector::Resolver to save + # timestamps of the last write in the session. + # + # The last_write is used to determine whether it's safe to read + # from the replica or the request needs to be sent to the primary. + class Session # :nodoc: + def self.build(request) + new(request.session) + end + + # Converts time to a timestamp that represents milliseconds since + # epoch. + def self.convert_time_to_timestamp(time) + time.to_i * 1000 + time.usec / 1000 + end + + # Converts milliseconds since epoch timestamp into a time object. + def self.convert_timestamp_to_time(timestamp) + timestamp ? Time.at(timestamp / 1000, (timestamp % 1000) * 1000) : Time.at(0) + end + + def initialize(session) + @session = session + end + + attr_reader :session + + def last_write_timestamp + self.class.convert_timestamp_to_time(session[:last_write]) + end + + def update_last_write_timestamp + session[:last_write] = self.class.convert_time_to_timestamp(Time.now) + end + end + end + end + end +end |