Added process action hook

Andrew Kane
1 parent 8d33e220
Showing 1 changed file with 12 additions and 2 deletions Show diff stats
lib/searchkick/logging.rb
@@ -15,6 +15,7 @@ module Searchkick
     alias_method_chain :execute, :instrumentation
   end
+  # https://github.com/rails/rails/blob/master/activerecord/lib/active_record/log_subscriber.rb
   class LogSubscriber < ActiveSupport::LogSubscriber
     def self.runtime=(value)
       Thread.current[:searchkick_runtime] = value
@@ -43,11 +44,22 @@ module Searchkick
     end
   end
+  # https://github.com/rails/rails/blob/master/activerecord/lib/active_record/railties/controller_runtime.rb
   module ControllerRuntime
     extend ActiveSupport::Concern
+    protected
+
     attr_internal :searchkick_runtime
+    def process_action(action, *args)
+      # We also need to reset the runtime before each action
+      # because of queries in middleware or in cases we are streaming
+      # and it won't be cleaned up by the method below.
+      Searchkick::LogSubscriber.reset_runtime
+      super
+    end
+
     def cleanup_view_runtime
       searchkick_rt_before_render = Searchkick::LogSubscriber.reset_runtime
       runtime = super
@@ -56,8 +68,6 @@ module Searchkick
       runtime - searchkick_rt_after_render
     end
-    protected
-
     def append_info_to_payload(payload)
       super
       payload[:searchkick_runtime] = (searchkick_runtime || 0) + Searchkick::LogSubscriber.reset_runtime