fix(instrumentation): redis multi typeerror

Author: dszakallas

circle.yml

@@ -1,4 +1,9 @@
 machine:
+  services:
+    - redis
+    - mysql
+    - mongodb
+    - postgresql
   node:
     version: 6.9.5
   environment:

lib/instrumentations/trace-instrumentation-redis.js

@@ -3,56 +3,74 @@ var flatMap = require('lodash.flatmap')
 var Shimmer = require('../utils/shimmer')
 var consts = require('../consts')
 var utils = require('./utils')
-var instrumentedCommands = require('./utils').redisTools.instrumentedCommands
+var Module = require('module')
+var redisCommands = require('./utils').redisTools.instrumentedCommands
 
-module.exports = function wrap (redis, agent, pkg) {
-  var _instrumentedCommands = flatMap(Object.keys(instrumentedCommands), function (key) {
-    return instrumentedCommands[key]
-  })
+var instrumentedCommands = flatMap(Object.keys(redisCommands), function (key) {
+  return redisCommands[key]
+})
 
-  Shimmer.wrap(redis.RedisClient.prototype,
-    _instrumentedCommands.concat(['multi']), function (original, name) {
-      return function () {
-        var host = this.address
-        var args = Array.prototype.slice.apply(arguments)
+function wrapRedisClient (redis, agent, pkg) {
+  Shimmer.wrap(redis.RedisClient.prototype, instrumentedCommands, function (original, name) {
+    return function () {
+      var host = this.address
+      var args = Array.prototype.slice.apply(arguments)
+      return utils.wrapQuery.call(this, original, args, agent, {
+        continuationMethod: 'callback',
+        protocol: consts.PROTOCOLS.REDIS,
+        host: host || 'unknown',
+        method: name,
+        url: 'unknown'
+      })
+    }
+  })
+  return redis
+}
 
-        if (name === 'multi') {
-          // start a multi
-          var multi = original.apply(this, args)
-          multi.__trace = []
-          var originalExec = multi.exec
-          multi.exec = function () {
-            var args = Array.prototype.slice.apply(arguments)
-            var commands = this.__trace
-            return utils.wrapQuery.call(this, originalExec, args, agent, {
-              continuationMethod: 'callback',
-              protocol: consts.PROTOCOLS.REDIS,
-              host: host || 'unknown',
-              method: 'multi: ' + commands.join(', '),
-              url: 'unknown'
-            })
-          }
-          return multi
-        } else {
-          return utils.wrapQuery.call(this, original, args, agent, {
-            continuationMethod: 'callback',
-            protocol: consts.PROTOCOLS.REDIS,
-            host: host || 'unknown',
-            method: name,
-            url: 'unknown'
-          })
-        }
+function wrapMulti (Multi, agent, pkg) {
+  Shimmer.wrap(Multi.prototype, ['exec', 'EXEC', 'exec_transaction'], function (original, name) {
+    return function () {
+      var args = Array.prototype.slice.apply(arguments)
+      var host = this._client && this._client.address
+      var commands = this.__trace || []
+      var methodString
+      if (commands.length) {
+        methodString = 'multi: ' + commands.join(', ')
+      } else {
+        methodString = 'multi'
       }
-    })
+      return utils.wrapQuery.call(this, original, args, agent, {
+        continuationMethod: 'callback',
+        protocol: consts.PROTOCOLS.REDIS,
+        host: host || 'unknown',
+        method: methodString,
+        url: 'unknown'
+      })
+    }
+  })
 
-  Shimmer.wrap(redis.Multi.prototype, _instrumentedCommands,
+  Shimmer.wrap(Multi.prototype, instrumentedCommands,
     function (original, name) {
       return function () {
         var args = Array.prototype.slice.apply(arguments)
+        this.__trace = this.__trace || []
         this.__trace.push(name)
         return original.apply(this, args)
       }
     })
+  return Multi
+}
 
-  return redis
+module.exports = {
+  instrumentations: [{
+    path: 'redis',
+    pre: function () {
+      Module._load('redis/lib/multi', arguments[3])
+      return Array.prototype.slice.call(arguments, 2)
+    },
+    post: wrapRedisClient
+  }, {
+    path: 'redis/lib/multi',
+    post: wrapMulti
+  }]
 }