Correctly track hijacked requests

lib/hijack.rb

@@ -8,27 +8,34 @@ module Hijack
   def hijack(&blk)
     controller_class = self.class
 
-    #env = request.env.dup
-    #request_copy = ActionDispatch::Request.new(env)
-    request_copy = self.request
-
     if hijack = request.env['rack.hijack']
       io = hijack.call
 
+      request.env['discourse.request_tracker.skip'] = true
+      request_tracker = request.env['discourse.request_tracker']
+
       # in prd the env object is re-used
       # make a copy of all strings
       env_copy = {}
       request.env.each do |k, v|
         env_copy[k] = v if String === v
       end
-
-      request_copy = ActionDispatch::Request.new(env_copy)
+      # we require that for request initialization
+      env_copy["rack.input"] = StringIO.new
 
       # params is generated per request so we can simply reuse it
       params_copy = params
 
+      env_copy["action_dispatch.request.parameters"] = params_copy
+
+      request_copy = ActionDispatch::Request.new(env_copy)
+
+      transfer_timings = MethodProfiler.transfer if defined? MethodProfiler
+
       Scheduler::Defer.later("hijack work") do
 
+        MethodProfiler.start(transfer_timings) if defined? MethodProfiler
+
         begin
           # do this first to confirm we have a working connection
           # before doing any work
@@ -74,6 +81,12 @@ module Hijack
           io.close
         rescue Errno::EPIPE, IOError
           # happens if client terminated before we responded, ignore
+        ensure
+          if request_tracker
+            status = instance.status rescue 500
+            timings = MethodProfiler.stop if defined? MethodProfiler
+            request_tracker.log_request_info(env_copy, [status, headers || {}, []], timings)
+          end
         end
       end
       # not leaked out, we use 418 ... I am a teapot to denote that we are hijacked

lib/method_profiler.rb

@@ -25,8 +25,14 @@ class MethodProfiler
     klass.class_eval patches
   end
 
-  def self.start
-    Thread.current[:_method_profiler] = {
+  def self.transfer
+    result = Thread.current[:_method_profiler]
+    Thread.current[:_method_profiler] = nil
+    result
+  end
+
+  def self.start(transfer = nil)
+    Thread.current[:_method_profiler] = transfer || {
       __start: Process.clock_gettime(Process::CLOCK_MONOTONIC)
     }
   end

lib/middleware/request_tracker.rb

@@ -105,12 +105,7 @@ class Middleware::RequestTracker
 
   end
 
-  def call(env)
-    MethodProfiler.start if @@detailed_request_loggers
-    result = @app.call(env)
-    info = MethodProfiler.stop if @@detailed_request_loggers
-    result
-  ensure
+  def log_request_info(env, result, info)
 
     # we got to skip this on error ... its just logging
     data = self.class.get_data(env, result, info) rescue nil
@@ -130,6 +125,16 @@ class Middleware::RequestTracker
 
   end
 
+  def call(env)
+    env["discourse.request_tracker"] = self
+    MethodProfiler.start if @@detailed_request_loggers
+    result = @app.call(env)
+    info = MethodProfiler.stop if @@detailed_request_loggers
+    result
+  ensure
+    log_request_info(env, result, info) unless env["discourse.request_tracker.skip"]
+  end
+
   def log_later(data, host)
     Scheduler::Defer.later("Track view", _db = nil) do
       self.class.log_request_on_site(data, host)

spec/components/hijack_spec.rb

@@ -6,15 +6,16 @@ describe Hijack do
 
     include Hijack
 
-    def initialize
+    def initialize(env = {})
       @io = StringIO.new
-      self.request = ActionController::TestRequest.new({
+
+      env.merge!(
         "rack.hijack" => lambda { @io },
         "rack.input" => StringIO.new
-      },
-        nil,
-        nil
       )
+
+      self.request = ActionController::TestRequest.new(env, nil, nil)
+
       # we need this for the 418
       self.response = ActionDispatch::Response.new
     end
@@ -29,6 +30,43 @@ describe Hijack do
     Hijack::Tester.new
   end
 
+  context "Request Tracker integration" do
+    let :logger do
+      lambda do |env, data|
+        @calls += 1
+        @status = data[:status]
+        @total = data[:timing][:total_duration]
+      end
+    end
+
+    before do
+      Middleware::RequestTracker.register_detailed_request_logger logger
+      @calls = 0
+    end
+
+    after do
+      Middleware::RequestTracker.unregister_detailed_request_logger logger
+    end
+
+    it "can properly track execution" do
+      app = lambda do |env|
+        tester = Hijack::Tester.new(env)
+        tester.hijack_test do
+          render body: "hello", status: 201
+        end
+      end
+
+      env = {}
+      middleware = Middleware::RequestTracker.new(app)
+
+      middleware.call(env)
+
+      expect(@calls).to eq(1)
+      expect(@status).to eq(201)
+      expect(@status).to be > 0
+    end
+  end
+
   it "dupes the request params and env" do
     orig_req = tester.request
     copy_req = nil

spec/components/method_profiler_spec.rb

@@ -0,0 +1,25 @@
+require 'rails_helper'
+require_dependency 'method_profiler'
+
+describe MethodProfiler do
+  class Sneetch
+    def beach
+    end
+  end
+
+  it "can transfer data between threads" do
+    MethodProfiler.patch(Sneetch, [:beach], :at_beach)
+
+    MethodProfiler.start
+    Sneetch.new.beach
+    data = MethodProfiler.transfer
+    result = nil
+    Thread.new do
+      MethodProfiler.start(data)
+      Sneetch.new.beach
+      result = MethodProfiler.stop
+    end.join
+
+    expect(result[:at_beach][:calls]).to eq(2)
+  end
+end