Skip to content

Commit 39eb842

Browse files
authoredJul 11, 2024··
fix: Updated redis v4 instrumentation to work with transactions(multi/exec) (#2343)
1 parent 5c9e3e6 commit 39eb842

File tree

5 files changed

+270
-54
lines changed

5 files changed

+270
-54
lines changed
 

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

+69-39
Original file line numberDiff line numberDiff line change
@@ -7,63 +7,78 @@
77

88
const {
99
OperationSpec,
10-
params: { DatastoreParameters }
10+
params: { DatastoreParameters },
11+
ClassWrapSpec
1112
} = require('../../shim/specs')
12-
const CLIENT_COMMANDS = ['select', 'quit', 'SELECT', 'QUIT']
13-
const opts = Symbol('clientOptions')
13+
const { redisClientOpts } = require('../../symbols')
1414

1515
module.exports = function initialize(_agent, redis, _moduleName, shim) {
1616
shim.setDatastore(shim.REDIS)
17-
const COMMANDS = Object.keys(shim.require('dist/lib/client/commands.js').default)
18-
const CMDS_TO_INSTRUMENT = [...COMMANDS, ...CLIENT_COMMANDS]
19-
shim.wrap(redis, 'createClient', function wrapCreateClient(shim, original) {
20-
return function wrappedCreateClient() {
21-
const client = original.apply(this, arguments)
22-
client[opts] = getRedisParams(client.options)
23-
CMDS_TO_INSTRUMENT.forEach(instrumentClientCommand.bind(null, shim, client))
24-
if (client.options.legacyMode) {
25-
client.v4[opts] = getRedisParams(client.options)
26-
CMDS_TO_INSTRUMENT.forEach(instrumentClientCommand.bind(null, shim, client.v4))
17+
const commandsQueue = shim.require('dist/lib/client/commands-queue.js')
18+
19+
shim.wrapClass(
20+
commandsQueue,
21+
'default',
22+
new ClassWrapSpec({
23+
post: function postConstructor(shim) {
24+
instrumentAddCommand({ shim, commandsQueue: this })
2725
}
26+
})
27+
)
28+
29+
shim.wrap(redis, 'createClient', function wrapCreateClient(_shim, createClient) {
30+
return function wrappedCreateClient(options) {
31+
// saving connection opts to shim
32+
// since the RedisCommandsQueue gets constructed at createClient
33+
// we can delete the symbol afterwards to ensure the appropriate
34+
// connection options are for the given RedisCommandsQueue
35+
shim[redisClientOpts] = getRedisParams(options)
36+
const client = createClient.apply(this, arguments)
37+
delete shim[redisClientOpts]
2838
return client
2939
}
3040
})
3141
}
3242

3343
/**
34-
* Instruments a given command on the client by calling `shim.recordOperation`
44+
* Instruments a given command when added to the command queue by calling `shim.recordOperation`
3545
*
36-
* @param {Shim} shim shim instance
37-
* @param {object} client redis client instance
38-
* @param {string} cmd command to instrument
46+
* @param {object} params
47+
* @param {Shim} params.shim shim instance
48+
* @param {object} params.commandsQueue instance
3949
*/
40-
function instrumentClientCommand(shim, client, cmd) {
50+
function instrumentAddCommand({ shim, commandsQueue }) {
4151
const { agent } = shim
52+
const clientOpts = shim[redisClientOpts]
4253

43-
shim.recordOperation(client, cmd, function wrapCommand(_shim, _fn, _fnName, args) {
44-
const [key, value] = args
45-
const parameters = Object.assign({}, client[opts])
46-
// If selecting a database, subsequent commands
47-
// will be using said database, update the clientOptions
48-
// but not the current parameters(feature parity with v3)
49-
if (cmd.toLowerCase() === 'select') {
50-
client[opts].database_name = key
51-
}
52-
if (agent.config.attributes.enabled) {
53-
if (key) {
54-
parameters.key = JSON.stringify(key)
54+
shim.recordOperation(
55+
commandsQueue,
56+
'addCommand',
57+
function wrapAddCommand(_shim, _fn, _fnName, args) {
58+
const [cmd, key, value] = args[0]
59+
const parameters = Object.assign({}, clientOpts)
60+
// If selecting a database, subsequent commands
61+
// will be using said database, update the clientOpts
62+
// but not the current parameters(feature parity with v3)
63+
if (cmd.toLowerCase() === 'select') {
64+
clientOpts.database_name = key
5565
}
56-
if (value) {
57-
parameters.value = JSON.stringify(value)
66+
if (agent.config.attributes.enabled) {
67+
if (key) {
68+
parameters.key = JSON.stringify(key)
69+
}
70+
if (value) {
71+
parameters.value = JSON.stringify(value)
72+
}
5873
}
59-
}
6074

61-
return new OperationSpec({
62-
name: (cmd && cmd.toLowerCase()) || 'other',
63-
parameters,
64-
promise: true
65-
})
66-
})
75+
return new OperationSpec({
76+
name: (cmd && cmd.toLowerCase()) || 'other',
77+
parameters,
78+
promise: true
79+
})
80+
}
81+
)
6782
}
6883

6984
/**
@@ -73,6 +88,21 @@ function instrumentClientCommand(shim, client, cmd) {
7388
* @returns {object} params
7489
*/
7590
function getRedisParams(clientOpts) {
91+
// need to replicate logic done in RedisClient
92+
// to parse the url to assign to socket.host/port
93+
// see: https://github.com/redis/node-redis/blob/5576a0db492cda2cd88e09881bc330aa956dd0f5/packages/client/lib/client/index.ts#L160
94+
if (clientOpts?.url) {
95+
const parsedURL = new URL(clientOpts.url)
96+
clientOpts.socket = { host: parsedURL.hostname }
97+
if (parsedURL.port) {
98+
clientOpts.socket.port = parsedURL.port
99+
}
100+
101+
if (parsedURL.pathname) {
102+
clientOpts.database = parsedURL.pathname.substring(1)
103+
}
104+
}
105+
76106
return new DatastoreParameters({
77107
host: clientOpts?.socket?.host || 'localhost',
78108
port_path_or_id: clientOpts?.socket?.path || clientOpts?.socket?.port || '6379',

‎lib/symbols.js

+1
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ module.exports = {
2727
langchainRunId: Symbol('runId'),
2828
prismaConnection: Symbol('prismaConnection'),
2929
prismaModelCall: Symbol('modelCall'),
30+
redisClientOpts: Symbol('clientOptions'),
3031
segment: Symbol('segment'),
3132
shim: Symbol('shim'),
3233
storeDatabase: Symbol('storeDatabase'),

‎test/unit/instrumentation/redis.test.js

+146-15
Original file line numberDiff line numberDiff line change
@@ -6,54 +6,55 @@
66
'use strict'
77

88
const tap = require('tap')
9-
const client = require('../../../lib/instrumentation/@node-redis/client')
9+
const sinon = require('sinon')
10+
const helper = require('../../lib/agent_helper')
11+
const DatastoreShim = require('../../../lib/shim/datastore-shim.js')
12+
const { redisClientOpts } = require('../../../lib/symbols')
1013

1114
tap.test('getRedisParams should behave as expected', function (t) {
12-
t.autoend()
13-
15+
const { getRedisParams } = require('../../../lib/instrumentation/@node-redis/client')
1416
t.test('given no opts, should return sensible defaults', function (t) {
15-
t.autoend()
16-
const params = client.getRedisParams()
17+
const params = getRedisParams()
1718
const expected = {
1819
host: 'localhost',
1920
port_path_or_id: '6379',
2021
database_name: 0
2122
}
2223
t.match(params, expected, 'redis client should be definable without params')
24+
t.end()
2325
})
2426
t.test('if host/port are defined incorrectly, should return expected defaults', function (t) {
25-
t.autoend()
26-
const params = client.getRedisParams({ host: 'myLocalHost', port: '1234' })
27+
const params = getRedisParams({ host: 'myLocalHost', port: '1234' })
2728
const expected = {
2829
host: 'localhost',
2930
port_path_or_id: '6379',
3031
database_name: 0
3132
}
3233
t.match(params, expected, 'should return sensible defaults if defined without socket')
34+
t.end()
3335
})
3436
t.test('if host/port are defined correctly, we should see them in config', function (t) {
35-
t.autoend()
36-
const params = client.getRedisParams({ socket: { host: 'myLocalHost', port: '1234' } })
37+
const params = getRedisParams({ socket: { host: 'myLocalHost', port: '1234' } })
3738
const expected = {
3839
host: 'myLocalHost',
3940
port_path_or_id: '1234',
4041
database_name: 0
4142
}
4243
t.match(params, expected, 'host/port should be returned when defined correctly')
44+
t.end()
4345
})
4446
t.test('path should be used if defined', function (t) {
45-
t.autoend()
46-
const params = client.getRedisParams({ socket: { path: '5678' } })
47+
const params = getRedisParams({ socket: { path: '5678' } })
4748
const expected = {
4849
host: 'localhost',
4950
port_path_or_id: '5678',
5051
database_name: 0
5152
}
5253
t.match(params, expected, 'path should show up in params')
54+
t.end()
5355
})
5456
t.test('path should be preferred over port', function (t) {
55-
t.autoend()
56-
const params = client.getRedisParams({
57+
const params = getRedisParams({
5758
socket: { host: 'myLocalHost', port: '1234', path: '5678' }
5859
})
5960
const expected = {
@@ -62,15 +63,145 @@ tap.test('getRedisParams should behave as expected', function (t) {
6263
database_name: 0
6364
}
6465
t.match(params, expected, 'path should show up in params')
66+
t.end()
6567
})
6668
t.test('database name should be definable', function (t) {
67-
t.autoend()
68-
const params = client.getRedisParams({ database: 12 })
69+
const params = getRedisParams({ database: 12 })
6970
const expected = {
7071
host: 'localhost',
7172
port_path_or_id: '6379',
7273
database_name: 12
7374
}
7475
t.match(params, expected, 'database should be definable')
76+
t.end()
77+
})
78+
79+
t.test('host/port/database should be extracted from url when it exists', function (t) {
80+
const params = getRedisParams({ url: 'redis://host:6369/db' })
81+
const expected = {
82+
host: 'host',
83+
port_path_or_id: '6369',
84+
database_name: 'db'
85+
}
86+
t.match(params, expected, 'host/port/database should match')
87+
t.end()
88+
})
89+
90+
t.test('should default port to 6379 when no port specified in URL', function (t) {
91+
const params = getRedisParams({ url: 'redis://host/db' })
92+
const expected = {
93+
host: 'host',
94+
port_path_or_id: '6379',
95+
database_name: 'db'
96+
}
97+
t.match(params, expected, 'host/port/database should match')
98+
t.end()
99+
})
100+
101+
t.test('should default database to 0 when no db pecified in URL', function (t) {
102+
const params = getRedisParams({ url: 'redis://host' })
103+
const expected = {
104+
host: 'host',
105+
port_path_or_id: '6379',
106+
database_name: 0
107+
}
108+
t.match(params, expected, 'host/port/database should match')
109+
t.end()
110+
})
111+
t.end()
112+
})
113+
114+
tap.test('createClient saves connection options', function (t) {
115+
t.beforeEach((t) => {
116+
t.context.sandbox = sinon.createSandbox()
117+
t.context.agent = helper.loadMockedAgent()
118+
t.context.shim = new DatastoreShim(t.context.agent, 'redis')
119+
t.context.instrumentation = require('../../../lib/instrumentation/@node-redis/client')
120+
t.context.clients = {
121+
1: { socket: { host: '1', port: 2 } },
122+
2: { socket: { host: '2', port: 3 } }
123+
}
124+
let i = 0
125+
class CommandQueueClass {
126+
constructor() {
127+
i++
128+
this.id = i
129+
const expectedValues = t.context.clients[this.id]
130+
t.match(t.context.shim[redisClientOpts], {
131+
host: expectedValues.socket.host,
132+
port_path_or_id: expectedValues.socket.port
133+
})
134+
}
135+
136+
async addCommand() {}
137+
}
138+
139+
const commandQueueStub = { default: CommandQueueClass }
140+
const redis = Object.create({
141+
createClient: function () {
142+
const instance = Object.create({})
143+
// eslint-disable-next-line new-cap
144+
instance.queue = new commandQueueStub.default()
145+
return instance
146+
}
147+
})
148+
149+
t.context.sandbox.stub(t.context.shim, 'require').returns(commandQueueStub)
150+
t.context.redis = redis
151+
})
152+
153+
t.afterEach((t) => {
154+
helper.unloadAgent(t.context.agent)
155+
t.context.sandbox.restore()
156+
})
157+
158+
t.test('should remove connect options after creation', function (t) {
159+
const { agent, redis, shim, instrumentation, clients } = t.context
160+
instrumentation(agent, redis, 'redis', shim)
161+
redis.createClient(clients[1])
162+
t.notOk(shim[redisClientOpts], 'should remove client options after creation')
163+
redis.createClient(clients[2])
164+
t.notOk(shim[redisClientOpts], 'should remove client options after creation')
165+
t.end()
166+
})
167+
168+
t.test('should keep the connection details per client', function (t) {
169+
const { agent, redis, shim, instrumentation, clients } = t.context
170+
instrumentation(agent, redis, 'redis', shim)
171+
const client = redis.createClient(clients[1])
172+
const client2 = redis.createClient(clients[2])
173+
helper.runInTransaction(agent, async function (tx) {
174+
await client.queue.addCommand(['test', 'key', 'value'])
175+
await client2.queue.addCommand(['test2', 'key2', 'value2'])
176+
const [redisSegment, redisSegment2] = tx.trace.root.children
177+
const attrs = redisSegment.getAttributes()
178+
t.same(
179+
attrs,
180+
{
181+
host: '1',
182+
port_path_or_id: 2,
183+
key: '"key"',
184+
value: '"value"',
185+
product: 'Redis',
186+
database_name: '0'
187+
},
188+
'should have appropriate segment attrs'
189+
)
190+
const attrs2 = redisSegment2.getAttributes()
191+
t.same(
192+
attrs2,
193+
{
194+
host: '2',
195+
port_path_or_id: 3,
196+
key: '"key2"',
197+
value: '"value2"',
198+
product: 'Redis',
199+
database_name: '0'
200+
},
201+
'should have appropriate segment attrs'
202+
)
203+
t.end()
204+
})
75205
})
206+
t.end()
76207
})

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

+24
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,30 @@ test('Redis instrumentation', function (t) {
116116
})
117117
})
118118

119+
t.test('should handle multi commands', function (t) {
120+
helper.runInTransaction(agent, async function transactionInScope() {
121+
const transaction = agent.getTransaction()
122+
const results = await client.multi().set('multi-key', 'multi-value').get('multi-key').exec()
123+
124+
t.same(results, ['OK', 'multi-value'], 'should return expected results')
125+
transaction.end()
126+
const unscoped = transaction.metrics.unscoped
127+
const expected = {
128+
'Datastore/all': 4,
129+
'Datastore/allWeb': 4,
130+
'Datastore/Redis/all': 4,
131+
'Datastore/Redis/allWeb': 4,
132+
'Datastore/operation/Redis/multi': 1,
133+
'Datastore/operation/Redis/set': 1,
134+
'Datastore/operation/Redis/get': 1,
135+
'Datastore/operation/Redis/exec': 1
136+
}
137+
expected['Datastore/instance/Redis/' + HOST_ID] = 4
138+
checkMetrics(t, unscoped, expected)
139+
t.end()
140+
})
141+
})
142+
119143
t.test('should add `key` attribute to trace segment', function (t) {
120144
t.notOk(agent.getTransaction(), 'no transaction should be in play')
121145
agent.config.attributes.enabled = true

‎test/versioned/redis/redis.tap.js

+30
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,36 @@ test('Redis instrumentation', { timeout: 20000 }, function (t) {
185185
})
186186
})
187187

188+
t.test('should handle multi commands', function (t) {
189+
helper.runInTransaction(agent, function transactionInScope() {
190+
const transaction = agent.getTransaction()
191+
client
192+
.multi()
193+
.set('multi-key', 'multi-value')
194+
.get('multi-key')
195+
.exec(function (error, data) {
196+
t.same(data, ['OK', 'multi-value'], 'should return expected results')
197+
t.error(error)
198+
199+
transaction.end()
200+
const unscoped = transaction.metrics.unscoped
201+
const expected = {
202+
'Datastore/all': 4,
203+
'Datastore/allWeb': 4,
204+
'Datastore/Redis/all': 4,
205+
'Datastore/Redis/allWeb': 4,
206+
'Datastore/operation/Redis/multi': 1,
207+
'Datastore/operation/Redis/set': 1,
208+
'Datastore/operation/Redis/get': 1,
209+
'Datastore/operation/Redis/exec': 1
210+
}
211+
expected['Datastore/instance/Redis/' + HOST_ID] = 4
212+
checkMetrics(t, unscoped, expected)
213+
t.end()
214+
})
215+
})
216+
})
217+
188218
t.test('should add `key` attribute to trace segment', function (t) {
189219
agent.config.attributes.enabled = true
190220

0 commit comments

Comments
 (0)
Please sign in to comment.