Hacked mongo instrumentation

module Mongoid
  class LogSubscriber < ActiveSupport::LogSubscriber
    def self.runtime=(value)
      Thread.current["mongoid_mongo_runtime"] = value
    end

    def self.runtime
      Thread.current["mongoid_mongo_runtime"] ||= 0
    end

    def self.reset_runtime
      rt, self.runtime = runtime, 0
      rt
    end

    def mongo(event)
      self.class.runtime += event.duration
    end
  end
end

Mongoid::LogSubscriber.attach_to :mongoid

module Mongoid
  module Railties
    module ControllerRuntime
      extend ActiveSupport::Concern

    protected

      attr_internal :mongo_runtime

      def cleanup_view_runtime
        mongo_rt_before_render = Mongoid::LogSubscriber.reset_runtime
        runtime = super
        mongo_rt_after_render = Mongoid::LogSubscriber.reset_runtime
        self.mongo_runtime = mongo_rt_before_render + mongo_rt_after_render
        runtime - mongo_rt_after_render
      end

      def append_info_to_payload(payload)
        super
        payload[:mongo_runtime] = mongo_runtime
      end

      module ClassMethods
        def log_process_action(payload)
          messages, mongo_runtime = super, payload[:mongo_runtime]
          messages << ("Mongoid: %.1fms" % mongo_runtime.to_f) if mongo_runtime
          messages
        end
      end
    end
  end
end

Mongo::Collection.module_eval do
  [:find, :save, :insert, :update].each do |m|
    class_eval %{def #{m}_with_instrumentation(*args)
      ActiveSupport::Notifications.instrumenter.instrument "mongo.mongoid", :name => "#{m}" do
        #{m}_without_instrumentation(*args)
      end
    end
    }

    alias_method_chain m, :instrumentation
  end
end

Mongo::DB.module_eval do
  [:command].each do |m|
    class_eval %{def #{m}_with_instrumentation(*args)
      ActiveSupport::Notifications.instrumenter.instrument "mongo.mongoid", :name => "#{m}" do
        #{m}_without_instrumentation(*args)
      end
    end
    }

    alias_method_chain m, :instrumentation
  end
end

ActiveSupport.on_load(:action_controller) do
  include Mongoid::Railties::ControllerRuntime
end