Merge remote-tracking branch 'upstream/main' into chore/standardrb

* upstream/main:
  Sidekiq middleware (#53)
  Methods for explicit start and stop flamegraph recording (#52)

Author: Envek

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"

Gemfile.lock

@@ -7,19 +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.3)
+    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.19.0)
+    drb (2.2.3)
+    globalid (1.3.0)
+      activesupport (>= 6.1)
+    i18n (1.14.8)
+      concurrent-ruby (~> 1.0)
+    json (2.19.1)
     language_server-protocol (3.17.0.5)
     lint_roller (1.1.0)
+    logger (1.7.0)
+    minitest (6.0.2)
+      drb (~> 2.0)
+      prism (~> 1.5)
     parallel (1.27.0)
     parser (3.3.10.2)
       ast (~> 2.4.1)
       racc
     prism (1.9.0)
     racc (1.8.1)
+    rack (3.2.5)
     rainbow (3.1.1)
     rake (13.0.6)
+    redis-client (0.26.4)
+      connection_pool
     regexp_parser (2.11.3)
     rspec (3.12.0)
       rspec-core (~> 3.12.0)
@@ -53,6 +85,13 @@ GEM
       rubocop (>= 1.75.0, < 2.0)
       rubocop-ast (>= 1.47.1, < 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.54.0)
       language_server-protocol (~> 3.17.0.2)
@@ -66,16 +105,21 @@ GEM
     standard-performance (1.9.0)
       lint_roller (~> 1.1)
       rubocop-performance (~> 1.26.0)
+    tzinfo (2.0.6)
+      concurrent-ruby (~> 1.0)
     unicode-display_width (3.2.0)
       unicode-emoji (~> 4.1)
     unicode-emoji (4.2.0)
+    uri (1.1.1)
 
 PLATFORMS
   ruby
 
 DEPENDENCIES
+  activejob
   rake (~> 13.0)
   rspec
+  sidekiq
   singed!
   standard
 

README.md

@@ -47,6 +47,25 @@ flamegraph(open: false) {
 }
 ```
 
+### Explicit start and stop
+
+You can also start and stop the flamegraph explicitly:
+
+```ruby
+# config/boot.rb
+require "singed"
+Singed.output_directory ||= Dir.pwd + "/tmp/speedscope"
+Singed.start
+# Let some code to run here...
+# and then stop the flamegraph with e.g. rails runner 'Singed.stop'
+flamegraph = Singed.stop
+# The flamegraph is saved to the output directory
+# Open it with your browser:
+flamegraph.open
+```
+
+Note that `Singed.start` can't be run multiple times in parallel, instantiate multiple `Singed::Flamegraph` objects instead and call `start` on them.
+
 ### RSpec
 
 If you are using RSpec, you can use the `flamegraph` metadata to capture it for you.
@@ -90,6 +109,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:

lib/singed.rb

@@ -8,7 +8,7 @@ module Singed
 
   # Where should flamegraphs be saved?
   def output_directory=(directory)
-    @output_directory = Pathname.new(directory)
+    @output_directory = directory && Pathname.new(directory)
   end
 
   def self.output_directory
@@ -45,6 +45,28 @@ def filter_line(line)
     line
   end
 
+  def start(label = nil, ignore_gc: false, interval: 1000)
+    return unless enabled?
+    return if profiling?
+
+    @current_flamegraph = Flamegraph.new(label: label, ignore_gc: ignore_gc, interval: interval)
+    @current_flamegraph.tap(&:start)
+  end
+
+  def stop
+    return nil unless profiling?
+
+    flamegraph = @current_flamegraph
+    @current_flamegraph = nil
+    flamegraph.stop
+    flamegraph.save
+    flamegraph
+  end
+
+  def profiling?
+    @current_flamegraph&.started? || false
+  end
+
   autoload :Flamegraph, "singed/flamegraph"
   autoload :Report, "singed/report"
   autoload :RackMiddleware, "singed/rack_middleware"

lib/singed/flamegraph.rb

@@ -23,14 +23,31 @@ def initialize(label: nil, ignore_gc: false, interval: 1000, filename: nil)
     end
 
     def record
-      return yield unless Singed.enabled?
-      return yield if filename.exist? # file existing means its been captured already
+      start
+      yield
+    ensure
+      stop
+    end
 
-      result = nil
-      @profile = StackProf.run(mode: :wall, raw: true, ignore_gc: @ignore_gc, interval: @interval) do
-        result = yield
-      end
-      result
+    def start
+      return false unless Singed.enabled?
+      return false if filename.exist? # file existing means its been captured already
+      return false if started?
+
+      StackProf.start(mode: :wall, raw: true, ignore_gc: @ignore_gc, interval: @interval)
+      @started = true
+    end
+
+    def stop
+      return nil unless started?
+
+      @started = false
+      StackProf.stop
+      @profile = StackProf.results
+    end
+
+    def started?
+      !!@started
     end
 
     def save

lib/singed/sidekiq.rb

@@ -0,0 +1,44 @@
+# 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 TRUTHY_STRINGS.include?(job_payload["x-singed"].to_s) if job_payload.key?("x-singed")
+
+        job_class = job_class(job_instance, job_payload)
+        return false unless job_class
+        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)
+      rescue NameError
+        nil
+      end
+    end
+  end
+end