Releasing connections in ActiveRecord 7.1+

[rsamoilov]

The ActiveRecord connection pool is designed so that if a thread (or a fiber) checks out a connection, the connection will remain attached to this thread until the server has finished processing the request.

This makes sense in a threaded environment but doesn't make sense for Rage. Consider the following code:

def index
  unless ApiToken.exists?(token: params[:token])
    head :forbidden
    return
  end

  Net::HTTP.get(URI("..."))
end

This code makes a request to the DB (ApiToken.exist?) and then sends an HTTP request. In Rails, the ActiveRecord connection, checked out with the ApiToken.exist? call, will remain attached to the thread which performs the request until the action (and consequently the HTTP request) has finished.

However, this doesn't make sense for Rage - while the current fiber is paused, waiting for the HTTP request to finish, other fibers could use the DB connection. Hence, we need to release ActiveRecord connections not once the whole action has finished but once a DB request has been completed.

Rage implements this using the Fiber.defer method. Unfortunately, something has changed in ActiveRecord 7.1, and while this approach works in ActiveRecord 7.0.8.1, it doesn't work reliably in ActiveRecord 7.1 anymore and thus has been disabled.

Goal: Make changes to allow releasing DB connections once a DB request has completed with ActiveRecord 7.1+.

Errors:

ActiveRecord::StatementInvalid (wrong argument type nil (expected PG::TypeMap)):
activerecord-7.1.0/lib/active_record/core.rb:411:in `rescue in cached_find_by'
activerecord-7.1.0/lib/active_record/core.rb:408:in `cached_find_by'
activerecord-7.1.0/lib/active_record/core.rb:251:in `find'

NoMethodError (undefined method `key?' for nil:NilClass):
activerecord-7.1.0/lib/active_record/connection_adapters/postgresql_adapter.rb:832:in `get_oid_type'
activerecord-7.1.0/lib/active_record/connection_adapters/postgresql/database_statements.rb:67:in `block (2 levels) in internal_exec_query'
activerecord-7.1.0/lib/active_record/connection_adapters/postgresql/database_statements.rb:64:in `each'
activerecord-7.1.0/lib/active_record/connection_adapters/postgresql/database_statements.rb:64:in `each_with_index'
activerecord-7.1.0/lib/active_record/connection_adapters/postgresql/database_statements.rb:64:in `block in internal_exec_query'
activerecord-7.1.0/lib/active_record/connection_adapters/postgresql_adapter.rb:880:in `execute_and_clear'
activerecord-7.1.0/lib/active_record/connection_adapters/postgresql/database_statements.rb:61:in `internal_exec_query'

[efibootmgr]

Probably caused by these changes? rails/rails@42fede5

[rsamoilov]

hey @efibootmgr, thank you for the suggestion!
It could be it, but I can reproduce the bug with AR 7.1.0, which doesn't have these changes.

[efibootmgr]

After some researching and help of Claude 3.5 hallucinations and intensive reverse prompting, i came back with this. Claude also suggested some fix, and why not, i will also post it below. Low chance it may help, but
rails/rails#50793
rails/rails@1dcb411
This may be related with the issue/cause additional problems. But the changes are not in 7.1 either. These are from 7.2.

[efibootmgr]

rails/rails@b36f918
And this is what causing current problem with 7.1 it seems. I'm 100% sure. Take a look at all the changes related to these two commits. That should help you :)

[efibootmgr]

The Claude said:
Thank you for providing the specific commit that's causing the issue. This commit indeed changes how the SchemaCache works, which can affect how connections are managed, especially in a fiber-based environment like Rage. Let's work on a solution to address this.

Here's an approach to fix this issue in Rage:

Create a custom SchemaReflection class that's fiber-aware:

# lib/fiber_aware_schema_reflection.rb
class FiberAwareSchemaReflection < ActiveRecord::ConnectionAdapters::SchemaReflection
  def initialize(connection_pool)
    @connection_pool = connection_pool
    super()
  end

  def connection
    Fiber.current[:ar_connection] || @connection_pool.connection
  end

  def schema_cache
    connection.schema_cache
  end

  # Override other methods as needed, always using the `connection` method to get the current connection
end

Modify the ActiveRecord connection handler for Rage:

# lib/active_record_fiber_handler.rb
module ActiveRecordFiberHandler
  def self.with_connection
    ActiveRecord::Base.connection_pool.with_connection do |conn|
      Fiber.current[:ar_connection] = conn
      begin
        # Ensure proper schema reflection
        conn.instance_variable_set(:@schema_reflection, FiberAwareSchemaReflection.new(ActiveRecord::Base.connection_pool))
        yield conn
      ensure
        Fiber.current[:ar_connection] = nil
      end
    end
  end

  def self.release_connection
    if conn = Fiber.current[:ar_connection]
      ActiveRecord::Base.connection_pool.release_connection(conn)
      Fiber.current[:ar_connection] = nil
    end
  end
end

Patch ActiveRecord::Base to use this handler:

# lib/active_record_fiber_patch.rb
module ActiveRecordFiberPatch
  def connection
    if Fiber.current[:ar_connection]
      Fiber.current[:ar_connection]
    else
      ActiveRecordFiberHandler.with_connection { |conn| conn }
    end
  end
end

ActiveRecord::Base.singleton_class.prepend(ActiveRecordFiberPatch)

Update the Rage::FiberWrapper:

class Rage::FiberWrapper
  def call(env)
    fiber = Fiber.schedule do
      @app.call(env)
    ensure
      ActiveRecordFiberHandler.release_connection
      Iodine.publish(Fiber.current.__get_id, "", Iodine::PubSub::PROCESS)
    end

    if fiber.alive?
      [:__http_defer__, fiber]
    else
      fiber.__get_result
    end
  end
end

Initialize this setup in your application:

# config/initializers/active_record_fiber_setup.rb
require 'fiber_aware_schema_reflection'
require 'active_record_fiber_handler'
require 'active_record_fiber_patch'

ActiveRecord::Base.singleton_class.prepend(ActiveRecordFiberPatch)

This solution addresses the changes introduced by the commit by:

Creating a fiber-aware SchemaReflection that always uses the correct connection for the current fiber.
Ensuring that each fiber gets its own connection with a properly initialized schema reflection.
Properly managing connections within the fiber lifecycle.

This approach should resolve the issues caused by the SchemaCache refactoring while maintaining compatibility with Rage's fiber-based concurrency model. Make sure to thoroughly test this solution in your environment, especially under various load conditions and with different types of database operations.

The code may be wrong, but he must be definitely right lmao. I suggest you to check on this issue before releasing websockets. I'm not using Action Cable, but instead trying combination of Crystal based websocket server and LavinMQ. But I'm sure that addition of Action Cable websockets will help with smaller projects.
p.s. I will continue to use Rage like forever. Hope this problem will be fixed now. Your project is the best.

[rsamoilov]

woah, this is wild!! Looking into it!

[rsamoilov]

Hey @efibootmgr

The Rails issue is not fixed yet, but we've reworked the way we release AR connections. The new approach is compatible with AR 6.0+ and it will go into v1.10.0 - see #103.

Thank you for your input! Couldn't have happened without you!