Helper middleware to preserve implicit context?

lib/sidekiq/middleware/current_attributes.rb

@@ -0,0 +1,48 @@
+require "active_support/current_attributes"
+
+module Sidekiq
+  ##
+  # Automatically save and load any current attributes in the execution context
+  # so context attributes "flow" from Rails actions into any associated jobs.
+  # This can be useful for multi-tenancy, i18n locale, timezone, any implicit
+  # per-request attribute. See +ActiveSupport::CurrentAttributes+.
+  #
+  # @example
+  #
+  #   # in your initializer
+  #   require "sidekiq/middleware/current_attributes"
+  #   Sidekiq::CurrentAttributes.persist(Myapp::Current)
+  #
+  module CurrentAttributes
+    class Save
+      def initialize(with:)
+        @klass = with
+      end
+
+      def call(_, job, _, _)
+        job["ctx"] = @klass.attributes
+        yield
+      end
+    end
+
+    class Load
+      def initialize(with:)
+        @klass = with
+      end
+
+      def call(_, job, _, &block)
+        @klass.set(job["ctx"], &block)
+      end
+    end
+
+    def self.persist(klass)
+      Sidekiq.configure_client do |config|
+        config.client_middleware.add Save, with: klass
+      end
+      Sidekiq.configure_server do |config|
+        config.client_middleware.add Save, with: klass
+        config.server_middleware.add Load, with: klass
+      end
+    end
+  end
+end

test/test_current_attributes.rb

@@ -0,0 +1,51 @@
+require_relative "./helper"
+require "sidekiq/middleware/current_attributes"
+
+module Myapp
+  class Current < ActiveSupport::CurrentAttributes
+    attribute :user_id
+  end
+end
+
+class TestCurrentAttributes < Minitest::Test
+  def test_save
+    cm = Sidekiq::CurrentAttributes::Save.new(with: Myapp::Current)
+    job = {}
+    with_context(:user_id, 123) do
+      cm.call(nil, job, nil, nil) do
+        assert_equal 123, job["ctx"][:user_id]
+      end
+    end
+  end
+
+  def test_load
+    cm = Sidekiq::CurrentAttributes::Load.new(with: Myapp::Current)
+
+    job = { "ctx" => { "user_id" => 123 } }
+    assert_nil Myapp::Current.user_id
+    cm.call(nil, job, nil) do
+      assert_equal 123, Myapp::Current.user_id
+    end
+    # the Rails reloader is responsible for reseting Current after every unit of work
+  end
+
+  def test_persist
+    begin
+      Sidekiq::CurrentAttributes.persist(Myapp::Current)
+    ensure
+      Sidekiq.client_middleware.clear
+      Sidekiq.server_middleware.clear
+    end
+  end
+
+  private
+
+  def with_context(attr, value)
+    begin
+      Myapp::Current.send("#{attr}=", value)
+      yield
+    ensure
+      Myapp::Current.reset_all
+    end
+  end
+end