rage-rb 1.2.2 → 1.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 75beeb9e9f98cc51b8b7979be7bbd65253a319092855dfb6bfed0a23fca3bc13
-  data.tar.gz: dc3d0676ca07217e325ec7af031848de7943d1655a3d27f8845500c5a92657b4
+  metadata.gz: cef03d55f96279c7e6d451dc7ce0abce04024a7ca30d07e9c48c0d054c530bb6
+  data.tar.gz: c399b5b3d060aee5367c2e3064435649bbb8eba7dcb2bf2bd215730a698f69d8
 SHA512:
-  metadata.gz: 874cec69c0eb95b78f015105ad75695ad147a190b50488e8bb2bddf2bab5eb9811140d70b6566e3002729892a92ce78f568e48842d3c207ae2db7a3e76b7af80
-  data.tar.gz: 848ad81e200d92fcbbeef7cf43f80cdac118bf41b559186e0976e59a79304fac20e183289a0d67316bc88b7d02dfaf98c4f832a1c6a034be63a2eb9582f637a4
+  metadata.gz: cc5578e1e51d557a8f30729b0eff9a41d4a651e62c837af53aff3b3f50caee2b827be47a239aeadd47f2e9ff15ff8263f35a032634c1fdb738cb2b2c71787780
+  data.tar.gz: 28a70f5c33752ea33dd1474bc4e481ba7b56fd76541544f838226c9b5ace3758b80d9a93b5947c1ea8cbf9374cbae204e9c0e6c5ff40f2adcbc40be3208f5597

data/CHANGELOG.md

@@ -1,5 +1,11 @@
 ## [Unreleased]
 
+## [1.3.0] - 2024-04-17
+
+### Added
+
+- Introduce the `ActiveRecord::ConnectionPool` patch (#78).
+
 ## [1.2.2] - 2024-04-03
 
 ### Fixed

data/lib/rage/configuration.rb

@@ -85,6 +85,18 @@
 #
 # > Specifies connection timeout.
 #
+# # Transient Settings
+#
+# The settings described in this section should be configured using **environment variables** and are either temporary or will become the default in the future.
+#
+# • _RAGE_DISABLE_IO_WRITE_
+#
+# > Disables the `io_write` hook to fix the ["zero-length iov"](https://bugs.ruby-lang.org/issues/19640) error on Ruby < 3.3.
+#
+# • _RAGE_PATCH_AR_POOL_
+#
+# > Enables the `ActiveRecord::ConnectionPool` patch to optimize database connection management. Use it to increase throughput under high load.
+#
 class Rage::Configuration
   attr_accessor :logger
   attr_reader :log_formatter, :log_level

data/lib/rage/controller/api.rb

@@ -75,7 +75,7 @@ class RageController::API
         ""
       end
 
-      activerecord_loaded = Rage.config.internal.rails_mode && defined?(::ActiveRecord)
+      activerecord_loaded = defined?(::ActiveRecord)
 
       class_eval <<~RUBY,  __FILE__, __LINE__ + 1
         def __run_#{action}

data/lib/rage/ext/active_record/connection_pool.rb

@@ -0,0 +1,277 @@
+# frozen_string_literal: true
+
+module Rage::Ext::ActiveRecord::ConnectionPool
+  # items can be added but not removed
+  class BlackHoleList
+    def initialize(arr)
+      @arr = arr
+    end
+
+    def <<(el)
+      @arr << el
+    end
+
+    def shift
+      nil
+    end
+
+    def length
+      0
+    end
+
+    def to_a
+      @arr
+    end
+  end
+
+  def self.extended(instance)
+    instance.class.alias_method :__checkout__, :checkout
+    instance.class.alias_method :__remove__, :remove
+
+    ActiveRecord::ConnectionAdapters::AbstractAdapter.attr_accessor(:__idle_since)
+  end
+
+  def __init_rage_extension
+    # a map of fibers that are currently waiting for a
+    # connection in the format of { Fiber => timestamp }
+    @__blocked = {}
+
+    # a map of fibers that are currently hodling connections
+    # in the format of { Fiber => Connection }
+    @__in_use = {}
+
+    # a list of all DB connections that are currently idle
+    @__connections = build_new_connections
+
+    # how long a fiber can wait for a connection to become available
+    @__checkout_timeout = checkout_timeout
+
+    # how long a connection can be idle for before disconnecting
+    @__idle_timeout = reaper.frequency
+
+    # how often should we check for fibers that wait for a connection for too long
+    @__timeout_worker_frequency = 0.5
+
+    # reject fibers that wait for a connection for more than `@__checkout_timeout`
+    Iodine.run_every((@__timeout_worker_frequency * 1_000).to_i) do
+      if @__blocked.length > 0
+        current_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        @__blocked.each do |fiber, blocked_since|
+          if (current_time - blocked_since) > @__checkout_timeout
+            @__blocked.delete(fiber)
+            fiber.raise(ActiveRecord::ConnectionTimeoutError, "could not obtain a connection from the pool within #{@__checkout_timeout} seconds; all pooled connections were in use")
+          end
+        end
+      end
+    end
+
+    # resume blocked fibers once connections become available
+    Iodine.subscribe("ext:ar-connection-released") do
+      if @__blocked.length > 0 && @__connections.length > 0
+        f, _ = @__blocked.shift
+        f.resume
+      end
+    end
+
+    # unsubscribe on shutdown
+    Iodine.on_state(:on_finish) do
+      Iodine.unsubscribe("ext:ar-connection-released")
+    end
+  end
+
+  # Returns true if there is an open connection being used for the current fiber.
+  def active_connection?
+    @__in_use[Fiber.current]
+  end
+
+  # Retrieve the connection associated with the current fiber, or obtain one if necessary.
+  def connection
+    @__in_use[Fiber.current] ||= @__connections.shift || begin
+      fiber, blocked_since = Fiber.current, Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      @__blocked[fiber] = blocked_since
+      Fiber.yield
+
+      @__connections.shift
+    end
+  end
+
+  # Signal that the fiber is finished with the current connection and it can be returned to the pool.
+  def release_connection(owner = Fiber.current)
+    if conn = @__in_use.delete(owner)
+      conn.__idle_since = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      @__connections << conn
+      Iodine.publish("ext:ar-connection-released", "", Iodine::PubSub::PROCESS) if @__blocked.length > 0
+    end
+
+    conn
+  end
+
+  # Recover lost connections for the pool.
+  def reap
+    @__in_use.each do |fiber, conn|
+      unless fiber.alive?
+        if conn.active?
+          conn.reset!
+          release_connection(fiber)
+        else
+          @__in_use.delete(fiber)
+          conn.disconnect!
+          __remove__(conn)
+          @__connections += build_new_connections(1)
+          Iodine.publish("ext:ar-connection-released", "", Iodine::PubSub::PROCESS) if @__blocked.length > 0
+        end
+      end
+    end
+  end
+
+  # Disconnect all connections that have been idle for at least
+  # `minimum_idle` seconds. Connections currently checked out, or that were
+  # checked in less than `minimum_idle` seconds ago, are unaffected.
+  def flush(minimum_idle = @__idle_timeout)
+    return if minimum_idle.nil? || @__connections.length == 0
+
+    current_time, i = Process.clock_gettime(Process::CLOCK_MONOTONIC), 0
+    while i < @__connections.length
+      conn = @__connections[i]
+      if conn.__idle_since && current_time - conn.__idle_since >= minimum_idle
+        conn.__idle_since = nil
+        conn.disconnect!
+      end
+      i += 1
+    end
+  end
+
+  # Disconnect all currently idle connections. Connections currently checked out are unaffected.
+  def flush!
+    reap
+    flush(-1)
+  end
+
+  # Yields a connection from the connection pool to the block.
+  def with_connection
+    yield connection
+  ensure
+    release_connection
+  end
+
+  # Returns an array containing the connections currently in the pool.
+  def connections
+    @__connections.to_a
+  end
+
+  # Returns true if a connection has already been opened.
+  def connected?
+    true
+  end
+
+  # Return connection pool's usage statistic.
+  def stat
+    {
+      size: size,
+      connections: size,
+      busy: @__in_use.count { |fiber, _| fiber.alive? },
+      dead: @__in_use.count { |fiber, _| !fiber.alive? },
+      idle: @__connections.length,
+      waiting: @__blocked.length,
+      checkout_timeout: @__checkout_timeout
+    }
+  end
+
+  # Disconnects all connections in the pool, and clears the pool.
+  # Raises `ActiveRecord::ExclusiveConnectionTimeoutError` if unable to gain ownership of all
+  # connections in the pool within a timeout interval (default duration is `checkout_timeout * 2` seconds).
+  def disconnect(raise_on_acquisition_timeout = true, disconnect_attempts = 0)
+    # allow request fibers to release connections, but block from acquiring new ones
+    if disconnect_attempts == 0
+      @__connections = BlackHoleList.new(@__connections)
+    end
+
+    # if some connections are in use, we will wait for up to `@__checkout_timeout * 2` seconds
+    if @__in_use.length > 0 && disconnect_attempts <= @__checkout_timeout * 4
+      Iodine.run_after(500) { disconnect(raise_on_acquisition_timeout, disconnect_attempts + 1) }
+      return
+    end
+
+    pool_connections = @__connections.to_a
+
+    # check if there are still some connections in use
+    if @__in_use.length > 0
+      raise(ActiveRecord::ExclusiveConnectionTimeoutError, "could not obtain ownership of all database connections") if raise_on_acquisition_timeout
+      pool_connections += @__in_use.values
+      @__in_use.clear
+    end
+
+    # disconnect all connections
+    pool_connections.each do |conn|
+      conn.disconnect!
+      __remove__(conn)
+    end
+
+    # create a new pool
+    @__connections = build_new_connections
+
+    # notify blocked fibers that there are new connections available
+    [@__blocked.length, @__connections.length].min.times do
+      Iodine.publish("ext:ar-connection-released", "", Iodine::PubSub::PROCESS)
+    end
+  end
+
+  # Disconnects all connections in the pool, and clears the pool.
+  # The pool first tries to gain ownership of all connections. If unable to
+  # do so within a timeout interval (default duration is `checkout_timeout * 2` seconds),
+  # then the pool is forcefully disconnected without any regard for other connection owning fibers.
+  def disconnect!
+    disconnect(false)
+  end
+
+  # Check out a database connection from the pool, indicating that you want
+  # to use it. You should call #checkin when you no longer need this.
+  def checkout(_ = nil)
+    connection
+  end
+
+  # Check in a database connection back into the pool, indicating that you no longer need this connection.
+  def checkin(conn)
+    fiber = @__in_use.key(conn)
+    release_connection(fiber)
+  end
+
+  # Remove a connection from the connection pool. The connection will
+  # remain open and active but will no longer be managed by this pool.
+  def remove(conn)
+    __remove__(conn)
+    @__in_use.delete_if { |_, c| c == conn }
+    @__connections.delete(conn)
+  end
+
+  def clear_reloadable_connections(raise_on_acquisition_timeout = true)
+    disconnect(raise_on_acquisition_timeout)
+  end
+
+  def clear_reloadable_connections!
+    disconnect(false)
+  end
+
+  def num_waiting_in_queue
+    @__blocked.length
+  end
+
+  # Discards all connections in the pool (even if they're currently in use!),
+  # along with the pool itself. Any further interaction with the pool is undefined.
+  def discard!
+    @__discarded = true
+    (@__connections + @__in_use.values).each { |conn| conn.discard! }
+  end
+
+  def discarded?
+    !!@__discarded
+  end
+
+  private
+
+  def build_new_connections(num_connections = size)
+    (1..num_connections).map do
+      __checkout__.tap { |conn| conn.__idle_since = Process.clock_gettime(Process::CLOCK_MONOTONIC) }
+    end
+  end
+end

data/lib/rage/ext/setup.rb

@@ -0,0 +1,36 @@
+# set ActiveSupport isolation level
+if defined?(ActiveSupport::IsolatedExecutionState)
+  ActiveSupport::IsolatedExecutionState.isolation_level = :fiber
+end
+
+# release ActiveRecord connections on yield
+if defined?(ActiveRecord)
+  class Fiber
+    def self.defer
+      res = Fiber.yield
+
+      if ActiveRecord::Base.connection_pool.active_connection?
+        ActiveRecord::Base.connection_handler.clear_active_connections!
+      end
+
+      res
+    end
+  end
+end
+
+# make `ActiveRecord::ConnectionPool` work correctly with fibers
+if defined?(ActiveRecord::ConnectionAdapters::ConnectionPool)
+  ActiveRecord::ConnectionAdapters::ConnectionPool
+  module ActiveRecord::ConnectionAdapters
+    class ConnectionPool
+      def connection_cache_key(_)
+        Fiber.current
+      end
+    end
+  end
+end
+
+# patch `ActiveRecord::ConnectionPool`
+if defined?(ActiveRecord) && ENV["RAGE_PATCH_AR_POOL"]
+  Rage.patch_active_record_connection_pool
+end

data/lib/rage/rails.rb

@@ -11,34 +11,6 @@ Iodine.patch_rack
 # configure the framework
 Rage.config.internal.rails_mode = true
 
-# patch ActiveRecord's connection pool
-if defined?(ActiveRecord)
-  Rails.configuration.after_initialize do
-    module ActiveRecord::ConnectionAdapters
-      class ConnectionPool
-        def connection_cache_key(_)
-          Fiber.current
-        end
-      end
-    end
-  end
-end
-
-# release ActiveRecord connections on yield
-if defined?(ActiveRecord)
-  class Fiber
-    def self.defer
-      res = Fiber.yield
-
-      if ActiveRecord::Base.connection_pool.active_connection?
-        ActiveRecord::Base.connection_handler.clear_active_connections!
-      end
-
-      res
-    end
-  end
-end
-
 # plug into Rails' Zeitwerk instance to reload the code
 Rails.autoloaders.main.on_setup do
   if Iodine.running?
@@ -71,6 +43,10 @@ end
 Rails.configuration.after_initialize do
   if Rails.logger && !Rage.logger
     rails_logdev = Rails.logger.instance_variable_get(:@logdev)
-    Rage.config.logger = Rage::Logger.new(rails_logdev.dev) if rails_logdev.is_a?(Logger::LogDevice)
+    Rage.configure do
+      config.logger = Rage::Logger.new(rails_logdev.dev) if rails_logdev.is_a?(Logger::LogDevice)
+    end
   end
 end
+
+require "rage/ext/setup"

data/lib/rage/setup.rb

@@ -9,3 +9,5 @@ Dir["#{Rage.root}/config/initializers/**/*.rb"].each { |initializer| load(initia
 Rage.code_loader.setup
 
 require_relative "#{Rage.root}/config/routes"
+
+require "rage/ext/setup"

data/lib/rage/templates/Gemfile

@@ -1,6 +1,6 @@
 source "https://rubygems.org"
 
-gem "rage-rb", "<%= Rage::VERSION %>"
+gem "rage-rb", "~> <%= Rage::VERSION[0..2] %>"
 
 # Build JSON APIs with ease
 # gem "alba"

data/lib/rage/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Rage
-  VERSION = "1.2.2"
+  VERSION = "1.3.0"
 end

data/lib/rage-rb.rb

@@ -71,10 +71,37 @@ module Rage
     @code_loader ||= Rage::CodeLoader.new
   end
 
+  def self.patch_active_record_connection_pool
+    patch = proc do
+      is_connected = ActiveRecord::Base.connection_pool rescue false
+      if is_connected
+        puts "INFO: Patching ActiveRecord::ConnectionPool"
+        Iodine.on_state(:on_start) do
+          ActiveRecord::Base.connection_pool.extend(Rage::Ext::ActiveRecord::ConnectionPool)
+          ActiveRecord::Base.connection_pool.__init_rage_extension
+        end
+      else
+        puts "WARNING: DB connection is not established - can't patch ActiveRecord::ConnectionPool"
+      end
+    end
+
+    if Rage.config.internal.rails_mode
+      Rails.configuration.after_initialize(&patch)
+    else
+      patch.call
+    end
+  end
+
   module Router
     module Strategies
     end
   end
+
+  module Ext
+    module ActiveRecord
+      autoload :ConnectionPool, "rage/ext/active_record/connection_pool"
+    end
+  end
 end
 
 module RageController

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rage-rb
 version: !ruby/object:Gem::Version
-  version: 1.2.2
+  version: 1.3.0
 platform: ruby
 authors:
 - Roman Samoilov
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2024-04-03 00:00:00.000000000 Z
+date: 2024-04-17 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: thor
@@ -107,6 +107,8 @@ files:
 - lib/rage/controller/api.rb
 - lib/rage/env.rb
 - lib/rage/errors.rb
+- lib/rage/ext/active_record/connection_pool.rb
+- lib/rage/ext/setup.rb
 - lib/rage/fiber.rb
 - lib/rage/fiber_scheduler.rb
 - lib/rage/logger/json_formatter.rb