mongo2.rb in oneapm_rpm-1.3.0.rc1

- old
+ new
@@ -1,87 +1,116 @@
 # encoding: utf-8
 
-LibraryDetection.defer do
-  named :mongo2
+module OneApm
+  module Agent
+    module Instrumentation
+      class MongoCommandSubscriber
 
-  depends_on do
-    require 'one_apm/agent/datastore/mongo'
-    defined?(::Mongo) &&  OneApm::Agent::Datastore::Mongo.is_version2? && !OneApm::Manager.config[:disable_mongo]
-  end
+        MONGODB = 'MongoDB'.freeze
+        STARTED = 'STARTED'.freeze
+        SUCCEEDED = 'SUCCEEDED'.freeze
+        FAILED = 'FAILED'.freeze
 
-  executes do
-    OneApm::Manager.logger.info 'Installing Mongo 2 instrumentation'
-    install_mongo_instrumentation
-  end
+        def started(event)
+          begin
+            return unless OneApm::Manager.tl_is_execution_traced?
+            operations[event.operation_id] = event
+          rescue Exception => e
+            log_operations_error(STARTED, e)
+          end
+        end
 
-  def install_mongo_instrumentation
-    require 'one_apm/agent/datastore/mongo/metric_translator'
-    require 'one_apm/agent/datastore/mongo/statement_formatter'
-    hook_instrument_method_for_collection
-    hook_instrument_method_for_view
-  end
+        def succeeded(event)
+          operator SUCCEEDED, event
+        end
 
-  def hook_instrument_method_for_collection
-    methods = [:create,:drop, :find, :indexes, :insert_one, :insert_many, :bulk_write, :parallel_scan]
-    instrument_methods_with_oneapm(::Mongo::Collection, methods)
-  end
+        def failed(event)
+          operator FAILED, event
+        end
 
+        private
 
-  def hook_instrument_method_for_view
-    methods = [:count, :distinct, :map_reduce,:find_one_and_delete, :find_one_and_replace, :find_one_and_update, 
-      :delete_many, :delete_one, :replace_one, :update_many, :update_one]
-    instrument_methods_with_oneapm(::Mongo::Collection::View, methods)
-  end
+        def operations
+          @operations ||= {}
+        end
 
+        def mertircs(event)
+          OneApm::Agent::Datastore::MetricHelper.metrics_for(
+            MONGODB, event.command_name, event.command.values.first
+          )
+        end
 
-  def instrument_methods_with_oneapm klass, methods
-    klass.class_eval do 
-      include OneApm::Support::MethodTracer
+        def operator(operator_type, event)
+          begin
+            state = OneApm::TransactionState.tl_get
+            return unless state.is_execution_traced?
+            stared_event = operations.delete(event.operation_id)
+            one_apm_generate_metrics(stared_event, event.duration)
+            one_apm_notice_statement(stared_event, operator_type, event.duration)
+            one_apm_notice_sql(state, stared_event, operator_type, event.duration)
+          rescue => e
+            log_operations_error(operator_type, e)
+          end
+        end
 
-      def payload_for_oneapm
-        collection = self.name rescue self.collection.name
-        @payload_for_oneapm ||= { :collection => collection, :database => self.database.name }
-      end
-
-      def one_apm_notice_statement(t0, name)
-        statement = OneApm::Agent::Datastore::Mongo::StatementFormatter.format(payload_for_oneapm, name)
-        if statement
-          OneApm::Manager.agent.transaction_sampler.notice_nosql_statement(statement, (Time.now - t0).to_f)
+        def one_apm_generate_metrics(event, duration)         
+          base, *other_metrics = mertircs(event)
+          OneApm::Manager.agent.stats_engine.tl_record_scoped_and_unscoped_metrics(
+            base, other_metrics, duration
+          )
         end
-      rescue => e
-        OneApm::Manager.logger.debug("Exception during Mongo statement gathering", e)
-      end
 
-      def one_apm_generate_metrics(operation)
-        OneApm::Agent::Datastore::Mongo::MetricTranslator.metrics_for(operation, payload_for_oneapm)
-      end
-    end
+        def one_apm_notice_sql(state, event, status, duration)
+          stack = state.traced_method_stack
+          base, *other_metrics = mertircs(event)
 
-    methods.each do |method_name|
-      next unless klass.method_defined?(method_name)
+          started_time = Time.now.to_f
+          frame = stack.push_frame(state, :mongo_tracer, started_time)
 
-      klass.class_eval do 
+          builder = state.transaction_sample_builder
+          format_sql = OneApm::Agent::Datastore::Mongo::CommandFormatter.format_sql(event, status)
+          OneApm::Manager.agent.transaction_sampler.send(:notice_extra_data, builder, format_sql, duration, :sql)
 
-        method_name_without_oneapm = :"#{method_name}_without_oneapm"
-        method_name_with_oneapm = :"#{method_name}_with_oneapm"
+          stack.pop_frame(state, frame, base, started_time + duration)
+        end
 
-        define_method method_name_with_oneapm do |*args, &block|
-          metrics = one_apm_generate_metrics(method_name)
-
-          trace_execution_scoped(metrics) do
-            t0 = Time.now
-            result = OneApm::Manager.disable_all_tracing do
-              send method_name_without_oneapm, *args, &block
-            end
-            one_apm_notice_statement(t0, method_name)
-            result
+        def one_apm_notice_statement(event, status, duration)
+          statement = OneApm::Agent::Datastore::Mongo::CommandFormatter.format(event, status)
+          if statement
+             OneApm::Manager.agent.transaction_sampler.notice_nosql_statement(statement, duration)
           end
-         
+          rescue => e
+            OneApm::Manager.logger.debug("Exception during Mongo statement gathering", e)
         end
-        alias_method method_name_without_oneapm, method_name.to_sym
-        alias_method method_name.to_sym, method_name_with_oneapm
+
+        def log_operations_error(event_type, error)
+          OneApm::Manager.logger.error("Error during MongoDB #{event_type} event:")
+          OneApm::Manager.logger.log_exception(:error, error)
+        end
+  
       end
     end
-
   end
- 
 end
+
+LibraryDetection.defer do
+  named :mongo2
+
+  depends_on do
+    require 'one_apm/agent/datastore/mongo'
+    defined?(::Mongo) && OneApm::Agent::Datastore::Mongo.is_support_version2? && !OneApm::Manager.config[:disable_mongo]
+  end
+
+  executes do
+    OneApm::Manager.logger.info 'Installing Mongo 2 instrumentation'
+    install_mongo_command_subscriber
+  end
+
+  def install_mongo_command_subscriber
+    require 'one_apm/agent/datastore/metric_helper'
+    require 'one_apm/agent/datastore/mongo/command_formatter'
+    ::Mongo::Monitoring::Global.subscribe(
+      ::Mongo::Monitoring::COMMAND,
+      OneApm::Agent::Instrumentation::MongoCommandSubscriber.new
+    )
+  end
+end
\ No newline at end of file