[RAILS] Added ActiveSupport instrumentation integration

Require the component in your `application.rb` file:

    require 'elasticsearch/rails/instrumentation'

You should see an output like this in your application log in development environment:

    Article Search (321.3ms) { index: "articles", type: "article", body: { query: ... } }

Also, the total duration of the request to Elasticsearch is displayed in the Rails request breakdown:

    Completed 200 OK in 615ms (Views: 230.9ms | ActiveRecord: 0.0ms | Elasticsearch: 321.3ms)

A special component for the Lograge logger implementation is provided as well.

See:

* http://edgeguides.rubyonrails.org/active_support_instrumentation.html
* http://api.rubyonrails.org/classes/ActiveSupport/Notifications.html
* https://github.com/roidrage/lograge

Author: karmi

elasticsearch-rails/README.md

@@ -46,6 +46,34 @@ Run this command to display usage instructions:
 $ bundle exec rake -D elasticsearch
 ```
 
+### ActiveSupport Instrumentation
+
+To display information about the search request (duration, search definition) during development,
+and to include the information in the Rails log file, require the component in your `application.rb` file:
+
+```ruby
+require 'elasticsearch/rails/instrumentation'
+```
+
+You should see an output like this in your application log in development environment:
+
+    Article Search (321.3ms) { index: "articles", type: "article", body: { query: ... } }
+
+Also, the total duration of the request to Elasticsearch is displayed in the Rails request breakdown:
+
+    Completed 200 OK in 615ms (Views: 230.9ms | ActiveRecord: 0.0ms | Elasticsearch: 321.3ms)
+
+There's a special component for the [Lograge](https://github.com/roidrage/lograge) logger.
+Require the component in your `application.rb` file (and set `config.lograge.enabled`):
+
+```ruby
+require 'elasticsearch/rails/lograge'
+```
+
+You should see the duration of the request to Elasticsearch as part of each log event:
+
+    method=GET path=/search ... status=200 duration=380.89 view=99.64 db=0.00 es=279.37
+
 ### Rails Application Templates
 
 You can generate a fully working example Ruby on Rails application, with an `Article` model and a search form,

elasticsearch-rails/elasticsearch-rails.gemspec

@@ -28,6 +28,9 @@ Gem::Specification.new do |s|
 
   s.add_development_dependency "oj"
   s.add_development_dependency "rails",  "> 3.0"
+
+  s.add_development_dependency "lograge"
+
   s.add_development_dependency "shoulda-context"
   s.add_development_dependency "mocha"
   s.add_development_dependency "turn"

elasticsearch-rails/lib/elasticsearch/rails/instrumentation.rb

@@ -0,0 +1,36 @@
+require 'elasticsearch/rails/instrumentation/railtie'
+require 'elasticsearch/rails/instrumentation/publishers'
+
+module Elasticsearch
+  module Rails
+
+    # This module adds support for displaying statistics about search duration in the Rails application log
+    # by integrating with the `ActiveSupport::Notifications` framework and `ActionController` logger.
+    #
+    # == Usage
+    #
+    # Require the component in your `application.rb` file:
+    #
+    #     require 'elasticsearch/rails/instrumentation'
+    #
+    # You should see an output like this in your application log in development environment:
+    #
+    #     Article Search (321.3ms) { index: "articles", type: "article", body: { query: ... } }
+    #
+    # Also, the total duration of the request to Elasticsearch is displayed in the Rails request breakdown:
+    #
+    #     Completed 200 OK in 615ms (Views: 230.9ms | ActiveRecord: 0.0ms | Elasticsearch: 321.3ms)
+    #
+    # @note The displayed duration includes the HTTP transfer -- the time it took Elasticsearch
+    #       to process your request is available in the `response.took` property.
+    #
+    # @see Elasticsearch::Rails::Instrumentation::Publishers
+    # @see Elasticsearch::Rails::Instrumentation::Railtie
+    #
+    # @see http://api.rubyonrails.org/classes/ActiveSupport/Notifications.html
+    #
+    #
+    module Instrumentation
+    end
+  end
+end

elasticsearch-rails/lib/elasticsearch/rails/instrumentation/controller_runtime.rb

@@ -0,0 +1,41 @@
+require 'active_support/core_ext/module/attr_internal'
+
+module Elasticsearch
+  module Rails
+    module Instrumentation
+
+      # Hooks into ActionController to display Elasticsearch runtime
+      #
+      # @see https://github.com/rails/rails/blob/master/activerecord/lib/active_record/railties/controller_runtime.rb
+      #
+      module ControllerRuntime
+        extend ActiveSupport::Concern
+
+        protected
+
+        attr_internal :elasticsearch_runtime
+
+        def cleanup_view_runtime
+          elasticsearch_rt_before_render = Elasticsearch::Rails::Instrumentation::LogSubscriber.reset_runtime
+          runtime = super
+          elasticsearch_rt_after_render = Elasticsearch::Rails::Instrumentation::LogSubscriber.reset_runtime
+          self.elasticsearch_runtime = elasticsearch_rt_before_render + elasticsearch_rt_after_render
+          runtime - elasticsearch_rt_after_render
+        end
+
+        def append_info_to_payload(payload)
+          super
+          payload[:elasticsearch_runtime] = (elasticsearch_runtime || 0) + Elasticsearch::Rails::Instrumentation::LogSubscriber.reset_runtime
+        end
+
+        module ClassMethods
+          def log_process_action(payload)
+            messages, elasticsearch_runtime = super, payload[:elasticsearch_runtime]
+            messages << ("Elasticsearch: %.1fms" % elasticsearch_runtime.to_f) if elasticsearch_runtime
+            messages
+          end
+        end
+      end
+    end
+  end
+end

elasticsearch-rails/lib/elasticsearch/rails/instrumentation/log_subscriber.rb

@@ -0,0 +1,41 @@
+module Elasticsearch
+  module Rails
+    module Instrumentation
+
+      # A log subscriber to attach to Elasticsearch related events
+      #
+      # @see https://github.com/rails/rails/blob/master/activerecord/lib/active_record/log_subscriber.rb
+      #
+      class LogSubscriber < ActiveSupport::LogSubscriber
+        def self.runtime=(value)
+          Thread.current["elasticsearch_runtime"] = value
+        end
+
+        def self.runtime
+          Thread.current["elasticsearch_runtime"] ||= 0
+        end
+
+        def self.reset_runtime
+          rt, self.runtime = runtime, 0
+          rt
+        end
+
+        # Intercept `search.elasticsearch` events, and display them in the Rails log
+        #
+        def search(event)
+          self.class.runtime += event.duration
+          return unless logger.debug?
+
+          payload = event.payload
+          name    = "#{payload[:klass]} #{payload[:name]} (#{event.duration.round(1)}ms)"
+          search  = payload[:search].inspect.gsub(/:(\w+)=>/, '\1: ')
+
+          debug %Q|  #{color(name, GREEN, true)} #{colorize_logging ? "\e[2m#{search}\e[0m" : search}|
+        end
+      end
+
+    end
+  end
+end
+
+Elasticsearch::Rails::Instrumentation::LogSubscriber.attach_to :elasticsearch

elasticsearch-rails/lib/elasticsearch/rails/instrumentation/publishers.rb

@@ -0,0 +1,32 @@
+module Elasticsearch
+  module Rails
+    module Instrumentation
+      module Publishers
+
+        # Wraps the `SearchRequest` methods to perform the instrumentation
+        #
+        # @see SearchRequest#execute_with_instrumentation!
+        # @see http://api.rubyonrails.org/classes/ActiveSupport/Notifications.html
+        #
+        module SearchRequest
+          extend ActiveSupport::Concern
+
+          included do
+            alias_method_chain :execute!, :instrumentation
+          end
+
+          # Wrap `Search#execute!` and perform instrumentation
+          #
+          def execute_with_instrumentation!
+            ActiveSupport::Notifications.instrument "search.elasticsearch",
+              name:   'Search',
+              klass:  (self.klass.is_a?(Elasticsearch::Model::Proxy::ClassMethodsProxy) ? self.klass.target.to_s : self.klass.to_s),
+              search: self.definition do
+              execute_without_instrumentation!
+            end
+          end
+        end
+      end
+    end
+  end
+end

elasticsearch-rails/lib/elasticsearch/rails/instrumentation/railtie.rb

@@ -0,0 +1,27 @@
+module Elasticsearch
+  module Rails
+    module Instrumentation
+
+      # Rails initializer class to require Elasticsearch::Rails::Instrumentation files,
+      # set up Elasticsearch::Model and hook into ActionController to display Elasticsearch-related duration
+      #
+      # @see http://edgeguides.rubyonrails.org/active_support_instrumentation.html
+      #
+      class Railtie < ::Rails::Railtie
+        initializer "elasticsearch.instrumentation" do |app|
+          require 'elasticsearch/rails/instrumentation/log_subscriber'
+          require 'elasticsearch/rails/instrumentation/controller_runtime'
+
+          Elasticsearch::Model::Searching::SearchRequest.class_eval do
+            include Elasticsearch::Rails::Instrumentation::Publishers::SearchRequest
+          end
+
+          ActiveSupport.on_load(:action_controller) do
+            include Elasticsearch::Rails::Instrumentation::ControllerRuntime
+          end
+        end
+      end
+
+    end
+  end
+end

elasticsearch-rails/lib/elasticsearch/rails/lograge.rb

@@ -0,0 +1,40 @@
+module Elasticsearch
+  module Rails
+    module Lograge
+
+      # Rails initializer class to require Elasticsearch::Rails::Instrumentation files,
+      # set up Elasticsearch::Model and add Lograge configuration to display Elasticsearch-related duration
+      #
+      # Require the component in your `application.rb` file and enable Lograge:
+      #
+      #     require 'elasticsearch/rails/lograge'
+      #
+      # You should see the full duration of the request to Elasticsearch as part of each log event:
+      #
+      #     method=GET path=/search ... status=200 duration=380.89 view=99.64 db=0.00 es=279.37
+      #
+      # @see https://github.com/roidrage/lograge
+      #
+      class Railtie < ::Rails::Railtie
+        initializer "elasticsearch.lograge" do |app|
+          require 'elasticsearch/rails/instrumentation/publishers'
+          require 'elasticsearch/rails/instrumentation/log_subscriber'
+          require 'elasticsearch/rails/instrumentation/controller_runtime'
+
+          config.lograge.custom_options = lambda do |event|
+            { es: event.payload[:elasticsearch_runtime].to_f.round(2) }
+          end
+
+          Elasticsearch::Model::Searching::SearchRequest.class_eval do
+            include Elasticsearch::Rails::Instrumentation::Publishers::SearchRequest
+          end
+
+          ActiveSupport.on_load(:action_controller) do
+            include Elasticsearch::Rails::Instrumentation::ControllerRuntime
+          end
+        end
+      end
+
+    end
+  end
+end

elasticsearch-rails/test/unit/instrumentation/instrumentation_test.rb

@@ -0,0 +1,61 @@
+require 'test_helper'
+
+require 'rails/railtie'
+require 'active_support/log_subscriber/test_helper'
+
+require 'elasticsearch/rails/instrumentation'
+
+class Elasticsearch::Rails::InstrumentationTest < Test::Unit::TestCase
+  include ActiveSupport::LogSubscriber::TestHelper
+
+  context "ActiveSupport::Instrumentation integration" do
+    class ::DummyInstrumentationModel
+      extend Elasticsearch::Model::Searching::ClassMethods
+
+      def self.index_name;    'foo'; end
+      def self.document_type; 'bar'; end
+    end
+
+    RESPONSE = { 'took' => '5ms', 'hits' => { 'total' => 123, 'max_score' => 456, 'hits' => [] } }
+
+    setup do
+      @search   = Elasticsearch::Model::Searching::SearchRequest.new ::DummyInstrumentationModel, '*'
+      @response = Elasticsearch::Model::Response::Response.new ::DummyInstrumentationModel, @search
+
+      @client = stub('client', search: @response)
+      DummyInstrumentationModel.stubs(:client).returns(@client)
+
+      Elasticsearch::Rails::Instrumentation::Railtie.run_initializers
+    end
+
+    should "wrap SearchRequest#execute! with instrumentation" do
+      s = Elasticsearch::Model::Searching::SearchRequest.new ::DummyInstrumentationModel, 'foo'
+      assert_respond_to s, :execute_without_instrumentation!
+      assert_respond_to s, :execute_with_instrumentation!
+    end
+
+    should "publish the notification" do
+      @query = { query: { match: { foo: 'bar' } } }
+
+      ActiveSupport::Notifications.expects(:instrument).with do |name, payload|
+        assert_equal "search.elasticsearch", name
+        assert_equal 'DummyInstrumentationModel', payload[:klass]
+        assert_equal @query, payload[:search][:body]
+      end
+
+      s = ::DummyInstrumentationModel.search @query
+      s.response
+    end
+
+    should "print the debug information to the Rails log" do
+      s = ::DummyInstrumentationModel.search query: { match: { moo: 'bam' } }
+      s.response
+
+      logged = @logger.logged(:debug).first
+
+      assert_not_nil logged
+      assert_match /DummyInstrumentationModel Search \(\d+\.\d+ms\)/,  logged
+      assert_match /body\: \{query\: \{match\: \{moo\: "bam"\}\}\}\}/, logged
+    end
+  end
+end

elasticsearch-rails/test/unit/instrumentation/lograge_test.rb

@@ -0,0 +1,20 @@
+require 'test_helper'
+
+require 'rails/railtie'
+require 'lograge'
+
+require 'elasticsearch/rails/lograge'
+
+class Elasticsearch::Rails::LogrageTest < Test::Unit::TestCase
+  context "Lograge integration" do
+    setup do
+      Elasticsearch::Rails::Lograge::Railtie.run_initializers
+    end
+
+    should "customize the Lograge configuration" do
+      assert_not_nil Elasticsearch::Rails::Lograge::Railtie.initializers
+                       .select { |i| i.name == 'elasticsearch.lograge' }
+                       .first
+    end
+  end
+end