Add Falcon support and per-fiber allocation counts

CHANGELOG.markdown

@@ -1,5 +1,8 @@
 # Pending
 
+- Add a Falcon server integration: report `falcon` as the app server and start the agent in each forked worker
+- Count allocations per fiber, so concurrent requests on one thread (Falcon, Async) are no longer charged for each other's allocations
+
 # 6.3.1
 
 - Capture route names for Grape 4 (#635)

ext/allocations/allocations.c

@@ -27,13 +27,49 @@ Init_hooks(VALUE module)
 #define ALLOCATIONS_ENABLED Qtrue
 
 static __thread uint64_t endpoint_allocations;
+static ID id_fiber_allocations;
+static ID id_running_fiber;
+static ID id_running_since;
+
 void increment_allocations() {
   endpoint_allocations++;
 }
 
+static uint64_t
+stored_count(VALUE obj, ID id)
+{
+  VALUE count = rb_ivar_get(obj, id);
+  return NIL_P(count) ? 0 : NUM2ULL(count);
+}
+
+// The counter above is per OS thread, but a thread can interleave many fibers
+// (one per request under Falcon/Async). Charge the allocations made since the
+// last switch to the fiber that was running, so each fiber sees only its own.
+static void
+charge_running_fiber(void)
+{
+  VALUE thread = rb_thread_current();
+  VALUE running = rb_ivar_get(thread, id_running_fiber);
+
+  if (!NIL_P(running)) {
+    uint64_t charged = stored_count(running, id_fiber_allocations) + endpoint_allocations - stored_count(thread, id_running_since);
+    rb_ivar_set(running, id_fiber_allocations, ULL2NUM(charged));
+  }
+
+  rb_ivar_set(thread, id_running_fiber, rb_fiber_current());
+  rb_ivar_set(thread, id_running_since, ULL2NUM(endpoint_allocations));
+}
+
 static VALUE
 get_allocation_count(VALUE klass) {
-  return ULL2NUM(endpoint_allocations);
+  charge_running_fiber();
+  return ULL2NUM(stored_count(rb_fiber_current(), id_fiber_allocations));
+}
+
+static void
+fiber_switch_hook(rb_event_flag_t event, VALUE data, VALUE self, ID mid, VALUE klass)
+{
+  charge_running_fiber();
 }
 
 static void
@@ -49,7 +85,6 @@ static VALUE
 set_gc_hook(rb_event_flag_t event)
 {
     VALUE tpval;
-    // TODO - need to prevent applying the same tracepoint multiple times?
     tpval = rb_tracepoint_new(0, event, tracepoint_handler, 0);
     rb_tracepoint_enable(tpval);
 
@@ -59,7 +94,11 @@ set_gc_hook(rb_event_flag_t event)
 void
 Init_hooks(VALUE module)
 {
+    id_fiber_allocations = rb_intern("__scout_apm_allocations");
+    id_running_fiber = rb_intern("__scout_apm_running_fiber");
+    id_running_since = rb_intern("__scout_apm_running_since");
     set_gc_hook(RUBY_INTERNAL_EVENT_NEWOBJ);
+    rb_add_event_hook(fiber_switch_hook, RUBY_EVENT_FIBER_SWITCH, Qnil);
 }
 
 #endif // _WIN32

lib/scout_apm.rb

@@ -51,6 +51,7 @@ module ScoutApm
 require 'scout_apm/layer_converters/histograms'
 require 'scout_apm/layer_converters/find_layer_by_type'
 
+require 'scout_apm/server_integrations/falcon'
 require 'scout_apm/server_integrations/passenger'
 require 'scout_apm/server_integrations/puma'
 require 'scout_apm/server_integrations/rainbows'

lib/scout_apm/environment.rb

@@ -14,6 +14,7 @@ class Environment
     # I've put Thin and Webrick last as they are often used in development and included in Gemfiles
     # but less likely used in production.
     SERVER_INTEGRATIONS = [
+      ScoutApm::ServerIntegrations::Falcon.new(STDOUT_LOGGER),
       ScoutApm::ServerIntegrations::Passenger.new(STDOUT_LOGGER),
       ScoutApm::ServerIntegrations::Unicorn.new(STDOUT_LOGGER),
       ScoutApm::ServerIntegrations::Rainbows.new(STDOUT_LOGGER),

lib/scout_apm/server_integrations/falcon.rb

@@ -0,0 +1,58 @@
+module ScoutApm
+  module ServerIntegrations
+    # Falcon preloads the app in a parent process, then forks worker processes
+    # (and re-forks them whenever its supervisor restarts one). Each worker
+    # serves many concurrent requests as fibers on a single thread.
+    class Falcon
+      attr_reader :logger
+
+      def initialize(logger)
+        @logger = logger
+      end
+
+      def name
+        :falcon
+      end
+
+      def forking?
+        true
+      end
+
+      def present?
+        defined?(::Falcon) && File.basename($PROGRAM_NAME).start_with?("falcon")
+      end
+
+      # Start the agent in every forked child. Ruby >= 3.1 routes all forks
+      # through Process._fork, so this also covers workers restarted later.
+      def install
+        if ::Process.respond_to?(:_fork)
+          logger.info "Installing Falcon worker hook."
+          ::Process.singleton_class.prepend(ForkHook)
+        else
+          logger.info "Falcon worker hook requires Ruby >= 3.1, the agent will start on the first request."
+        end
+      end
+
+      # Don't start in the preloading parent: it serves no requests, and a
+      # reporting thread caught mid-request when Falcon forks can leave locks
+      # (OpenSSL, Net::HTTP) held forever in the child.
+      def found?
+        false
+      end
+
+      module ForkHook
+        def _fork
+          pid = super
+          if pid == 0
+            begin
+              ScoutApm::Agent.instance.start
+            rescue => e
+              ScoutApm::Agent.instance.logger.warn("Failed to start in Falcon worker: #{e.message}")
+            end
+          end
+          pid
+        end
+      end
+    end
+  end
+end

test/unit/allocations_test.rb

@@ -0,0 +1,40 @@
+require 'test_helper'
+
+class AllocationsTest < Minitest::Test
+  def count
+    ScoutApm::Instruments::Allocations.count
+  end
+
+  def test_counts_only_the_running_fibers_allocations
+    skip "Allocation tracking is not enabled on this platform" unless ScoutApm::Instruments::Allocations::ENABLED
+
+    heavy = Fiber.new do
+      start = count
+      5.times do
+        20_000.times { Object.new }
+        Fiber.yield
+      end
+      count - start
+    end
+
+    light_start = count
+    5.times do
+      heavy.resume
+      10.times { Object.new }
+    end
+    heavy_allocations = heavy.resume
+    light_allocations = count - light_start
+
+    assert_operator heavy_allocations, :>=, 100_000
+    assert_operator light_allocations, :<, 1_000
+  end
+
+  def test_counts_keep_increasing_within_a_fiber
+    skip "Allocation tracking is not enabled on this platform" unless ScoutApm::Instruments::Allocations::ENABLED
+
+    before = count
+    1_000.times { Object.new }
+
+    assert_operator count - before, :>=, 1_000
+  end
+end

test/unit/server_integrations/falcon_test.rb

@@ -0,0 +1,56 @@
+require 'test_helper'
+
+class FalconTest < Minitest::Test
+  def setup
+    @program_name = $PROGRAM_NAME
+    @integration = ScoutApm::ServerIntegrations::Falcon.new(ScoutApm::Agent.instance.logger)
+    Object.const_set(:Falcon, Module.new) unless defined?(::Falcon)
+  end
+
+  def teardown
+    $PROGRAM_NAME = @program_name
+    Object.send(:remove_const, :Falcon) if defined?(::Falcon)
+  end
+
+  def test_present_when_running_under_falcon
+    $PROGRAM_NAME = "/app/bin/falcon"
+
+    assert @integration.present?
+    assert_equal :falcon, @integration.name
+  end
+
+  def test_not_present_in_other_processes_that_load_falcon
+    $PROGRAM_NAME = "/app/bin/sidekiq"
+
+    refute @integration.present?
+  end
+
+  def test_defers_starting_in_the_preloading_parent
+    assert @integration.forking?
+    refute @integration.found?
+  end
+
+  def test_starts_the_agent_in_a_forked_worker
+    skip "Process._fork requires Ruby 3.1" unless Process.respond_to?(:_fork)
+
+    agent = ScoutApm::Agent.instance
+    agent.define_singleton_method(:start) { |*| @started_by_fork = true }
+    @integration.install
+
+    reader, writer = IO.pipe
+    pid = fork do
+      reader.close
+      writer.write(agent.instance_variable_get(:@started_by_fork).inspect)
+      writer.close
+      exit!(0)
+    end
+    writer.close
+    started_in_child = reader.read
+    Process.wait(pid)
+
+    assert_equal "true", started_in_child
+    assert_nil agent.instance_variable_get(:@started_by_fork)
+  ensure
+    agent.singleton_class.send(:remove_method, :start)
+  end
+end