Merge pull request newrelic#502 from ruby-agent/RUBY-1294-mongo-again

RUBY-1294 Fixes for usage with mongod 2.6+

lib/new_relic/agent/datastores/mongo/metric_translator.rb

@@ -12,19 +12,19 @@ module MetricTranslator
           def self.metrics_for(name, payload, request_type = :web)
             payload ||= {}
 
-            # The 1.10.0 version of the mongo driver renamed 'remove' to
-            # 'delete', but for metric consistency with previous versions we
-            # want to keep it as 'remove'.
-            name = 'remove' if name.to_s == 'delete'
-
-            collection = payload[:collection]
-
-            if collection_in_selector?(collection, payload)
+            if collection_in_selector?(payload)
               command_key = command_key_from_selector(payload)
               name        = get_name_from_selector(command_key, payload)
               collection  = get_collection_from_selector(command_key, payload)
+            else
+              collection = payload[:collection]
             end
 
+            # The 1.10.0 version of the mongo driver renamed 'remove' to
+            # 'delete', but for metric consistency with previous versions we
+            # want to keep it as 'remove'.
+            name = 'remove' if name.to_s == 'delete'
+
             if self.find_one?(name, payload)
               name = 'findOne'
             elsif self.find_and_remove?(name, payload)
@@ -74,8 +74,8 @@ def self.instance_metric(host, port, database)
             "Datastore/instance/MongoDB/#{host}:#{port}/#{database}"
           end
 
-          def self.collection_in_selector?(collection, payload)
-            collection == '$cmd' && payload[:selector]
+          def self.collection_in_selector?(payload)
+            payload[:collection] == '$cmd' && payload[:selector]
           end
 
           NAMES_IN_SELECTOR = [

lib/new_relic/agent/datastores/mongo/statement_formatter.rb

@@ -24,10 +24,11 @@ module StatementFormatter
             :selector
           ]
 
-          def self.format(statement)
+          def self.format(statement, operation)
             return nil unless NewRelic::Agent.config[:'mongo.capture_queries']
 
-            result = {}
+            result = { :operation => operation }
+
             PLAINTEXT_KEYS.each do |key|
               result[key] = statement[key] if statement.key?(key)
             end

lib/new_relic/agent/instrumentation/mongo.rb

@@ -28,17 +28,24 @@
   end
 
   def install_mongo_instrumentation
-    setup_logging_for_instrumentation
-    instrument_mongo_logging
+    require 'new_relic/agent/datastores/mongo/metric_generator'
+    require 'new_relic/agent/datastores/mongo/statement_formatter'
+
+    instrument_via_mongo_logging
     instrument_save
     instrument_ensure_index
   end
 
-  def setup_logging_for_instrumentation
-    ::Mongo::Logging.class_eval do
+  def instrument_via_mongo_logging
+    hook_instrument_method(::Mongo::Collection)
+    hook_instrument_method(::Mongo::Connection)
+    hook_instrument_method(::Mongo::Cursor)
+    hook_instrument_method(::Mongo::CollectionWriter) if defined?(::Mongo::CollectionWriter)
+  end
+
+  def hook_instrument_method(target_class)
+    target_class.class_eval do
       include NewRelic::Agent::MethodTracer
-      require 'new_relic/agent/datastores/mongo/metric_generator'
-      require 'new_relic/agent/datastores/mongo/statement_formatter'
 
       def new_relic_instance_metric_builder
         Proc.new do
@@ -55,9 +62,8 @@ def new_relic_instance_metric_builder
 
       # It's key that this method eats all exceptions, as it rests between the
       # Mongo operation the user called and us returning them the data. Be safe!
-      def new_relic_notice_statement(t0, payload, operation)
-        payload[:operation] = operation
-        statement = NewRelic::Agent::Datastores::Mongo::StatementFormatter.format(payload)
+      def new_relic_notice_statement(t0, payload, name)
+        statement = NewRelic::Agent::Datastores::Mongo::StatementFormatter.format(payload, name)
         if statement
           NewRelic::Agent.instance.transaction_sampler.notice_nosql_statement(statement, (Time.now - t0).to_f)
         end
@@ -67,27 +73,19 @@ def new_relic_notice_statement(t0, payload, operation)
 
       def new_relic_generate_metrics(operation, payload = nil)
         payload ||= { :collection => self.name, :database => self.db.name }
-        metrics = NewRelic::Agent::Datastores::Mongo::MetricGenerator.generate_metrics_for(operation, payload)
+        NewRelic::Agent::Datastores::Mongo::MetricGenerator.generate_metrics_for(operation, payload)
       end
-    end
 
-    ::Mongo::Collection.class_eval { include ::Mongo::Logging }
-    ::Mongo::Connection.class_eval { include ::Mongo::Logging }
-    ::Mongo::Cursor.class_eval     { include ::Mongo::Logging }
-
-    if defined?(::Mongo::CollectionOperationWriter)
-      ::Mongo::CollectionOperationWriter.class_eval { include ::Mongo::Logging }
-    end
-  end
-
-  def instrument_mongo_logging
-    ::Mongo::Logging.class_eval do
       def instrument_with_new_relic_trace(name, payload = {}, &block)
         metrics = new_relic_generate_metrics(name, payload)
 
         trace_execution_scoped(metrics, :additional_metrics_callback => new_relic_instance_metric_builder) do
           t0 = Time.now
-          result = instrument_without_new_relic_trace(name, payload, &block)
+
+          result = NewRelic::Agent.disable_all_tracing do
+            instrument_without_new_relic_trace(name, payload, &block)
+          end
+
           new_relic_notice_statement(t0, payload, name)
           result
         end
@@ -105,13 +103,8 @@ def save_with_new_relic_trace(doc, opts = {}, &block)
         trace_execution_scoped(metrics, :additional_metrics_callback => new_relic_instance_metric_builder) do
           t0 = Time.now
 
-          transaction_state = NewRelic::Agent::TransactionState.get
-          transaction_state.push_traced(false)
-
-          begin
-            result = save_without_new_relic_trace(doc, opts, &block)
-          ensure
-            transaction_state.pop_traced
+          result = NewRelic::Agent.disable_all_tracing do
+            save_without_new_relic_trace(doc, opts, &block)
           end
 
           new_relic_notice_statement(t0, doc, :save)
@@ -131,13 +124,8 @@ def ensure_index_with_new_relic_trace(spec, opts = {}, &block)
         trace_execution_scoped(metrics, :additional_metrics_callback => new_relic_instance_metric_builder) do
           t0 = Time.now
 
-          transaction_state = NewRelic::Agent::TransactionState.get
-          transaction_state.push_traced(false)
-
-          begin
-            result = ensure_index_without_new_relic_trace(spec, opts, &block)
-          ensure
-            transaction_state.pop_traced
+          result = NewRelic::Agent.disable_all_tracing do
+            ensure_index_without_new_relic_trace(spec, opts, &block)
           end
 
           spec = case spec

test/multiverse/suites/mongo/helpers/mongo_operation_tests.rb

@@ -134,8 +134,6 @@ def test_records_metrics_for_find_and_remove
   def test_records_metrics_for_create_index
     @collection.create_index([[unique_field_name, Mongo::ASCENDING]])
 
-    metrics = metrics_with_attributes(build_test_metrics(:createIndex))
-
     # The createIndexes command was added to the mongo server in version 2.6.
     # As of version 1.10.0 of the Ruby driver, the driver will attempt to
     # service a create_index call by first issuing a createIndexes command to
@@ -144,12 +142,20 @@ def test_records_metrics_for_create_index
     #
     # So, if we're running with version 1.10.0 or later of the driver, we expect
     # some additional metrics to be recorded.
-    if NewRelic::Agent::Datastores::Mongo.is_version_1_10_or_later?
-      metrics["Datastore/statement/MongoDB/#{@collection_name}/createIndexes"] = { :call_count => 1 }
-      metrics['Datastore/operation/MongoDB/createIndexes']                     = { :call_count => 1 }
+    client_is_1_10_or_later = NewRelic::Agent::Datastores::Mongo.is_version_1_10_or_later?
+
+    create_index_metrics   = metrics_with_attributes(build_test_metrics(:createIndex))
+    create_indexes_metrics = metrics_with_attributes(build_test_metrics(:createIndexes))
+
+    if !client_is_1_10_or_later
+      metrics = create_index_metrics
+    elsif client_is_1_10_or_later && !server_is_2_6_or_later?
+      metrics = create_index_metrics.merge(create_indexes_metrics)
       metrics['ActiveRecord/all'][:call_count] += 1
       metrics['Datastore/allWeb'][:call_count] += 1
       metrics['Datastore/all'][:call_count]    += 1
+    elsif client_is_1_10_or_later && server_is_2_6_or_later?
+      metrics = create_indexes_metrics
     end
 
     assert_metrics_recorded(metrics)
@@ -284,16 +290,18 @@ def test_notices_nosql
 
     in_transaction do
       @collection.insert(@tribble)
+
       segment = find_last_transaction_segment
     end
 
-    expected = { :database   => @database_name,
-                 :collection => @collection_name,
-                 :operation  => :insert}
+    expected = {
+      :database   => @database_name,
+      :collection => @collection_name,
+      :operation  => :insert
+    }
 
     result = segment.params[:statement]
-
-    assert_equal expected, result, "Expected result (#{result}) to be #{expected}"
+    assert_equal expected, result
   end
 
   def test_noticed_nosql_includes_operation
@@ -304,12 +312,25 @@ def test_noticed_nosql_includes_operation
       segment = find_last_transaction_segment
     end
 
-    expected = :insert
+    query = segment.params[:statement]
+
+    assert_equal :insert, query[:operation]
+  end
+
+  def test_noticed_nosql_includes_update_operation
+    segment = nil
+
+    in_transaction do
+      updated = @tribble.dup
+      updated['name'] = 't-rex'
+      @collection.update(@tribble, updated)
+
+      segment = find_last_transaction_segment
+    end
 
     query = segment.params[:statement]
-    result = query[:operation]
 
-    assert_equal expected, result
+    assert_equal :update, query[:operation]
   end
 
   def test_noticed_nosql_includes_save_operation
@@ -320,12 +341,8 @@ def test_noticed_nosql_includes_save_operation
       segment = find_last_transaction_segment
     end
 
-    expected = :save
-
     query = segment.params[:statement]
-    result = query[:operation]
-
-    assert_equal expected, result
+    assert_equal :save, query[:operation]
   end
 
   def test_noticed_nosql_includes_ensure_index_operation
@@ -392,6 +409,8 @@ def test_noticed_nosql_does_not_contain_selector_values
 
     statement = segment.params[:statement]
 
+    refute statement.inspect.include?('$secret')
+
     assert_equal '?', statement[:selector]['password']
   end
 
@@ -453,4 +472,9 @@ def ensure_collection_exists
     NewRelic::Agent.drop_buffered_data
   end
 
+  def server_is_2_6_or_later?
+    client = @collection.db.respond_to?(:client) && @collection.db.client
+    return false unless client
+    client.respond_to?(:max_wire_version) && client.max_wire_version >= 2
+  end
 end

test/multiverse/suites/mongo/helpers/mongo_server.rb

@@ -115,18 +115,20 @@ def wait_until(seconds = 10)
   def startup_command
     pid_file = "--pidfilepath #{pid_path}"
     log_file = "--logpath #{log_path} "
+    fork     = "--fork"
 
     dbpath = "--dbpath #{db_path}"
     port_flag = "--port #{self.port}"
     small_mongo = "--oplogSize 128 --smallfiles"
-    repl_set = "--fork --replSet multiverse"
+    repl_set = "--replSet multiverse"
 
-    base = "#{port_flag} #{pid_file} #{log_file} #{small_mongo} #{dbpath}"
+    base = "#{port_flag} #{fork} #{pid_file} #{log_file} #{small_mongo} #{dbpath}"
 
+    mongod_path = ENV['MONGOD_PATH'] || 'mongod'
     if self.type == :single
-      "mongod #{base} &"
+      "#{mongod_path} #{base}"
     elsif self.type == :replica
-      "mongod #{repl_set} #{base} &"
+      "#{mongod_path} #{repl_set} #{base}"
     end
   end
 

test/new_relic/agent/datastores/mongo/statement_formatter_test.rb

@@ -30,37 +30,38 @@ class StatementFormatterTest < Minitest::Test
                                                   :_id       => "BSON::ObjectId('?')" } }.freeze
 
           def test_doesnt_modify_incoming_statement
-            formatted = StatementFormatter.format(DOC_STATEMENT)
+            formatted = StatementFormatter.format(DOC_STATEMENT, :find)
             refute_same DOC_STATEMENT, formatted
           end
 
           def test_statement_formatter_removes_unwhitelisted_keys
-            formatted = StatementFormatter.format(DOC_STATEMENT)
+            formatted = StatementFormatter.format(DOC_STATEMENT, :find)
             assert_equal_unordered(formatted.keys, StatementFormatter::PLAINTEXT_KEYS)
           end
 
           def test_can_disable_statement_capturing_queries
             with_config(:'mongo.capture_queries' => false) do
-              formatted = StatementFormatter.format(DOC_STATEMENT)
+              formatted = StatementFormatter.format(DOC_STATEMENT, :find)
               assert_nil formatted
             end
           end
 
           def test_statement_formatter_obfuscates_by_default
             expected = { :database   => 'multiverse',
                          :collection => 'tribbles',
+                         :operation  => :find,
                          :selector   => { 'name'     => '?',
                                           :operation => :find,
                                           :_id       => '?' } }
 
-            result = StatementFormatter.format(SELECTOR_STATEMENT)
+            result = StatementFormatter.format(SELECTOR_STATEMENT, :find)
             assert_equal expected, result
           end
 
           def test_statement_formatter_raw_selectors
             with_config(:'mongo.obfuscate_queries' => false) do
-              result = StatementFormatter.format(SELECTOR_STATEMENT)
-              assert_equal SELECTOR_STATEMENT, result
+              result = StatementFormatter.format(SELECTOR_STATEMENT, :find)
+              assert_equal SELECTOR_STATEMENT.merge(:operation => :find), result
             end
           end