From 7b0fe4aef97c6a5f73a03146b669a415f396799c Mon Sep 17 00:00:00 2001 From: Eugen Rochko Date: Fri, 29 Apr 2022 22:43:07 +0200 Subject: [PATCH] Fix opening and closing Redis connections instead of using a pool (#18171) * Fix opening and closing Redis connections instead of using a pool * Fix Redis connections not being returned to the pool in CLI commands --- app/lib/redis_configuration.rb | 7 ++++++- app/models/concerns/redisable.rb | 2 +- config/initializers/stoplight.rb | 6 ++++-- lib/mastodon/cli_helper.rb | 12 +++++++++--- lib/mastodon/rack_middleware.rb | 2 +- lib/mastodon/search_cli.rb | 7 ++++--- lib/mastodon/sidekiq_middleware.rb | 2 +- spec/services/resolve_account_service_spec.rb | 2 ++ 8 files changed, 28 insertions(+), 12 deletions(-) diff --git a/app/lib/redis_configuration.rb b/app/lib/redis_configuration.rb index fc8cf2f80..e14d6c8b6 100644 --- a/app/lib/redis_configuration.rb +++ b/app/lib/redis_configuration.rb @@ -2,12 +2,17 @@ class RedisConfiguration class << self + def establish_pool(new_pool_size) + @pool&.shutdown(&:close) + @pool = ConnectionPool.new(size: new_pool_size) { new.connection } + end + def with pool.with { |redis| yield redis } end def pool - @pool ||= ConnectionPool.new(size: pool_size) { new.connection } + @pool ||= establish_pool(pool_size) end def pool_size diff --git a/app/models/concerns/redisable.rb b/app/models/concerns/redisable.rb index cce8efb86..8d76b6b82 100644 --- a/app/models/concerns/redisable.rb +++ b/app/models/concerns/redisable.rb @@ -6,6 +6,6 @@ module Redisable private def redis - Thread.current[:redis] ||= RedisConfiguration.new.connection + Thread.current[:redis] ||= RedisConfiguration.pool.checkout end end diff --git a/config/initializers/stoplight.rb b/config/initializers/stoplight.rb index d9ebca76c..8c3c5755a 100644 --- a/config/initializers/stoplight.rb +++ b/config/initializers/stoplight.rb @@ -1,4 +1,6 @@ require 'stoplight' -Stoplight::Light.default_data_store = Stoplight::DataStore::Redis.new(RedisConfiguration.new.connection) -Stoplight::Light.default_notifiers = [Stoplight::Notifier::Logger.new(Rails.logger)] +Rails.application.reloader.to_prepare do + Stoplight::Light.default_data_store = Stoplight::DataStore::Redis.new(RedisConfiguration.new.connection) + Stoplight::Light.default_notifiers = [Stoplight::Notifier::Logger.new(Rails.logger)] +end diff --git a/lib/mastodon/cli_helper.rb b/lib/mastodon/cli_helper.rb index aaee1fa91..a78a28e27 100644 --- a/lib/mastodon/cli_helper.rb +++ b/lib/mastodon/cli_helper.rb @@ -19,15 +19,18 @@ module Mastodon ProgressBar.create(total: total, format: '%c/%u |%b%i| %e') end + def reset_connection_pools! + ActiveRecord::Base.establish_connection(ActiveRecord::Base.configurations[Rails.env].dup.tap { |config| config['pool'] = options[:concurrency] + 1 }) + RedisConfiguration.establish_pool(options[:concurrency]) + end + def parallelize_with_progress(scope) if options[:concurrency] < 1 say('Cannot run with this concurrency setting, must be at least 1', :red) exit(1) end - db_config = ActiveRecord::Base.configurations[Rails.env].dup - db_config['pool'] = options[:concurrency] + 1 - ActiveRecord::Base.establish_connection(db_config) + reset_connection_pools! progress = create_progress_bar(scope.count) pool = Concurrent::FixedThreadPool.new(options[:concurrency]) @@ -52,6 +55,9 @@ module Mastodon result = ActiveRecord::Base.connection_pool.with_connection do yield(item) + ensure + RedisConfiguration.pool.checkin if Thread.current[:redis] + Thread.current[:redis] = nil end aggregate.increment(result) if result.is_a?(Integer) diff --git a/lib/mastodon/rack_middleware.rb b/lib/mastodon/rack_middleware.rb index 619a2c36d..8aa7911fe 100644 --- a/lib/mastodon/rack_middleware.rb +++ b/lib/mastodon/rack_middleware.rb @@ -19,7 +19,7 @@ class Mastodon::RackMiddleware end def clean_up_redis_socket! - Thread.current[:redis]&.close + RedisConfiguration.pool.checkin if Thread.current[:redis] Thread.current[:redis] = nil end diff --git a/lib/mastodon/search_cli.rb b/lib/mastodon/search_cli.rb index 6ad9d7b6a..74f980ba1 100644 --- a/lib/mastodon/search_cli.rb +++ b/lib/mastodon/search_cli.rb @@ -59,9 +59,7 @@ module Mastodon index.specification.lock! end - db_config = ActiveRecord::Base.configurations[Rails.env].dup - db_config['pool'] = options[:concurrency] + 1 - ActiveRecord::Base.establish_connection(db_config) + reset_connection_pools! pool = Concurrent::FixedThreadPool.new(options[:concurrency]) added = Concurrent::AtomicFixnum.new(0) @@ -139,6 +137,9 @@ module Mastodon sleep 1 rescue => e progress.log pastel.red("Error importing #{index}: #{e}") + ensure + RedisConfiguration.pool.checkin if Thread.current[:redis] + Thread.current[:redis] = nil end end end diff --git a/lib/mastodon/sidekiq_middleware.rb b/lib/mastodon/sidekiq_middleware.rb index 7ec4097df..c75e8401f 100644 --- a/lib/mastodon/sidekiq_middleware.rb +++ b/lib/mastodon/sidekiq_middleware.rb @@ -26,7 +26,7 @@ class Mastodon::SidekiqMiddleware end def clean_up_redis_socket! - Thread.current[:redis]&.close + RedisConfiguration.pool.checkin if Thread.current[:redis] Thread.current[:redis] = nil end diff --git a/spec/services/resolve_account_service_spec.rb b/spec/services/resolve_account_service_spec.rb index 7b1e8885c..8c302e1d8 100644 --- a/spec/services/resolve_account_service_spec.rb +++ b/spec/services/resolve_account_service_spec.rb @@ -220,6 +220,8 @@ RSpec.describe ResolveAccountService, type: :service do return_values << described_class.new.call('foo@ap.example.com') rescue ActiveRecord::RecordNotUnique fail_occurred = true + ensure + RedisConfiguration.pool.checkin if Thread.current[:redis] end end end