From 66a4484bbb7a80b28374a3f266a9333cb1d03756 Mon Sep 17 00:00:00 2001 From: Andrey Novikov Date: Thu, 26 Feb 2026 22:28:49 +0900 Subject: [PATCH] Sidekiq middleware --- Gemfile | 2 + Gemfile.lock | 47 ++++++++++- README.md | 32 ++++++++ lib/singed/sidekiq.rb | 41 ++++++++++ spec/singed/sidekiq_spec.rb | 158 ++++++++++++++++++++++++++++++++++++ spec/support/sidekiq.rb | 93 +++++++++++++++++++++ 6 files changed, 372 insertions(+), 1 deletion(-) create mode 100644 lib/singed/sidekiq.rb create mode 100644 spec/singed/sidekiq_spec.rb create mode 100644 spec/support/sidekiq.rb diff --git a/Gemfile b/Gemfile index 3ee0307..ec1b8a4 100644 --- a/Gemfile +++ b/Gemfile @@ -5,5 +5,7 @@ source "https://rubygems.org" # Specify your gem's dependencies in singed.gemspec gemspec +gem "activejob" gem "rake", "~> 13.0" +gem "sidekiq" gem "standard" diff --git a/Gemfile.lock b/Gemfile.lock index 04b22de..4b44c7b 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -7,18 +7,51 @@ PATH GEM remote: https://rubygems.org/ specs: + activejob (8.1.2) + activesupport (= 8.1.2) + globalid (>= 0.3.6) + activesupport (8.1.2) + base64 + bigdecimal + concurrent-ruby (~> 1.0, >= 1.3.1) + connection_pool (>= 2.2.5) + drb + i18n (>= 1.6, < 2) + json + logger (>= 1.4.2) + minitest (>= 5.1) + securerandom (>= 0.3) + tzinfo (~> 2.0, >= 2.0.5) + uri (>= 0.13.1) ast (2.4.2) + base64 (0.3.0) + bigdecimal (4.0.1) + concurrent-ruby (1.3.6) + connection_pool (3.0.2) diff-lcs (1.5.0) - json (2.7.2) + drb (2.2.3) + globalid (1.3.0) + activesupport (>= 6.1) + i18n (1.14.8) + concurrent-ruby (~> 1.0) + json (2.18.1) language_server-protocol (3.17.0.3) lint_roller (1.1.0) + logger (1.7.0) + minitest (6.0.2) + drb (~> 2.0) + prism (~> 1.5) parallel (1.24.0) parser (3.3.1.0) ast (~> 2.4.1) racc + prism (1.9.0) racc (1.7.3) + rack (3.2.5) rainbow (3.1.1) rake (13.0.6) + redis-client (0.26.4) + connection_pool regexp_parser (2.9.0) rexml (3.2.6) rspec (3.12.0) @@ -51,6 +84,13 @@ GEM rubocop (>= 1.48.1, < 2.0) rubocop-ast (>= 1.30.0, < 2.0) ruby-progressbar (1.13.0) + securerandom (0.4.1) + sidekiq (8.1.1) + connection_pool (>= 3.0.0) + json (>= 2.16.0) + logger (>= 1.7.0) + rack (>= 3.2.0) + redis-client (>= 0.26.0) stackprof (0.2.17) standard (1.35.1) language_server-protocol (~> 3.17.0.2) @@ -64,14 +104,19 @@ GEM standard-performance (1.3.1) lint_roller (~> 1.1) rubocop-performance (~> 1.20.2) + tzinfo (2.0.6) + concurrent-ruby (~> 1.0) unicode-display_width (2.5.0) + uri (1.1.1) PLATFORMS ruby DEPENDENCIES + activejob rake (~> 13.0) rspec + sidekiq singed! standard diff --git a/README.md b/README.md index a52fee2..f852552 100644 --- a/README.md +++ b/README.md @@ -90,6 +90,38 @@ PROTIP: use Chrome Developer Tools to record network activity, and copy requests This can also be enabled to always run by setting `SINGED_MIDDLEWARE_ALWAYS_CAPTURE=1` in the environment. +### Sidekiq + +If you are using Sidekiq, you can use the `Singed::Sidekiq::ServerMiddleware` to capture flamegraphs for you. + +```ruby +require "singed/sidekiq" + +Sidekiq.configure_server do |config| + config.server_middleware do |chain| + chain.add Singed::Sidekiq::ServerMiddleware + end +end +``` + +To capture flamegraphs for all jobs, you can set the `SINGED_MIDDLEWARE_ALWAYS_CAPTURE` environment variable to `true` the same way as the Rack middleware. + +To capture flamegraphs for a specific job, you can set the `x-singed` key in the job payload to `true`. + +```ruby +MyJob.set(x-singed: true).perform_async +``` + +Or define a `capture_flamegraph?` method on the job class: + +```ruby +class MyJob + def self.capture_flamegraph?(payload) + payload["flamegraph"] + end +end +``` + ### Command Line There is a `singed` command line you can use that will record a flamegraph from the entirety of a command run: diff --git a/lib/singed/sidekiq.rb b/lib/singed/sidekiq.rb new file mode 100644 index 0000000..e4e612d --- /dev/null +++ b/lib/singed/sidekiq.rb @@ -0,0 +1,41 @@ +# frozen_string_literal: true + +module Singed + module Sidekiq + class ServerMiddleware + include ::Sidekiq::ServerMiddleware + + def call(job_instance, job_payload, queue, &block) + return block.call unless capture_flamegraph?(job_instance, job_payload) + + flamegraph(flamegraph_label(job_instance, job_payload), &block) + end + + private + + TRUTHY_STRINGS = %w[true 1 yes].freeze + + def capture_flamegraph?(job_instance, job_payload) + return job_payload["x-singed"] if job_payload.key?("x-singed") + + job_class = self.job_class(job_instance, job_payload) + return job_class.capture_flamegraph?(job_payload) if job_class.respond_to?(:capture_flamegraph?) + + TRUTHY_STRINGS.include?(ENV.fetch("SINGED_MIDDLEWARE_ALWAYS_CAPTURE", "false")) + end + + def flamegraph_label(job_instance, job_payload) + [job_class(job_instance, job_payload), job_payload["jid"]].compact.join("--") + end + + def job_class(job_instance, job_payload) + job_class = job_payload.fetch("wrapped", job_instance) # ActiveJob + return job_class if job_class.is_a?(Class) + return job_class.class if job_class.is_a?(::Sidekiq::Job) + return job_class.constantize if job_class.respond_to?(:constantize) + + Object.const_get(job_class.to_s) + end + end + end +end diff --git a/spec/singed/sidekiq_spec.rb b/spec/singed/sidekiq_spec.rb new file mode 100644 index 0000000..827436a --- /dev/null +++ b/spec/singed/sidekiq_spec.rb @@ -0,0 +1,158 @@ +# frozen_string_literal: true + +require "spec_helper" +require "sidekiq" +require "active_job" +require "singed/sidekiq" +require_relative "../support/sidekiq" + +RSpec.describe Singed::Sidekiq::ServerMiddleware do + subject { job_class.set(job_modifiers).perform_async(*job_args) } + + let(:job_class) { SidekiqPlainJob } + let(:job_args) { [] } + let(:job_modifiers) { {} } + + before do + allow_any_instance_of(described_class).to receive(:flamegraph) { |*, &block| block.call } + allow_any_instance_of(job_class).to receive(:perform).and_call_original + end + + context "with plain Sidekiq jobs" do + it "doesn't capture flamegraph by default" do + expect_any_instance_of(described_class).not_to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + + context "when x-singed payload is true" do + let(:job_modifiers) { {"x-singed" => true} } + + it "wraps execution in flamegraph when x-singed is true" do + expect_any_instance_of(described_class).to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + end + end + + context "with class-level capture_flamegraph?" do + let(:job_class) { SidekiqFlamegraphJob } + + it "doesn't capture when capture_flamegraph? returns false" do + expect_any_instance_of(described_class).not_to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + + context "when payload satisfies capture_flamegraph?" do + let(:job_modifiers) { {"x-flamegraph" => true} } + + it "wraps execution in flamegraph when capture_flamegraph? returns true" do + expect_any_instance_of(described_class).to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + end + end + + context "when SINGED_MIDDLEWARE_ALWAYS_CAPTURE env var is set" do + around do |example| + original = ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] + example.run + ensure + if original.nil? + ENV.delete("SINGED_MIDDLEWARE_ALWAYS_CAPTURE") + else + ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] = original + end + end + + context "when SINGED_MIDDLEWARE_ALWAYS_CAPTURE=true" do + before { ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] = "true" } + + it "wraps execution in flamegraph" do + expect_any_instance_of(described_class).to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + end + + context "when SINGED_MIDDLEWARE_ALWAYS_CAPTURE is false" do + before { ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] = "false" } + + it "doesn't capture flamegraph" do + expect_any_instance_of(described_class).not_to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + end + end + + context "with ActiveJob jobs" do + subject { job_class.set(job_modifiers).perform_later(*job_args) } + + context "with plain ActiveJob" do + let(:job_class) { ActiveJobPlainJob } + + it "doesn't capture flamegraph by default" do + expect_any_instance_of(described_class).not_to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + + context "with ActiveJob class where capture_flamegraph? is true" do + let(:job_class) { ActiveJobFlamegraphJob } + + it "wraps execution in flamegraph when capture_flamegraph? returns true" do + expect_any_instance_of(described_class).to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + end + + context "with ActiveJob class where capture_flamegraph? is false" do + let(:job_class) { ActiveJobNoFlamegraphJob } + + it "doesn't capture when capture_flamegraph? returns false" do + expect_any_instance_of(described_class).not_to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + end + + context "when SINGED_MIDDLEWARE_ALWAYS_CAPTURE env var is set" do + around do |example| + original = ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] + example.run + ensure + if original.nil? + ENV.delete("SINGED_MIDDLEWARE_ALWAYS_CAPTURE") + else + ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] = original + end + end + + context "when SINGED_MIDDLEWARE_ALWAYS_CAPTURE=true" do + before { ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] = "true" } + + it "wraps execution in flamegraph" do + expect_any_instance_of(described_class).to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + end + + context "when SINGED_MIDDLEWARE_ALWAYS_CAPTURE is false" do + before { ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] = "false" } + + it "doesn't capture flamegraph" do + expect_any_instance_of(described_class).not_to receive(:flamegraph) + expect_any_instance_of(job_class).to receive(:perform) + subject + end + end + end + end + end +end diff --git a/spec/support/sidekiq.rb b/spec/support/sidekiq.rb new file mode 100644 index 0000000..a8e2c55 --- /dev/null +++ b/spec/support/sidekiq.rb @@ -0,0 +1,93 @@ +require "singed/sidekiq" +require "tempfile" + +RSpec.configure do |config| + config.before(:suite) do + Sidekiq.testing!(:inline) + + Sidekiq::Client.prepend(SidekiqTestingInlineWithMiddlewares) + + Sidekiq.configure_client do |config| + config.server_middleware do |chain| + chain.add Singed::Sidekiq::ServerMiddleware + end + end + + ActiveJob::Base.queue_adapter = :sidekiq + ActiveJob::Base.logger = Logger.new(nil) + + Singed.output_directory = Dir.mktmpdir("singed-sidekiq-spec") + end +end + +# Sidekiq doesn't invoke middlewares in inline testingmode, so we need to invoke it oursleves +module SidekiqTestingInlineWithMiddlewares + # rubocop:disable Metrics/AbcSize + def push(job) + return super unless Sidekiq::Testing.inline? + + job = Sidekiq.load_json(Sidekiq.dump_json(job)) + job["jid"] ||= SecureRandom.hex(12) + job_class = Object.const_get(job["class"]) + job_instance = job_class.new + queue = (job_instance.sidekiq_options_hash || {}).fetch("queue", "default") + server = Sidekiq.respond_to?(:default_configuration) ? Sidekiq.default_configuration : Sidekiq + server.server_middleware.invoke(job_instance, job, queue) do + job_instance.perform(*job["args"]) + end + job["jid"] + end + # rubocop:enable Metrics/AbcSize +end + +class SidekiqPlainJob + include Sidekiq::Job + + def perform(*_args) + "My job is simple" + end +end + +class SidekiqFlamegraphJob + include Sidekiq::Job + + def self.capture_flamegraph?(payload) + !!payload["x-flamegraph"] + end + + def perform(*_args) + "Phew, I'm done!" + end +end + +class ActiveJobPlainJob < ActiveJob::Base + self.queue_adapter = :sidekiq + + def perform(*_args) + "My job is simple" + end +end + +class ActiveJobFlamegraphJob < ActiveJob::Base + self.queue_adapter = :sidekiq + + def self.capture_flamegraph?(_payload) + true + end + + def perform(*_args) + "Phew, I'm done!" + end +end + +class ActiveJobNoFlamegraphJob < ActiveJob::Base + self.queue_adapter = :sidekiq + + def self.capture_flamegraph?(_payload) + false + end + + def perform(*_args) + "Phew, I'm done!" + end +end