From 7270be0a662e2539fc24ce42886b52bf7090f7f0 Mon Sep 17 00:00:00 2001
From: Joshua Young <djry1999@gmail.com>
Date: Sat, 14 Sep 2024 08:25:56 +1000
Subject: [PATCH] test

---
 lib/puma_cfs_rails/middleware.rb | 47 ++++++++++++++++++++++++++------
 lib/puma_cfs_rails/railtie.rb    |  4 +--
 2 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/lib/puma_cfs_rails/middleware.rb b/lib/puma_cfs_rails/middleware.rb
index 6fb9ee7..14ff6f6 100644
--- a/lib/puma_cfs_rails/middleware.rb
+++ b/lib/puma_cfs_rails/middleware.rb
@@ -5,6 +5,11 @@
 module PumaCFSRails
   class Middleware
     EMPTY_KEY = ""
+    PATH_INFO = "PATH_INFO"
+    REQUEST_METHOD = "REQUEST_METHOD"
+    MIN_VRUNTIME = "MIN_VRUNTIME"
+    SERVER_TIMING = ActionDispatch::Constants::SERVER_TIMING
+    PUMA_CFS_RAILS = "PUMA-CFS-RAILS"
 
     def initialize(app)
       @app = app
@@ -13,16 +18,31 @@ def initialize(app)
 
     def call(env)
       path_params = begin
-        Rails.application.routes.recognize_path(env["PATH_INFO"], method: env["REQUEST_METHOD"])
+        Rails.application.routes.recognize_path(env[PATH_INFO], method: env[REQUEST_METHOD])
       rescue ActionController::RoutingError
         {}
       end
       key = key_from(path_params)
 
-      env["MIN_VRUNTIME"] = @min_vruntimes[key] ||= 0
-      duration, response = with_time_diff { @app.call(env) }
-      unless key == EMPTY_KEY || (@min_vruntimes[key] > 0 && @min_vruntimes[key] < duration)
-        @min_vruntimes[key] = duration
+      env[MIN_VRUNTIME] = @min_vruntimes[key] ||= 0
+
+      response = @app.call(env)
+      headers = response[1]
+
+      if key == EMPTY_KEY
+        headers[PUMA_CFS_RAILS] = "Could not determine controller and action from path"
+      else
+        vruntime = cpu_time_sum(parsed_server_timing(headers))
+        pp vruntime
+        if vruntime > 0
+          headers[PUMA_CFS_RAILS] = "Success"
+        else
+          headers[PUMA_CFS_RAILS] = "No server timing data, check if `config.server_timings = true`"
+        end
+
+        if vruntime < @min_vruntimes[key]
+          @min_vruntimes[key] = vruntime
+        end
       end
 
       response
@@ -38,10 +58,19 @@ def key_from(path_params)
       end
     end
 
-    def with_time_diff
-      before = Time.now
-      result = yield
-      [Time.now - before, result]
+    def parsed_server_timing(headers)
+      (headers[SERVER_TIMING] || "").split(', ').to_h do |entry|
+        key, value = entry.split(';dur=')
+        [key, value.to_f]
+      end
+    end
+
+    def cpu_time_sum(timings)
+      (
+        (timings["start_processing.action_controller"] || 0) +
+        (timings["render_template.action_view"] || 0) +
+        (timings["process_action.action_controller"] || 0)
+      ) - (timings["sql.active_record"] || 0)
     end
   end
 end
diff --git a/lib/puma_cfs_rails/railtie.rb b/lib/puma_cfs_rails/railtie.rb
index 64acf6b..195dd96 100644
--- a/lib/puma_cfs_rails/railtie.rb
+++ b/lib/puma_cfs_rails/railtie.rb
@@ -3,9 +3,9 @@
 require_relative 'middleware'
 
 module PumaCFSRails
-  class Railtie < Rails::Railtie
+  class Railtie < ::Rails::Railtie
     initializer "puma_cfs_rails.use_middleware" do |app|
-      app.middleware.use PumaCFSRails::Middleware
+      app.middleware.insert_before ::ActionDispatch::ServerTiming, ::PumaCFSRails::Middleware
     end
   end
 end