lib/new_relic/agent/instrumentation/mongo.rb in newrelic_rpm-3.16.0.318 vs lib/new_relic/agent/instrumentation/mongo.rb in newrelic_rpm-3.16.1.320

@@ -63,27 +63,34 @@
         end
       rescue => e
         NewRelic::Agent.logger.debug("Exception during Mongo statement gathering", e)
       end
 
-      def new_relic_generate_metrics(operation, payload = nil)
-        payload ||= { :collection => self.name, :database => self.db.name }
-        NewRelic::Agent::Datastores::Mongo::MetricTranslator.metrics_for(operation, payload)
+      def new_relic_default_payload
+        { :collection => self.name, :database => self.db.name }
       end
 
+      def new_relic_start_segment name, payload
+        product = NewRelic::Agent::Datastores::Mongo::MetricTranslator::MONGO_PRODUCT_NAME
+        op_and_col = NewRelic::Agent::Datastores::Mongo::MetricTranslator.operation_and_collection_for name, payload
+        if op_and_col
+          NewRelic::Agent::Transaction.start_datastore_segment product, *op_and_col
+        end
+      end
+
       def instrument_with_new_relic_trace(name, payload = {}, &block)
-        metrics = new_relic_generate_metrics(name, payload)
+        segment = new_relic_start_segment name, payload
 
-        trace_execution_scoped(metrics) do
-          t0 = Time.now
-
+        begin
           result = NewRelic::Agent.disable_all_tracing do
             instrument_without_new_relic_trace(name, payload, &block)
           end
 
-          new_relic_notice_statement(t0, payload, name)
+          new_relic_notice_statement(segment.start_time, payload, name) if segment
           result
+        ensure
+          segment.finish if segment
         end
       end
 
       alias_method :instrument_without_new_relic_trace, :instrument
       alias_method :instrument, :instrument_with_new_relic_trace
@@ -91,20 +98,21 @@
   end
 
   def instrument_save
     ::Mongo::Collection.class_eval do
       def save_with_new_relic_trace(doc, opts = {}, &block)
-        metrics = new_relic_generate_metrics(:save)
-        trace_execution_scoped(metrics) do
-          t0 = Time.now
+        segment = new_relic_start_segment :save, new_relic_default_payload
 
+        begin
           result = NewRelic::Agent.disable_all_tracing do
             save_without_new_relic_trace(doc, opts, &block)
           end
 
-          new_relic_notice_statement(t0, doc, :save)
+          new_relic_notice_statement(segment.start_time, doc, :save) if segment
           result
+        ensure
+          segment.finish if segment
         end
       end
 
       alias_method :save_without_new_relic_trace, :save
       alias_method :save, :save_with_new_relic_trace
@@ -112,14 +120,13 @@
   end
 
   def instrument_ensure_index
     ::Mongo::Collection.class_eval do
       def ensure_index_with_new_relic_trace(spec, opts = {}, &block)
-        metrics = new_relic_generate_metrics(:ensureIndex)
-        trace_execution_scoped(metrics) do
-          t0 = Time.now
+        segment = new_relic_start_segment :ensureIndex, new_relic_default_payload
 
+        begin
           result = NewRelic::Agent.disable_all_tracing do
             ensure_index_without_new_relic_trace(spec, opts, &block)
           end
 
           spec = case spec
@@ -129,11 +136,13 @@
                    { spec => 1 }
                  else
                    spec.dup
                  end
 
-          new_relic_notice_statement(t0, spec, :ensureIndex)
+          new_relic_notice_statement(segment.start_time, spec, :ensureIndex) if segment
           result
+        ensure
+          segment.finish if segment
         end
       end
 
       alias_method :ensure_index_without_new_relic_trace, :ensure_index
       alias_method :ensure_index, :ensure_index_with_new_relic_trace