fix(redis): support new internal methods

Author: gergelyke

lib/instrumentations/core/http/util.js

@@ -1,3 +1,4 @@
+// eslint-disable-next-line no-useless-escape
 var DATA_URI_REGEX = /(data:([a-z]+\/[a-z0-9\-\+]+(;[a-z\-]+\=[a-z0-9\-]+)?)?(;base64)?,)([a-z0-9\!\$\&\'\,\(\)\*\+\,\;\=\-\.\_\~\:\@\/\?\%\s]*\s*)$/i
 
 function formatDataUrl (text) {

lib/instrumentations/redis.js

@@ -1,59 +1,57 @@
+var flatMap = require('lodash.flatmap')
+
 var Shimmer = require('../utils/shimmer')
 var consts = require('../consts')
+var utils = require('./utils')
+var instrumentedCommands = require('./utils').redisTools.instrumentedCommands
 
-module.exports = function wrap (redis, agent) {
-  Shimmer.wrap(redis.RedisClient.prototype, 'redis.RedisClient.prototype', 'send_command', function (original) {
-    return function () {
-      var spanId = agent.generateSpanId()
-      var clientSendTime = agent.getMicrotime()
-      var requestId = agent.getTransactionId()
-
-      var host = this.address
-
-      var args = Array.prototype.slice.apply(arguments)
-      var command = args[0]
-      var last = args[args.length - 1]
+module.exports = function wrap (redis, agent, pkg) {
+  var _instrumentedCommands = flatMap(Object.keys(instrumentedCommands), function (key) {
+    return instrumentedCommands[key]
+  })
 
-      var wrappedCallback = function (original) {
-        return function (err) {
-          agent.clientReceive({
+  Shimmer.wrap(redis.RedisClient.prototype, 'redis.RedisClient.prototype',
+    _instrumentedCommands.concat(['multi']), function (original, name) {
+      return function () {
+        var host = this.address
+        var args = Array.prototype.slice.apply(arguments)
+
+        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 host = this.address
+            var commands = this.__trace
+            return utils.wrapQuery.call(this, originalExec, args, agent, {
+              protocol: consts.PROTOCOLS.REDIS,
+              host: host,
+              method: 'multi: ' + commands.join(', '),
+              url: 'unknown'
+            })
+          }
+          return multi
+        } else {
+          return utils.wrapQuery.call(this, original, args, agent, {
             protocol: consts.PROTOCOLS.REDIS,
-            id: requestId,
-            spanId: spanId,
             host: host,
-            time: clientSendTime,
-            url: 'unknown', // TODO(c/KRU7H3D1): add support for url parameter with redis key
-            method: command,
-            mustCollect: err ? consts.MUST_COLLECT.ERROR : undefined,
-            responseTime: agent.getMicrotime() - clientSendTime,
-            status: err ? consts.EDGE_STATUS.NOT_OK : consts.EDGE_STATUS.OK,
-            statusCode: err ? err.code : 200
+            method: name,
+            url: 'unknown'
           })
-          return original.apply(this, arguments)
         }
       }
-
-      if (last && typeof last === 'function') {
-        args[args.length - 1] = wrappedCallback(last)
-      } else if (Array.isArray(last) && typeof last[last.length - 1] === 'function') {
-        last[last.length - 1] = wrappedCallback(last[last.length - 1])
-      } else {
-        args.push(wrappedCallback(function () { }))
+    })
+
+  Shimmer.wrap(redis.Multi.prototype, 'redis.Multi.prototype', _instrumentedCommands,
+    function (original, name) {
+      return function () {
+        var args = Array.prototype.slice.apply(arguments)
+        this.__trace.push(name)
+        return original.apply(this, args)
       }
-
-      agent.clientSend({
-        id: requestId,
-        spanId: spanId,
-        host: host,
-        time: clientSendTime,
-        method: command,
-        type: agent.CLIENT_SEND,
-        url: 'unknown'
-      })
-
-      return original.apply(this, args)
-    }
-  })
+    })
 
   return redis
 }

lib/instrumentations/redis.spec.e2e.js

@@ -23,7 +23,9 @@ describe('redis module wrapper', function () {
 
   it('should instrument operation test #1', function (done) {
     var agent = fakeAgent(this.sandbox)
-    var redis = wrap(require('redis'), agent)
+    var redis = wrap(require('redis'), agent, {
+      version: '2.0.0'
+    })
     var client = redis.createClient(REDIS_URL)
 
     client.sadd('x', 6, function () {
@@ -45,7 +47,9 @@ describe('redis module wrapper', function () {
 
   it('should instrument operation test #2', function () {
     var agent = fakeAgent(this.sandbox)
-    var redis = wrap(require('redis'), agent)
+    var redis = wrap(require('redis'), agent, {
+      version: '2.0.0'
+    })
     var client = redis.createClient(REDIS_URL)
 
     for (var i = 0; i < 9; ++i) {
@@ -57,18 +61,28 @@ describe('redis module wrapper', function () {
 
   it('should work with multi', function (done) {
     var agent = fakeAgent(this.sandbox)
-    var redis = wrap(require('redis'), agent)
+    var redis = wrap(require('redis'), agent, {
+      version: '2.0.0'
+    })
 
     var client = redis.createClient(REDIS_URL)
 
-    client.multi()
-      .sadd('x', 6)
-      .srem('x', 7)
-      .exec(function () {
+    var multi1 = client.multi().sadd('x', 6)
+    var multi2 = client.multi().sadd('x', 5)
+
+    multi1.srem('x', 6)
+    multi2.zadd('x', 5)
+
+    multi1.exec(function () {
+      expect(agent.clientReceive).to.have.been.called
+      expect(agent.clientReceive.args[0][0].method)
+        .to.eql('multi: sadd, srem')
+      multi2.exec(function () {
+        expect(agent.clientReceive.args[1][0].method)
+          .to.eql('multi: sadd, zadd')
         expect(agent.clientReceive).to.have.been.called
         done()
       })
-
-    expect(agent.clientSend).to.have.callCount(4)
+    })
   })
 })