minor fixes for promises when using APM

Author: christkv

lib/apm.js

@@ -132,13 +132,18 @@ var Instrumentation = function(core, options) {
 
         // Get the callback
         var callback = args.pop();
-        args.push(function(err, r) {
-          // Return to caller
-          callback(err, r);
-        });
-
-        // Apply the call
-        func.apply(this, args);
+        // If we have a callback use this
+        if(typeof callback == 'function') {
+          args.push(function(err, r) {
+            // Return to caller
+            callback(err, r);
+          });
+
+          // Apply the call
+          func.apply(this, args);
+        } else {
+          return func.apply(this, args);
+        }
       }
     });
   });
@@ -174,14 +179,14 @@ var Instrumentation = function(core, options) {
         var ourOpId = operationIdGenerator.next();
         var parts = this.ns.split('.');
         var db = parts[0];
+
         // Get the collection
         parts.shift();
         var collection = parts.join('.');
 
         // If we have a find method, set the operationId on the cursor
         if(x == '_find') {
-          this.operationId = ourOpId;
-          this.startTime = new Date();
+          cursor.operationId = ourOpId;
         }
 
         // Set the command
@@ -242,6 +247,18 @@ 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.
@@ -258,50 +275,96 @@ var Instrumentation = function(core, options) {
           // Returns the connection id for the command. For languages that do not have this,
           // this MUST return the driver equivalent which MUST include the server address and port.
           // The name of this field is flexible to match the object that is returned from the driver.
-          connectionId: this.server.getConnection()
+          connectionId: connectionId
         };
 
-        // Emit the started event
-        self.emit('started', command)
-
         // Get the aruments
         var args = Array.prototype.slice.call(arguments, 0);
+
         // Get the callback
         var callback = args.pop();
-        args.push(function(err, r) {
-          if(err) {
-            // Command
-            var command = {
-              duration: (new Date().getTime() - cursor.startTime.getTime()),
-              commandName: commandTranslation[x],
-              requestId: requestId,
-              operationId: ourOpId,
-              connectionId: cursor.server.getConnection(),
-              failure: err };
-
-            // Emit the command
-            self.emit('failed', command)
-          } else {
-            // cursor id is zero, we can issue success command
-            var command = {
-              duration: (new Date().getTime() - cursor.startTime.getTime()),
-              commandName: commandTranslation[x],
-              requestId: requestId,
-              operationId: cursor.operationId,
-              connectionId: cursor.server.getConnection(),
-              reply: cursor.cursorState.documents
-            };
-
-            // Emit the command
-            self.emit('succeeded', command)
-          }
 
-          // Return to caller
-          callback(err, r);
-        });
-
-        // Apply the call
-        func.apply(this, args);
+        // We do not have a callback but a Promise
+        if(typeof callback == 'function') {
+          var startTime = new Date();
+          // 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()),
+                commandName: commandTranslation[x],
+                requestId: requestId,
+                operationId: ourOpId,
+                connectionId: cursor.server.getConnection(),
+                failure: err };
+
+              // Emit the command
+              self.emit('failed', command)
+            } else {
+              // cursor id is zero, we can issue success command
+              var command = {
+                duration: (new Date().getTime() - startTime.getTime()),
+                commandName: commandTranslation[x],
+                requestId: requestId,
+                operationId: cursor.operationId,
+                connectionId: cursor.server.getConnection(),
+                reply: cursor.cursorState.documents
+              };
+
+              // Emit the command
+              self.emit('succeeded', command)
+            }
+
+            // Return to caller
+            callback(err, r);
+          });
+
+          // Apply the call
+          func.apply(this, args);
+        } else {
+          // Assume promise, push back the missing value
+          args.push(callback);
+          // Get the promise
+          var promise = func.apply(this, args);
+          // Return a new promise
+          return new cursor.s.promiseLibrary(function(resolve, reject) {
+            var startTime = new Date();
+            // 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()),
+                commandName: commandTranslation[x],
+                requestId: requestId,
+                operationId: cursor.operationId,
+                connectionId: cursor.server.getConnection(),
+                reply: cursor.cursorState.documents
+              };
+
+              // Emit the command
+              self.emit('succeeded', command)
+            }).catch(function(err) {
+              // Command
+              var command = {
+                duration: (new Date().getTime() - startTime.getTime()),
+                commandName: commandTranslation[x],
+                requestId: requestId,
+                operationId: ourOpId,
+                connectionId: cursor.server.getConnection(),
+                failure: err };
+
+              // Emit the command
+              self.emit('failed', command)
+              // reject the promise
+              reject(err);
+            });
+          });
+        }
       }
     });
   });

test/runner.js

@@ -50,6 +50,12 @@ var startupOptions = {
   , skip: false
 }
 
+// var listener = require('../').instrument();
+// listener.on('started', function(event) {
+//   // console.log("-------------------------------------------- started")
+//   // console.log(JSON.stringify(event, null, 2))
+// });
+
 /**
  * Standalone MongoDB Configuration
  */
@@ -211,7 +217,7 @@ var createConfiguration = function(options) {
 
 // Set up the runner
 var runner = new Runner({
-  logLevel:'error',
+  logLevel:'info',
   runners: 1,
   failFast: true
 });