Add db_runtime instrumentation for ActiveGraph

lib/active_graph/railtie.rb

@@ -64,7 +64,13 @@ def empty_config
         config.app_middleware.insert_after ::ActionDispatch::Callbacks, ActiveGraph::Migrations::CheckPending
       end
     end
+    initializer 'activegraph.controller_runtime' do
+      require 'active_graph/railties/controller_runtime'
 
+      ActiveSupport.on_load(:action_controller) do
+        include ActiveGraph::Railties::ControllerRuntime
+      end
+    end
     initializer "activegraph.deprecator" do |app|
       app.deprecators[:activegraph] = ActiveGraph.deprecator
     end

lib/active_graph/railties/controller_runtime.rb

@@ -0,0 +1,22 @@
+require 'active_graph/runtime_registry'
+
+module ActiveGraph
+  module Railties
+    module ControllerRuntime
+      extend ActiveSupport::Concern
+
+      private
+
+      def process_action(*)
+        ActiveGraph::RuntimeRegistry.reset
+        super
+      end
+
+      def append_info_to_payload(payload)
+        super
+
+        payload[:db_runtime] = (payload[:db_runtime] || 0) + ActiveGraph::RuntimeRegistry.reset
+      end
+    end
+  end
+end

lib/active_graph/runtime_registry.rb

@@ -0,0 +1,38 @@
+require 'active_support/isolated_execution_state'
+
+module ActiveGraph
+  module RuntimeRegistry
+    class Stats
+      attr_accessor :db_runtime
+
+      def initialize
+        @db_runtime = 0.0
+      end
+
+      def reset_runtime
+        db_runtime_was = @db_runtime
+        @db_runtime = 0.0
+        db_runtime_was
+      end
+    end
+
+    extend self
+
+    def call(_name, start, finish, _id, _payload)
+      stats.db_runtime += (finish - start) * 1_000.0
+    end
+
+    def stats
+      ActiveSupport::IsolatedExecutionState[:active_graph_runtime] ||= Stats.new
+    end
+
+    def reset
+      stats.reset_runtime
+    end
+  end
+
+  ActiveSupport::Notifications.monotonic_subscribe(
+    'neo4j.core.bolt.request',
+    RuntimeRegistry
+  )
+end

spec/unit/railties/controller_runtime_spec.rb

@@ -0,0 +1,72 @@
+require 'active_graph/railties/controller_runtime'
+
+RSpec.describe ActiveGraph::Railties::ControllerRuntime do
+  let(:controller_class) do
+    Class.new(ActionController::Base) do
+      include ActiveGraph::Railties::ControllerRuntime
+
+      def index
+      end
+    end
+  end
+
+  let(:controller) { controller_class.new }
+
+  before do
+    ActiveGraph::RuntimeRegistry.reset
+  end
+
+  after do
+    ActiveGraph::RuntimeRegistry.reset
+  end
+
+  describe '#append_info_to_payload' do
+    it 'adds database runtime to the payload' do
+      ActiveGraph::RuntimeRegistry.call(
+        'neo4j.core.bolt.request',
+        1.0,
+        1.25,
+        'id',
+        {}
+      )
+
+      payload = {}
+
+      controller.send(:append_info_to_payload, payload)
+
+      expect(payload[:db_runtime]).to eq(250.0)
+    end
+
+    it 'preserves existing database runtime in the payload' do
+      ActiveGraph::RuntimeRegistry.call(
+        'neo4j.core.bolt.request',
+        1.0,
+        1.25,
+        'id',
+        {}
+      )
+
+      payload = { db_runtime: 100.0 }
+
+      controller.send(:append_info_to_payload, payload)
+
+      expect(payload[:db_runtime]).to eq(350.0)
+    end
+  end
+
+  describe '#process_action' do
+    it 'resets database runtime before processing the action' do
+      ActiveGraph::RuntimeRegistry.call(
+        'neo4j.core.bolt.request',
+        1.0,
+        1.25,
+        'id',
+        {}
+      )
+
+      expect(ActiveGraph::RuntimeRegistry).to receive(:reset).and_call_original
+
+      controller.send(:process_action, :index)
+    end
+  end
+end

spec/unit/runtime_registry_spec.rb

@@ -0,0 +1,36 @@
+require 'active_graph/runtime_registry'
+
+RSpec.describe ActiveGraph::RuntimeRegistry do
+  before do
+    described_class.reset
+  end
+
+  after do
+    described_class.reset
+  end
+
+  describe '.call' do
+    it 'accumulates database runtime' do
+      described_class.call('neo4j.core.bolt.request', 1.0, 1.25, 'id', {})
+
+      expect(described_class.reset).to eq(250.0)
+    end
+
+    it 'accumulates runtime across multiple database requests' do
+      described_class.call('neo4j.core.bolt.request', 1.0, 1.25, 'id', {})
+      described_class.call('neo4j.core.bolt.request', 2.0, 2.5, 'id', {})
+
+      expect(described_class.reset).to eq(750.0)
+    end
+  end
+
+  describe '.reset' do
+    it 'clears the accumulated runtime' do
+      described_class.call('neo4j.core.bolt.request', 1.0, 1.25, 'id', {})
+
+      described_class.reset
+
+      expect(described_class.reset).to eq(0.0)
+    end
+  end
+end