fix: Updated redis v4 instrumentation to work with transactions(multi/exec)

lib/instrumentation/@node-redis/client.js

@@ -7,63 +7,78 @@
 
 const {
   OperationSpec,
-  params: { DatastoreParameters }
+  params: { DatastoreParameters },
+  ClassWrapSpec
 } = require('../../shim/specs')
-const CLIENT_COMMANDS = ['select', 'quit', 'SELECT', 'QUIT']
-const opts = Symbol('clientOptions')
+const { redisClientOpts } = require('../../symbols')
 
 module.exports = function initialize(_agent, redis, _moduleName, shim) {
   shim.setDatastore(shim.REDIS)
-  const COMMANDS = Object.keys(shim.require('dist/lib/client/commands.js').default)
-  const CMDS_TO_INSTRUMENT = [...COMMANDS, ...CLIENT_COMMANDS]
-  shim.wrap(redis, 'createClient', function wrapCreateClient(shim, original) {
-    return function wrappedCreateClient() {
-      const client = original.apply(this, arguments)
-      client[opts] = getRedisParams(client.options)
-      CMDS_TO_INSTRUMENT.forEach(instrumentClientCommand.bind(null, shim, client))
-      if (client.options.legacyMode) {
-        client.v4[opts] = getRedisParams(client.options)
-        CMDS_TO_INSTRUMENT.forEach(instrumentClientCommand.bind(null, shim, client.v4))
+  const commandsQueue = shim.require('dist/lib/client/commands-queue.js')
+
+  shim.wrapClass(
+    commandsQueue,
+    'default',
+    new ClassWrapSpec({
+      post: function postConstructor(shim) {
+        instrumentAddCommand({ shim, commandsQueue: this })
       }
+    })
+  )
+
+  shim.wrap(redis, 'createClient', function wrapCreateClient(_shim, createClient) {
+    return function wrappedCreateClient(options) {
+      // saving connection opts to shim
+      // since the RedisCommandsQueue gets constructed at createClient
+      // we can delete the symbol afterwards to ensure the appropriate
+      // connection options are for the given RedisCommandsQueue
+      shim[redisClientOpts] = getRedisParams(options)
+      const client = createClient.apply(this, arguments)
+      delete shim[redisClientOpts]
       return client
     }
   })
 }
 
 /**
- * Instruments a given command on the client by calling `shim.recordOperation`
+ * Instruments a given command when added to the command queue by calling `shim.recordOperation`
  *
- * @param {Shim} shim shim instance
- * @param {object} client redis client instance
- * @param {string} cmd command to instrument
+ * @param {object} params
+ * @param {Shim} params.shim shim instance
+ * @param {object} params.commandsQueue instance
  */
-function instrumentClientCommand(shim, client, cmd) {
+function instrumentAddCommand({ shim, commandsQueue }) {
   const { agent } = shim
+  const clientOpts = shim[redisClientOpts]
 
-  shim.recordOperation(client, cmd, function wrapCommand(_shim, _fn, _fnName, args) {
-    const [key, value] = args
-    const parameters = Object.assign({}, client[opts])
-    // If selecting a database, subsequent commands
-    // will be using said database, update the clientOptions
-    // but not the current parameters(feature parity with v3)
-    if (cmd.toLowerCase() === 'select') {
-      client[opts].database_name = key
-    }
-    if (agent.config.attributes.enabled) {
-      if (key) {
-        parameters.key = JSON.stringify(key)
+  shim.recordOperation(
+    commandsQueue,
+    'addCommand',
+    function wrapAddCommand(_shim, _fn, _fnName, args) {
+      const [cmd, key, value] = args[0]
+      const parameters = Object.assign({}, clientOpts)
+      // If selecting a database, subsequent commands
+      // will be using said database, update the clientOpts
+      // but not the current parameters(feature parity with v3)
+      if (cmd.toLowerCase() === 'select') {
+        clientOpts.database_name = key
       }
-      if (value) {
-        parameters.value = JSON.stringify(value)
+      if (agent.config.attributes.enabled) {
+        if (key) {
+          parameters.key = JSON.stringify(key)
+        }
+        if (value) {
+          parameters.value = JSON.stringify(value)
+        }
       }
-    }
 
-    return new OperationSpec({
-      name: (cmd && cmd.toLowerCase()) || 'other',
-      parameters,
-      promise: true
-    })
-  })
+      return new OperationSpec({
+        name: (cmd && cmd.toLowerCase()) || 'other',
+        parameters,
+        promise: true
+      })
+    }
+  )
 }
 
 /**
@@ -73,6 +88,21 @@ function instrumentClientCommand(shim, client, cmd) {
  * @returns {object} params
  */
 function getRedisParams(clientOpts) {
+  // need to replicate logic done in RedisClient
+  // to parse the url to assign to socket.host/port
+  // see: https://github.com/redis/node-redis/blob/5576a0db492cda2cd88e09881bc330aa956dd0f5/packages/client/lib/client/index.ts#L160
+  if (clientOpts?.url) {
+    const parsedURL = new URL(clientOpts.url)
+    clientOpts.socket = { host: parsedURL.hostname }
+    if (parsedURL.port) {
+      clientOpts.socket.port = parsedURL.port
+    }
+
+    if (parsedURL.pathname) {
+      clientOpts.database = parsedURL.pathname.substring(1)
+    }
+  }
+
   return new DatastoreParameters({
     host: clientOpts?.socket?.host || 'localhost',
     port_path_or_id: clientOpts?.socket?.path || clientOpts?.socket?.port || '6379',

lib/symbols.js

@@ -27,6 +27,7 @@ module.exports = {
   langchainRunId: Symbol('runId'),
   prismaConnection: Symbol('prismaConnection'),
   prismaModelCall: Symbol('modelCall'),
+  redisClientOpts: Symbol('clientOptions'),
   segment: Symbol('segment'),
   shim: Symbol('shim'),
   storeDatabase: Symbol('storeDatabase'),

test/unit/instrumentation/redis.test.js

@@ -6,54 +6,55 @@
 'use strict'
 
 const tap = require('tap')
-const client = require('../../../lib/instrumentation/@node-redis/client')
+const sinon = require('sinon')
+const helper = require('../../lib/agent_helper')
+const DatastoreShim = require('../../../lib/shim/datastore-shim.js')
+const { redisClientOpts } = require('../../../lib/symbols')
 
 tap.test('getRedisParams should behave as expected', function (t) {
-  t.autoend()
-
+  const { getRedisParams } = require('../../../lib/instrumentation/@node-redis/client')
   t.test('given no opts, should return sensible defaults', function (t) {
-    t.autoend()
-    const params = client.getRedisParams()
+    const params = getRedisParams()
     const expected = {
       host: 'localhost',
       port_path_or_id: '6379',
       database_name: 0
     }
     t.match(params, expected, 'redis client should be definable without params')
+    t.end()
   })
   t.test('if host/port are defined incorrectly, should return expected defaults', function (t) {
-    t.autoend()
-    const params = client.getRedisParams({ host: 'myLocalHost', port: '1234' })
+    const params = getRedisParams({ host: 'myLocalHost', port: '1234' })
     const expected = {
       host: 'localhost',
       port_path_or_id: '6379',
       database_name: 0
     }
     t.match(params, expected, 'should return sensible defaults if defined without socket')
+    t.end()
   })
   t.test('if host/port are defined correctly, we should see them in config', function (t) {
-    t.autoend()
-    const params = client.getRedisParams({ socket: { host: 'myLocalHost', port: '1234' } })
+    const params = getRedisParams({ socket: { host: 'myLocalHost', port: '1234' } })
     const expected = {
       host: 'myLocalHost',
       port_path_or_id: '1234',
       database_name: 0
     }
     t.match(params, expected, 'host/port should be returned when defined correctly')
+    t.end()
   })
   t.test('path should be used if defined', function (t) {
-    t.autoend()
-    const params = client.getRedisParams({ socket: { path: '5678' } })
+    const params = getRedisParams({ socket: { path: '5678' } })
     const expected = {
       host: 'localhost',
       port_path_or_id: '5678',
       database_name: 0
     }
     t.match(params, expected, 'path should show up in params')
+    t.end()
   })
   t.test('path should be preferred over port', function (t) {
-    t.autoend()
-    const params = client.getRedisParams({
+    const params = getRedisParams({
       socket: { host: 'myLocalHost', port: '1234', path: '5678' }
     })
     const expected = {
@@ -62,15 +63,145 @@ tap.test('getRedisParams should behave as expected', function (t) {
       database_name: 0
     }
     t.match(params, expected, 'path should show up in params')
+    t.end()
   })
   t.test('database name should be definable', function (t) {
-    t.autoend()
-    const params = client.getRedisParams({ database: 12 })
+    const params = getRedisParams({ database: 12 })
     const expected = {
       host: 'localhost',
       port_path_or_id: '6379',
       database_name: 12
     }
     t.match(params, expected, 'database should be definable')
+    t.end()
+  })
+
+  t.test('host/port/database should be extracted from url when it exists', function (t) {
+    const params = getRedisParams({ url: 'redis://host:6369/db' })
+    const expected = {
+      host: 'host',
+      port_path_or_id: '6369',
+      database_name: 'db'
+    }
+    t.match(params, expected, 'host/port/database should match')
+    t.end()
+  })
+
+  t.test('should default port to 6379 when no port specified in URL', function (t) {
+    const params = getRedisParams({ url: 'redis://host/db' })
+    const expected = {
+      host: 'host',
+      port_path_or_id: '6379',
+      database_name: 'db'
+    }
+    t.match(params, expected, 'host/port/database should match')
+    t.end()
+  })
+
+  t.test('should default database to 0 when no db pecified in URL', function (t) {
+    const params = getRedisParams({ url: 'redis://host' })
+    const expected = {
+      host: 'host',
+      port_path_or_id: '6379',
+      database_name: 0
+    }
+    t.match(params, expected, 'host/port/database should match')
+    t.end()
+  })
+  t.end()
+})
+
+tap.test('createClient saves connection options', function (t) {
+  t.beforeEach((t) => {
+    t.context.sandbox = sinon.createSandbox()
+    t.context.agent = helper.loadMockedAgent()
+    t.context.shim = new DatastoreShim(t.context.agent, 'redis')
+    t.context.instrumentation = require('../../../lib/instrumentation/@node-redis/client')
+    t.context.clients = {
+      1: { socket: { host: '1', port: 2 } },
+      2: { socket: { host: '2', port: 3 } }
+    }
+    let i = 0
+    class CommandQueueClass {
+      constructor() {
+        i++
+        this.id = i
+        const expectedValues = t.context.clients[this.id]
+        t.match(t.context.shim[redisClientOpts], {
+          host: expectedValues.socket.host,
+          port_path_or_id: expectedValues.socket.port
+        })
+      }
+
+      async addCommand() {}
+    }
+
+    const commandQueueStub = { default: CommandQueueClass }
+    const redis = Object.create({
+      createClient: function () {
+        const instance = Object.create({})
+        // eslint-disable-next-line new-cap
+        instance.queue = new commandQueueStub.default()
+        return instance
+      }
+    })
+
+    t.context.sandbox.stub(t.context.shim, 'require').returns(commandQueueStub)
+    t.context.redis = redis
+  })
+
+  t.afterEach((t) => {
+    helper.unloadAgent(t.context.agent)
+    t.context.sandbox.restore()
+  })
+
+  t.test('should remove connect options after creation', function (t) {
+    const { agent, redis, shim, instrumentation, clients } = t.context
+    instrumentation(agent, redis, 'redis', shim)
+    redis.createClient(clients[1])
+    t.notOk(shim[redisClientOpts], 'should remove client options after creation')
+    redis.createClient(clients[2])
+    t.notOk(shim[redisClientOpts], 'should remove client options after creation')
+    t.end()
+  })
+
+  t.test('should keep the connection details per client', function (t) {
+    const { agent, redis, shim, instrumentation, clients } = t.context
+    instrumentation(agent, redis, 'redis', shim)
+    const client = redis.createClient(clients[1])
+    const client2 = redis.createClient(clients[2])
+    helper.runInTransaction(agent, async function (tx) {
+      await client.queue.addCommand(['test', 'key', 'value'])
+      await client2.queue.addCommand(['test2', 'key2', 'value2'])
+      const [redisSegment, redisSegment2] = tx.trace.root.children
+      const attrs = redisSegment.getAttributes()
+      t.same(
+        attrs,
+        {
+          host: '1',
+          port_path_or_id: 2,
+          key: '"key"',
+          value: '"value"',
+          product: 'Redis',
+          database_name: '0'
+        },
+        'should have appropriate segment attrs'
+      )
+      const attrs2 = redisSegment2.getAttributes()
+      t.same(
+        attrs2,
+        {
+          host: '2',
+          port_path_or_id: 3,
+          key: '"key2"',
+          value: '"value2"',
+          product: 'Redis',
+          database_name: '0'
+        },
+        'should have appropriate segment attrs'
+      )
+      t.end()
+    })
   })
+  t.end()
 })

test/versioned/redis/redis-v4.tap.js

@@ -116,6 +116,30 @@ test('Redis instrumentation', function (t) {
     })
   })
 
+  t.test('should handle multi commands', function (t) {
+    helper.runInTransaction(agent, async function transactionInScope() {
+      const transaction = agent.getTransaction()
+      const results = await client.multi().set('multi-key', 'multi-value').get('multi-key').exec()
+
+      t.same(results, ['OK', 'multi-value'], 'should return expected results')
+      transaction.end()
+      const unscoped = transaction.metrics.unscoped
+      const expected = {
+        'Datastore/all': 4,
+        'Datastore/allWeb': 4,
+        'Datastore/Redis/all': 4,
+        'Datastore/Redis/allWeb': 4,
+        'Datastore/operation/Redis/multi': 1,
+        'Datastore/operation/Redis/set': 1,
+        'Datastore/operation/Redis/get': 1,
+        'Datastore/operation/Redis/exec': 1
+      }
+      expected['Datastore/instance/Redis/' + HOST_ID] = 4
+      checkMetrics(t, unscoped, expected)
+      t.end()
+    })
+  })
+
   t.test('should add `key` attribute to trace segment', function (t) {
     t.notOk(agent.getTransaction(), 'no transaction should be in play')
     agent.config.attributes.enabled = true