Experimental Mongo Instrumentation for Rails 3

instrumentation.rb

module Mongo
  module Instrumentation
    def self.instrument(clazz, *methods)
      clazz.module_eval do
        methods.each do |m|
          class_eval %{def #{m}_with_instrumentation(*args, &block)
            ActiveSupport::Notifications.instrumenter.instrument "mongo.mongo", :name => "#{m}" do
              #{m}_without_instrumentation(*args, &block)
            end
          end
          }

          alias_method_chain m, :instrumentation
        end
      end
    end

    class Railtie < Rails::Railtie
      initializer "mongo.instrumentation" do |app|
        Mongo::Instrumentation.instrument Mongo::Connection, :send_message, :send_message_with_safe_check, :receive_message
        
        #Mongo::Instrumentation.instrument Mongo::Collection, :find, :save, :insert, :update
        #Mongo::Instrumentation.instrument Mongo::DB, :command

        ActiveSupport.on_load(:action_controller) do
          include Mongo::Instrumentation::ControllerRuntime
        end

        Mongo::Instrumentation::LogSubscriber.attach_to :mongo
      end
    end

    module ControllerRuntime
      extend ActiveSupport::Concern

      protected

      attr_internal :mongo_runtime

      def cleanup_view_runtime
        mongo_rt_before_render = Mongo::Instrumentation::LogSubscriber.reset_runtime
        runtime = super
        mongo_rt_after_render = Mongo::Instrumentation::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 << ("Mongo: %.1fms" % mongo_runtime.to_f) if mongo_runtime
          messages
        end
      end
    end

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

      def self.runtime
        Thread.current["mongo_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
end

Hi Tom,
I would have commented on your blog but it seems it's not set up for comments?
Anyway, I was wondering if you think this code is supposed to do the trick for apps that use MongoMapper as well? I'm guessing the answer is no because I tried it and the output looks okay, except that the values are not too credible (much lower than I would expect). Is you app exclusively using the raw Mongo ruby library?
Thanks,
Sean

The values may well be too low. I don't think I'm capturing every call that's made for some reason (maybe more methods need instrumenting?). I'm using Mongoid, but as it hooks onto the underlying mongo library I'd have thought it would work for MongoMapper too. I'll be looking into improving it based on the feedback I'm getting.

Another thing to bear in mind is that currently it's only attempting to measure access to the underlying mongo datastore. Mongo Mapper and Mongoid both do a lot of other work before returning models to your application. For example, if you have a page request that accidentally loads the same record by id 1000 times, mongo itself is likely to be very quick, but the cost of instantiating 1000 models will be significant (each model may be made of 100 different objects, so you're effectively instantiating 100,000 things).

sohara: I've changed it to instrument only methods on the Mongo::Connection and it seems to be giving me better results. Let me know if that's any better for you

I'm getting a "block not supplied" argument error on rails 3.0.3? What version are you on?

My bad. An extra comma creeped in there when I cut and paste into the gist. Should work now (I'm on 3.0.3 too)

That's working for me. Yep, the stats seem much more realistic with this change (around 10 * more than before). Thanks!

This is great. I was just about to suggest not instrumenting :find :save :insert :update, but you beat me to it. Thanks! I hope you comment every time you make an improvement.

This code is now a gem, 'mongo-rails-instrumentation', and available on github here: https://github.com/tomafro/mongo-rails-instrumentation.