Implement workaround for Pitchfork PID provider compatibility (#62)

* Implement workaround for Pitchfork PID provider compatibility

* Use Puma method for PID provider compatibility, adding a fallback

* Fix Namespace ref

* Ignore rubocop

I think the complexity in the setup method is unavoidable,
and trying to factor this logic into another method doesn't improve
anything.

The assignment in if conditions stuff is a bit suspect, but I don't
really mind it ...

* Don't do code coverage on specs

* Add a spec to check the pid_provider selection logic

* Improve code coverage

* Improve code coverage

---------

Co-authored-by: Ed Robinson <[email protected]>

lib/promenade/pitchfork/worker_pid_provider.rb

@@ -0,0 +1,31 @@
+module Promenade
+  module Pitchfork
+    class WorkerPidProvider
+      def self.fetch
+        worker_id || "process_id_#{Process.pid}"
+      end
+
+      def self.object_based_worker_id
+        return unless defined?(::Pitchfork::Worker)
+
+        workers = ObjectSpace.each_object(::Pitchfork::Worker)
+        return if workers.nil?
+
+        workers_first = workers.first
+        workers_first&.nr
+      end
+
+      def self.program_name
+        $PROGRAM_NAME
+      end
+
+      def self.worker_id
+        if matchdata = program_name.match(/pitchfork.*worker\[(.+)\]/) # rubocop:disable Lint/AssignmentInCondition
+          "pitchfork_#{matchdata[1]}"
+        elsif object_worker_id = object_based_worker_id # rubocop:disable Lint/AssignmentInCondition
+          "pitchfork_#{object_worker_id}"
+        end
+      end
+    end
+  end
+end

lib/promenade/setup.rb

@@ -27,11 +27,31 @@ def setup
     ENV["prometheus_multiproc_dir"] = multiprocess_files_dir.to_s
 
     require "prometheus/client"
-    require "prometheus/client/support/unicorn"
 
     ::Prometheus::Client.configure do |config|
       config.multiprocess_files_dir = multiprocess_files_dir
-      config.pid_provider = ::Prometheus::Client::Support::Unicorn.method(:worker_pid_provider)
+
+      config.pid_provider = pid_provider_method
+    end
+  end
+
+  def pid_provider_method
+    # This workaround enables us to utilize the same PID provider for both Unicorn, Pitchfork and Puma.
+    # We cannot employ the same method directly because Unicorn and Pitchfork are not loaded simultaneously.
+    # Instead, we define a method that dynamically loads the appropriate PID provider based on the active server.
+    # As a fallback, we use the process ID.
+
+    if defined?(::Unicorn)
+      require "prometheus/client/support/unicorn"
+      ::Prometheus::Client::Support::Unicorn.method(:worker_pid_provider)
+    elsif defined?(::Pitchfork)
+      require "promenade/pitchfork/worker_pid_provider"
+      Pitchfork::WorkerPidProvider.method(:fetch)
+    elsif defined?(::Puma)
+      require "prometheus/client/support/puma"
+      ::Prometheus::Client::Support::Puma.method(:worker_pid_provider)
+    else
+      -> { "process_id_#{Process.pid}" }
     end
   end
 end

spec/promenade/pitchfork/worker_pid_provider_spec.rb

@@ -0,0 +1,82 @@
+require "spec_helper"
+require "promenade/pitchfork/worker_pid_provider"
+
+RSpec.describe Promenade::Pitchfork::WorkerPidProvider do
+  describe ".fetch" do
+    subject { described_class.fetch }
+
+    context "when worker_id is defined" do
+      let(:worker_id) { "worker_id" }
+
+      before do
+        allow(described_class).to receive(:worker_id).and_return(worker_id)
+      end
+
+      it { is_expected.to eq(worker_id) }
+    end
+
+    context "when worker_id is not defined" do
+      before do
+        allow(described_class).to receive(:worker_id).and_return(nil)
+        allow(Process).to receive(:pid).and_return(123)
+      end
+
+      it { is_expected.to eq("process_id_123") }
+    end
+  end
+
+  describe ".worker_id" do
+    subject { described_class.send(:worker_id) }
+    around(:example) do |ex|
+      old_name = $PROGRAM_NAME
+      $PROGRAM_NAME = program_name
+      ex.run
+      $PROGRAM_NAME = old_name
+    end
+
+    context "when program_name matches pitchfork worker" do
+      let(:program_name) { "pitchfork (gen:0) worker[1] - requests: 13, waiting" }
+
+
+      it { is_expected.to eq("pitchfork_1") }
+
+      context "when program_name matches pitchfork worker" do
+        let(:program_name) { "pitchfork worker[1]" }
+
+        it { is_expected.to eq("pitchfork_1") }
+      end
+
+      context "when program_name doesn't match pitchfork worker" do
+        let(:program_name) { "something else" }
+
+        let(:worker) { double("Pitchfork::Worker", nr: 2) }
+
+        before do
+          stub_const("Pitchfork::Worker", Class.new)
+          allow(ObjectSpace).to receive(:each_object).with(Pitchfork::Worker).and_return([worker])
+        end
+
+        it { is_expected.to eq("pitchfork_2") }
+      end
+    end
+  end
+
+  describe ".object_based_worker_id" do
+    subject { described_class.send(:object_based_worker_id) }
+
+    context "when Pitchfork::Worker is defined" do
+      let(:worker) { double("Pitchfork::Worker", nr: 1) }
+
+      before do
+        stub_const("Pitchfork::Worker", Class.new)
+        allow(ObjectSpace).to receive(:each_object).with(Pitchfork::Worker).and_return([worker])
+      end
+
+      it { is_expected.to eq(1) }
+    end
+
+    context "when Pitchfork::Worker is not defined" do
+      it { is_expected.to be_nil }
+    end
+  end
+end

spec/setup_spec.rb

@@ -2,6 +2,33 @@
 
 RSpec.describe Promenade do
   describe ".setup" do
+    context "pid_provider" do
+      before do
+        stub_const(self.class.description, true)
+      end
+
+      context "Unicorn" do
+        it "uses the unicorn pid provider" do
+          Promenade.setup
+          expect(Prometheus::Client.configuration.pid_provider.to_s).to match "Prometheus::Client::Support::Unicorn"
+        end
+      end
+
+      context "Pitchfork" do
+        it "uses the pitchfork pid provider" do
+          Promenade.setup
+          expect(Prometheus::Client.configuration.pid_provider.to_s).to match "Promenade::Pitchfork::WorkerPidProvider"
+        end
+      end
+
+      context "Puma" do
+        it "uses the puma pid provider" do
+          Promenade.setup
+          expect(Prometheus::Client.configuration.pid_provider.to_s).to match "Prometheus::Client::Support::Puma"
+        end
+      end
+    end
+
     context "without rails" do
       context "environment variable set" do
         let(:multiproc_root) { Pathname.new(Dir.mktmpdir) }

spec/spec_helper.rb

@@ -1,6 +1,8 @@
 require "simplecov"
 SimpleCov.minimum_coverage 99
-SimpleCov.start
+SimpleCov.start do
+  add_filter "/spec/"
+end
 
 if ENV["CI"]
   require "simplecov-cobertura"

spec/yjit_spec.rb

@@ -1,4 +1,5 @@
 require "promenade/yjit/stats"
+require "promenade/yjit/middleware"
 require "open3"
 
 RSpec.describe Promenade::YJIT::Stats do
@@ -56,3 +57,17 @@ def parse_number(string)
     string.to_f
   end
 end
+
+RSpec.describe Promenade::YJIT::Middlware do
+  let(:app) { double(:app, call: nil) }
+
+  it "is adds it's instrumentation method to the rack.after_reply array" do
+    stats = class_spy("Promenade::YJIT::Stats").as_stubbed_const
+
+    after_reply = []
+    described_class.new(app).call({ "rack.after_reply" => after_reply })
+    after_reply.each(&:call)
+
+    expect(stats).to have_received(:instrument)
+  end
+end