Skip to content

Commit

Permalink
Don't clear ActiveRecord connections when in nested job
Browse files Browse the repository at this point in the history
  • Loading branch information
tomgi committed Jun 27, 2024
1 parent e675fee commit 62517fd
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 11 deletions.
45 changes: 34 additions & 11 deletions lib/que/active_record/connection.rb
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@
module Que
module ActiveRecord
class << self
def active_rails_executor?
defined?(::Rails.application.executor) && ::Rails.application.executor.active?
end

def wrap_in_rails_executor(&block)
if defined?(::Rails.application.executor)
::Rails.application.executor.wrap(&block)
Expand Down Expand Up @@ -39,17 +43,36 @@ def call(job)
yield
end

# ActiveRecord will check out connections to the current thread when
# queries are executed and not return them to the pool until
# explicitly requested to. I'm not wild about this API design, and
# it doesn't pose a problem for the typical case of workers using a
# single PG connection (since we ensure that connection is checked
# in and checked out responsibly), but since ActiveRecord supports
# connections to multiple databases, it's easy for people using that
# feature to unknowingly leak connections to other databases. So,
# take the additional step of telling ActiveRecord to check in all
# of the current thread's connections after each job is run.
::ActiveRecord::Base.clear_active_connections! unless job.class.resolve_que_setting(:run_synchronously)
clear_active_connections_if_needed!(job)
end

private

# ActiveRecord will check out connections to the current thread when
# queries are executed and not return them to the pool until
# explicitly requested to. I'm not wild about this API design, and
# it doesn't pose a problem for the typical case of workers using a
# single PG connection (since we ensure that connection is checked
# in and checked out responsibly), but since ActiveRecord supports
# connections to multiple databases, it's easy for people using that
# feature to unknowingly leak connections to other databases. So,
# take the additional step of telling ActiveRecord to check in all
# of the current thread's connections after each job is run.
def clear_active_connections_if_needed!(job)
# don't clean in synchronous mode
# see https://github.com/que-rb/que/pull/393
return if job.class.resolve_que_setting(:run_synchronously)

# don't clear connections in nested jobs,
# i.e. clear only if this is the outermost instance of
# Que::ActiveRecord::Connection::JobMiddleware
return if Que::ActiveRecord.active_rails_executor?

if ::ActiveRecord.version >= Gem::Version.new('7.1')
::ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
else
::ActiveRecord::Base.clear_active_connections!
end
end
end
end
Expand Down
19 changes: 19 additions & 0 deletions spec/que/active_record/connection_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -62,5 +62,24 @@ def run(*args)

assert SecondDatabaseModel.connection_handler.active_connections?
end

it "shouldn't clear connections to secondary DBs if within an active rails executor" do
# This is a hacky spec, but it's better than requiring Rails.
rails, application, executor = 3.times.map { Object.new }
application.define_singleton_method(:executor) { executor }
rails.define_singleton_method(:application) { application }
executor.define_singleton_method(:wrap) { |&block| block.call }
executor.define_singleton_method(:active?) { true }

refute defined?(::Rails)
::Rails = rails

SecondDatabaseModelJob.run

assert SecondDatabaseModel.connection_handler.active_connections?

Object.send :remove_const, :Rails
refute defined?(::Rails)
end
end
end

0 comments on commit 62517fd

Please sign in to comment.