Reworked APM to allow custom operationId generator as well as timestamp generator

Author: christkv

index.js

@@ -37,39 +37,10 @@ connect.Timestamp = core.BSON.Timestamp;
 // Add connect method
 connect.connect = connect;
 
-// Instrumentation instance
-var instrumentation = null;
-
 // Set up the instrumentation method
-connect.instrument = function(options) {
-  if(!instrumentation) instrumentation = new Instrumentation(core, options)
-  return instrumentation;
-}
-// Get prototype
-var AggregationCursor = require('./lib/aggregation_cursor'),
-  CommandCursor = require('./lib/command_cursor'),
-  OrderedBulkOperation = require('./lib/bulk/ordered').OrderedBulkOperation,
-  UnorderedBulkOperation = require('./lib/bulk/unordered').UnorderedBulkOperation,
-  Admin = require('./lib/admin');
-
-// Instrument Hook
-connect.instrument = function(callback) {
-  var instrumentations = []
-
-  // Classes to support
-  var classes = [connect.GridStore, connect.Server, connect.ReplSet, connect.Mongos,
-    OrderedBulkOperation, UnorderedBulkOperation, CommandCursor, AggregationCursor,
-    connect.Cursor, connect.Collection, connect.Db];
-
-  // Add instrumentations to the available list
-  for(var i = 0; i < classes.length; i++) {
-    if(classes[i].define) {
-      instrumentations.push(classes[i].define.generate());
-    }
-  }
-
-  // Return the list of instrumentation points
-  callback(null, instrumentations);
+connect.instrument = function(options, callback) {
+  if(typeof options == 'function') callback = options, options = {};
+  return new Instrumentation(core, options, callback);
 }
 
 // Set our exports to be the connect function

lib/apm.js

@@ -1,6 +1,20 @@
 var EventEmitter = require('events').EventEmitter,
   inherits = require('util').inherits;
 
+// Get prototypes
+var AggregationCursor = require('./aggregation_cursor'),
+  CommandCursor = require('./command_cursor'),
+  OrderedBulkOperation = require('./bulk/ordered').OrderedBulkOperation,
+  UnorderedBulkOperation = require('./bulk/unordered').UnorderedBulkOperation,
+  GridStore = require('./gridfs/grid_store'),
+  Server = require('./server'),
+  ReplSet = require('./replset'),
+  Mongos = require('./mongos'),
+  Cursor = require('./cursor'),
+  Collection = require('./collection'),
+  Db = require('./db'),
+  Admin = require('./admin');
+
 var basicOperationIdGenerator = {
   operationId: 1,
 
@@ -9,12 +23,58 @@ var basicOperationIdGenerator = {
   }
 }
 
-var Instrumentation = function(core, options) {
+var basicTimestampGenerator = {
+  current: function() {
+    return new Date().getTime();
+  },
+  duration: function(start, end) {
+    return end - start;
+  }
+}
+
+var Instrumentation = function(core, options, callback) {
   options = options || {};
+
+  // Optional id generators
   var operationIdGenerator = options.operationIdGenerator || basicOperationIdGenerator;
+  // Optional timestamp generator
+  var timestampGenerator = options.timestampGenerator || basicTimestampGenerator;
   // Extend with event emitter functionality
   EventEmitter.call(this);
 
+  // Contains all the instrumentation overloads
+  this.overloads = [];
+
+  // ---------------------------------------------------------
+  //
+  // Instrument prototype
+  //
+  // ---------------------------------------------------------
+
+  var instrumentPrototype = function(callback) {
+    var instrumentations = []
+
+    // Classes to support
+    var classes = [GridStore, Server, ReplSet, Mongos,
+      OrderedBulkOperation, UnorderedBulkOperation, CommandCursor, AggregationCursor,
+      Cursor, Collection, Db];
+
+    // Add instrumentations to the available list
+    for(var i = 0; i < classes.length; i++) {
+      if(classes[i].define) {
+        instrumentations.push(classes[i].define.generate());
+      }
+    }
+
+    // Return the list of instrumentation points
+    callback(null, instrumentations);
+  }
+
+  // Did the user want to instrument the prototype
+  if(typeof callback == 'function') {
+    instrumentPrototype(callback);
+  }
+
   // ---------------------------------------------------------
   //
   // Server
@@ -31,6 +91,9 @@ var Instrumentation = function(core, options) {
   methods.forEach(function(x) {
     var func = proto[x];
 
+    // Add to overloaded methods
+    self.overloads.push({proto: proto, name:x, func:func});
+
     // The actual prototype
     proto[x] = function() {
       var requestId = core.Query.nextRequestId();
@@ -50,6 +113,7 @@ var Instrumentation = function(core, options) {
 
       // Get a connection reference for this server instance
       var connection = this.s.pool.get()
+      
       // Emit the start event for the command
       var command = {
         // Returns the command.
@@ -73,13 +137,13 @@ var Instrumentation = function(core, options) {
       self.emit('started', command)
 
       // Start time
-      var startTime = new Date().getTime();
+      var startTime = timestampGenerator.current();
 
       // Push our handler callback
       args.push(function(err, r) {
-        var endTime = new Date().getTime();
+        var endTime = timestampGenerator.current();
         var command = {
-          duration: (endTime - startTime),
+          duration: timestampGenerator.duration(startTime, endTime),
           commandName: commandName,
           requestId: requestId,
           operationId: ourOpId,
@@ -122,6 +186,9 @@ var Instrumentation = function(core, options) {
     methods.forEach(function(x) {
       var func = proto[x];
 
+      // Add to overloaded methods
+      self.overloads.push({proto: proto, name:x, func:func});
+
       // The actual prototype
       proto[x] = function() {
         var bulk = this;
@@ -168,10 +235,14 @@ var Instrumentation = function(core, options) {
   }
 
   prototypes.forEach(function(proto) {
+
     // Core server method we are going to wrap
     methods.forEach(function(x) {
       var func = proto[x];
 
+      // Add to overloaded methods
+      self.overloads.push({proto: proto, name:x, func:func});
+
       // The actual prototype
       proto[x] = function() {
         var cursor = this;
@@ -247,18 +318,12 @@ var Instrumentation = function(core, options) {
           }
         }
 
-        // console.log("#############################################################")
-        // console.dir(this)
-
         // Set up the connection
         var connectionId = null;
         // Set local connection
         if(this.connection) connectionId = this.connection;
         if(!connectionId && this.server && this.server.getConnection) connectionId = this.server.getConnection();
 
-        // var connections = this.connections();
-        // var connectionId = connections.length > 0 ? connections[0] :  null;
-
         // Emit the start event for the command
         var command = {
           // Returns the command.
@@ -286,15 +351,15 @@ var Instrumentation = function(core, options) {
 
         // We do not have a callback but a Promise
         if(typeof callback == 'function') {
-          var startTime = new Date();
+          var startTime = timestampGenerator.current();
           // Emit the started event
           self.emit('started', command)
           // Add our callback handler
           args.push(function(err, r) {
             if(err) {
               // Command
               var command = {
-                duration: (new Date().getTime() - startTime.getTime()),
+                duration: timestampGenerator.duration(startTime, timestampGenerator.current()),
                 commandName: commandTranslation[x],
                 requestId: requestId,
                 operationId: ourOpId,
@@ -306,7 +371,7 @@ var Instrumentation = function(core, options) {
             } else {
               // cursor id is zero, we can issue success command
               var command = {
-                duration: (new Date().getTime() - startTime.getTime()),
+                duration: timestampGenerator.duration(startTime, timestampGenerator.current()),
                 commandName: commandTranslation[x],
                 requestId: requestId,
                 operationId: cursor.operationId,
@@ -331,14 +396,14 @@ var Instrumentation = function(core, options) {
           var promise = func.apply(this, args);
           // Return a new promise
           return new cursor.s.promiseLibrary(function(resolve, reject) {
-            var startTime = new Date();
+            var startTime = timestampGenerator.current();
             // Emit the started event
             self.emit('started', command)
             // Execute the function
             promise.then(function(r) {
               // cursor id is zero, we can issue success command
               var command = {
-                duration: (new Date().getTime() - startTime.getTime()),
+                duration: timestampGenerator.duration(startTime, timestampGenerator.current()),
                 commandName: commandTranslation[x],
                 requestId: requestId,
                 operationId: cursor.operationId,
@@ -351,7 +416,7 @@ var Instrumentation = function(core, options) {
             }).catch(function(err) {
               // Command
               var command = {
-                duration: (new Date().getTime() - startTime.getTime()),
+                duration: timestampGenerator.duration(startTime, timestampGenerator.current()),
                 commandName: commandTranslation[x],
                 requestId: requestId,
                 operationId: ourOpId,
@@ -372,4 +437,11 @@ var Instrumentation = function(core, options) {
 
 inherits(Instrumentation, EventEmitter);
 
-module.exports = Instrumentation;
+Instrumentation.prototype.uninstrument = function() {
+  for(var i = 0; i < this.overloads.length; i++) {
+    var obj = this.overloads[i];
+    obj.proto[obj.name] = obj.func;
+  }
+}
+
+module.exports = Instrumentation;