Fix view runtime for controllers with async queries

Author: fatkodima

activerecord/lib/active_record/railties/controller_runtime.rb

@@ -37,9 +37,10 @@ def cleanup_view_runtime
             db_rt_before_render = ActiveRecord::RuntimeRegistry.reset
             self.db_runtime = (db_runtime || 0) + db_rt_before_render
             runtime = super
+            queries_rt = ActiveRecord::RuntimeRegistry.sql_runtime - ActiveRecord::RuntimeRegistry.async_sql_runtime
             db_rt_after_render = ActiveRecord::RuntimeRegistry.reset
             self.db_runtime += db_rt_after_render
-            runtime - db_rt_after_render
+            runtime - queries_rt
           else
             super
           end

activerecord/lib/active_record/runtime_registry.rb

@@ -17,13 +17,26 @@ def sql_runtime=(runtime)
       ActiveSupport::IsolatedExecutionState[:active_record_sql_runtime] = runtime
     end
 
+    def async_sql_runtime
+      ActiveSupport::IsolatedExecutionState[:active_record_async_sql_runtime] ||= 0.0
+    end
+
+    def async_sql_runtime=(runtime)
+      ActiveSupport::IsolatedExecutionState[:active_record_async_sql_runtime] = runtime
+    end
+
     def reset
       rt, self.sql_runtime = sql_runtime, 0.0
+      self.async_sql_runtime = 0.0
       rt
     end
   end
 end
 
 ActiveSupport::Notifications.monotonic_subscribe("sql.active_record") do |name, start, finish, id, payload|
-  ActiveRecord::RuntimeRegistry.sql_runtime += (finish - start)
+  runtime = finish - start
+  if payload[:async]
+    ActiveRecord::RuntimeRegistry.async_sql_runtime += (runtime - payload[:lock_wait])
+  end
+  ActiveRecord::RuntimeRegistry.sql_runtime += runtime
 end