Add support for redis 2.6

googleapis · Jun 2, 2016 · 0af3c57 · 0af3c57
1 parent c70ec53
commit 0af3c57
Showing 1 changed file with 62 additions and 23 deletions.
diff --git a/lib/hooks/userspace/hook-redis.js b/lib/hooks/userspace/hook-redis.js
@@ -22,7 +22,7 @@ var semver = require('semver');
 var SpanData = require('../../span-data.js');
 var agent;
 
-var SUPPORTED_VERSIONS = '<=2.x';
+var SUPPORTED_VERSIONS = '<=2.6.x';
 
 function createClientWrap(createClient) {
   return function createClientTrace() {
@@ -56,31 +56,61 @@ function streamListenersWrap(install_stream_listeners) {
   };
 }
 
+function setupSpan(cmd, args, skipped_frames) {
+  var labels = { command: cmd };
+  if (agent.config_.enhancedDatabaseReporting) {
+    labels.arguments = JSON.stringify(args);
+  }
+  return agent.startSpan('redis-' + cmd, labels, skipped_frames + 1);
+}
+
+function startSpanFromArguments(cmd, args, cb, send_command) {
+  if (!cmd || !args || typeof cmd !== 'string' || !Array.isArray(args) ||
+    (cb && typeof cb !== 'function')) {
+    return send_command(cmd, args, cb);
+  }
+  if (!cb) {
+    if (typeof args[args.length - 1] === 'function' ||
+        typeof args[args.length - 1] === 'undefined') {
+      cb = args.pop();
+    }
+  }
+  var span = setupSpan(cmd, args, 1);
+  return send_command(cmd, args, wrapCallback(span, cb));
+}
+
+function validRootSpan(root, cmd, args) {
+  if (!root) {
+    agent.logger.debug('Untraced redis command:', cmd, args);
+    return false;
+  } else if (root === SpanData.nullSpan) {
+    return false;
+  }
+  return true;
+}
+
+function internalSendCommandWrap(internal_send_command) {
+  return function internal_send_command_trace(cmd, args, cb) {
+    var root = cls.getRootContext();
+    if (!validRootSpan(root, cmd, args)) {
+      return internal_send_command.call(this, cmd, args, cb);
+    }
+    if (arguments.length === 1 && typeof cmd === 'object') {
+      var span = setupSpan(cmd.command, cmd.args, 0);
+      cmd.callback = wrapCallback(span, cmd.callback);
+      return internal_send_command.call(this, cmd);
+    }
+    return startSpanFromArguments(cmd, args, cb, internal_send_command.bind(this));
+  };
+}
+
 function sendCommandWrap(send_command) {
   return function send_command_trace(cmd, args, cb) {
     var root = cls.getRootContext();
-    if (!root) {
-      agent.logger.debug('Untraced redis command:', cmd, args);
-      return send_command.call(this, cmd, args, cb);
-    } else if (root === SpanData.nullSpan) {
+    if (!validRootSpan(root, cmd, args)) {
       return send_command.call(this, cmd, args, cb);
     }
-    if (!cmd || !args || typeof cmd !== 'string' || !Array.isArray(args) ||
-        (cb && typeof cb !== 'function')) {
-      return send_command.call(this, cmd, args, cb);
-    }
-    if (!cb) {
-      if (typeof args[args.length - 1] === 'function' ||
-          typeof args[args.length - 1] === 'undefined') {
-        cb = args.pop();
-      }
-    }
-    var labels = { command: cmd };
-    if (agent.config_.enhancedDatabaseReporting) {
-      labels.arguments = JSON.stringify(args);
-    }
-    var span = agent.startSpan('redis-' + cmd, labels);
-    return send_command.call(this, cmd, args, wrapCallback(span, cb));
+    return startSpanFromArguments(cmd, args, cb, send_command.bind(this));
   };
 }
 
@@ -113,7 +143,12 @@ module.exports = function(version_, agent_) {
     '': {
       patch: function(redis) {
         agent = agent_;
-        shimmer.wrap(redis.RedisClient.prototype, 'send_command', sendCommandWrap);
+        if (semver.satisfies(version_, '<2.6')) {
+          shimmer.wrap(redis.RedisClient.prototype, 'send_command', sendCommandWrap);
+        } else {
+          shimmer.wrap(redis.RedisClient.prototype, 'internal_send_command',
+            internalSendCommandWrap);
+        }
         if (semver.satisfies(version_, '<=2.3.x')) {
           shimmer.wrap(redis.RedisClient.prototype, 'install_stream_listeners',
             streamListenersWrap);
@@ -124,7 +159,11 @@ module.exports = function(version_, agent_) {
         shimmer.wrap(redis, 'createClient', createClientWrap);
       },
       unpatch: function(redis) {
-        shimmer.unwrap(redis.RedisClient.prototype, 'send_command');
+        if (semver.satisfies(version_, '<2.6')) {
+          shimmer.unwrap(redis.RedisClient.prototype, 'send_command');
+        } else {
+          shimmer.unwrap(redis.RedisClient.prototype, 'internal_send_command');
+        }
         if (semver.satisfies(version_, '<=2.3.x')) {
           shimmer.unwrap(redis.RedisClient.prototype, 'install_stream_listeners');
         } else {