Skip to content

Commit 4349e82

Browse files
wconti27wconti27
authored
remove async storage from redis instrumentations (#5921)
* remove async storage from redis instrumentations
1 parent 2c87998 commit 4349e82

File tree

5 files changed

+66
-71
lines changed

5 files changed

+66
-71
lines changed

‎packages/datadog-instrumentations/src/ioredis.js

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,7 @@
22

33
const {
44
channel,
5-
addHook,
6-
AsyncResource
5+
addHook
76
} = require('./helpers/instrument')
87
const shimmer = require('../../datadog-shimmer')
98

@@ -22,14 +21,9 @@ addHook({ name: 'ioredis', versions: ['>=2'] }, Redis => {
2221
const db = options.db
2322
const connectionOptions = { host: options.host, port: options.port }
2423

25-
const asyncResource = new AsyncResource('bound-anonymous-fn')
26-
return asyncResource.runInAsyncScope(() => {
27-
startCh.publish({ db, command: command.name, args: command.args, connectionOptions, connectionName })
28-
29-
const onResolve = asyncResource.bind(() => finish(finishCh, errorCh))
30-
const onReject = asyncResource.bind(err => finish(finishCh, errorCh, err))
31-
32-
command.promise.then(onResolve, onReject)
24+
const ctx = { db, command: command.name, args: command.args, connectionOptions, connectionName }
25+
return startCh.runStores(ctx, () => {
26+
command.promise.then(() => finish(finishCh, errorCh, ctx), err => finish(finishCh, errorCh, ctx, err))
3327

3428
try {
3529
return sendCommand.apply(this, arguments)
@@ -43,9 +37,10 @@ addHook({ name: 'ioredis', versions: ['>=2'] }, Redis => {
4337
return Redis
4438
})
4539

46-
function finish (finishCh, errorCh, error) {
40+
function finish (finishCh, errorCh, ctx, error) {
4741
if (error) {
48-
errorCh.publish(error)
42+
ctx.error = error
43+
errorCh.publish(ctx)
4944
}
50-
finishCh.publish()
45+
finishCh.publish(ctx)
5146
}

‎packages/datadog-instrumentations/src/iovalkey.js

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,7 @@
22

33
const {
44
channel,
5-
addHook,
6-
AsyncResource
5+
addHook
76
} = require('./helpers/instrument')
87
const shimmer = require('../../datadog-shimmer')
98

@@ -22,19 +21,15 @@ addHook({ name: 'iovalkey', versions: ['>=0.0.1'] }, Valkey => {
2221
const db = options.db
2322
const connectionOptions = { host: options.host, port: options.port }
2423

25-
const asyncResource = new AsyncResource('bound-anonymous-fn')
26-
return asyncResource.runInAsyncScope(() => {
27-
startCh.publish({ db, command: command.name, args: command.args, connectionOptions, connectionName })
28-
29-
const onResolve = asyncResource.bind(() => finishCh.publish())
30-
const onReject = asyncResource.bind(err => finish(finishCh, errorCh, err))
31-
32-
command.promise.then(onResolve, onReject)
24+
const ctx = { db, command: command.name, args: command.args, connectionOptions, connectionName }
25+
return startCh.runStores(ctx, () => {
26+
command.promise.then(() => finish(finishCh, errorCh, ctx), err => finish(finishCh, errorCh, ctx, err))
3327

3428
try {
3529
return sendCommand.apply(this, arguments)
3630
} catch (err) {
37-
errorCh.publish(err)
31+
ctx.error = err
32+
errorCh.publish(ctx)
3833

3934
throw err
4035
}
@@ -43,9 +38,10 @@ addHook({ name: 'iovalkey', versions: ['>=0.0.1'] }, Valkey => {
4338
return Valkey
4439
})
4540

46-
function finish (finishCh, errorCh, error) {
41+
function finish (finishCh, errorCh, ctx, error) {
4742
if (error) {
48-
errorCh.publish(error)
43+
ctx.error = error
44+
errorCh.publish(ctx)
4945
}
50-
finishCh.publish()
46+
finishCh.publish(ctx)
5147
}

‎packages/datadog-instrumentations/src/redis.js

Lines changed: 30 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,7 @@
22

33
const {
44
channel,
5-
addHook,
6-
AsyncResource
5+
addHook
76
} = require('./helpers/instrument')
87
const shimmer = require('../../datadog-shimmer')
98

@@ -22,15 +21,11 @@ function wrapAddCommand (addCommand) {
2221
const name = command[0]
2322
const args = command.slice(1)
2423

25-
const asyncResource = new AsyncResource('bound-anonymous-fn')
26-
return asyncResource.runInAsyncScope(() => {
27-
start(this, name, args, this._url)
28-
24+
const ctx = getStartCtx(this, name, args, this._url)
25+
return startCh.runStores(ctx, () => {
2926
const res = addCommand.apply(this, arguments)
30-
const onResolve = asyncResource.bind(() => finish(finishCh, errorCh))
31-
const onReject = asyncResource.bind(err => finish(finishCh, errorCh, err))
3227

33-
res.then(onResolve, onReject)
28+
res.then(() => finish(finishCh, errorCh, ctx), err => finish(finishCh, errorCh, ctx, err))
3429

3530
return res
3631
})
@@ -94,14 +89,9 @@ addHook({ name: 'redis', versions: ['>=2.6 <4'] }, redis => {
9489

9590
if (!options.callback) return internalSendCommand.apply(this, arguments)
9691

97-
const callbackResource = new AsyncResource('bound-anonymous-fn')
98-
const asyncResource = new AsyncResource('bound-anonymous-fn')
99-
const cb = callbackResource.bind(options.callback)
100-
101-
return asyncResource.runInAsyncScope(() => {
102-
start(this, options.command, options.args)
103-
104-
options.callback = asyncResource.bind(wrapCallback(finishCh, errorCh, cb))
92+
const ctx = getStartCtx(this, options.command, options.args)
93+
return startCh.runStores(ctx, () => {
94+
options.callback = wrapCallback(finishCh, errorCh, ctx, options.callback)
10595

10696
try {
10797
return internalSendCommand.apply(this, arguments)
@@ -121,26 +111,21 @@ addHook({ name: 'redis', versions: ['>=0.12 <2.6'] }, redis => {
121111
return sendCommand.apply(this, arguments)
122112
}
123113

124-
const callbackResource = new AsyncResource('bound-anonymous-fn')
125-
const asyncResource = new AsyncResource('bound-anonymous-fn')
126-
127-
return asyncResource.runInAsyncScope(() => {
128-
start(this, command, args)
129-
114+
const ctx = getStartCtx(this, command, args)
115+
return startCh.runStores(ctx, () => {
130116
if (typeof callback === 'function') {
131-
const cb = callbackResource.bind(callback)
132-
arguments[2] = asyncResource.bind(wrapCallback(finishCh, errorCh, cb))
117+
arguments[2] = wrapCallback(finishCh, errorCh, ctx, callback)
133118
} else if (Array.isArray(args) && typeof args.at(-1) === 'function') {
134-
const cb = callbackResource.bind(args.at(-1))
135-
args[args.length - 1] = asyncResource.bind(wrapCallback(finishCh, errorCh, cb))
119+
args[args.length - 1] = wrapCallback(finishCh, errorCh, ctx, args.at(-1))
136120
} else {
137-
arguments[2] = asyncResource.bind(wrapCallback(finishCh, errorCh))
121+
arguments[2] = wrapCallback(finishCh, errorCh, ctx)
138122
}
139123

140124
try {
141125
return sendCommand.apply(this, arguments)
142126
} catch (err) {
143-
errorCh.publish(err)
127+
ctx.error = err
128+
errorCh.publish(ctx)
144129

145130
throw err
146131
}
@@ -149,24 +134,28 @@ addHook({ name: 'redis', versions: ['>=0.12 <2.6'] }, redis => {
149134
return redis
150135
})
151136

152-
function start (client, command, args, url = {}) {
153-
const db = client.selected_db
154-
const connectionOptions = client.connection_options || client.connection_option || client.connectionOption || url
155-
startCh.publish({ db, command, args, connectionOptions })
137+
function getStartCtx (client, command, args, url = {}) {
138+
return {
139+
db: client.selected_db,
140+
command,
141+
args,
142+
connectionOptions: client.connection_options || client.connection_option || client.connectionOption || url
143+
}
156144
}
157145

158-
function wrapCallback (finishCh, errorCh, callback) {
146+
function wrapCallback (finishCh, errorCh, ctx, callback) {
159147
return shimmer.wrapFunction(callback, callback => function (err) {
160-
finish(finishCh, errorCh, err)
161-
if (callback) {
162-
return callback.apply(this, arguments)
163-
}
148+
return finish(finishCh, errorCh, ctx, err, callback, this, arguments)
164149
})
165150
}
166151

167-
function finish (finishCh, errorCh, error) {
152+
function finish (finishCh, errorCh, ctx, error, callback, thisArg, args) {
168153
if (error) {
169-
errorCh.publish(error)
154+
ctx.error = error
155+
errorCh.publish(ctx)
156+
}
157+
if (callback) {
158+
return finishCh.runStores(ctx, callback, thisArg, ...args)
170159
}
171-
finishCh.publish()
160+
finishCh.publish(ctx)
172161
}

‎packages/datadog-plugin-redis/src/index.js

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,14 @@ class RedisPlugin extends CachePlugin {
1313
this._spanType = 'redis'
1414
}
1515

16-
start ({ db, command, args, connectionOptions = {}, connectionName }) {
16+
bindStart (ctx) {
17+
const { db, command, args, connectionOptions, connectionName } = ctx
18+
1719
const resource = command
1820
const normalizedCommand = command.toUpperCase()
19-
if (!this.config.filter(normalizedCommand)) return this.skip()
21+
if (!this.config.filter(normalizedCommand)) {
22+
return { noop: true }
23+
}
2024

2125
this.startSpan({
2226
resource,
@@ -29,7 +33,9 @@ class RedisPlugin extends CachePlugin {
2933
'out.host': connectionOptions.host,
3034
[CLIENT_PORT_KEY]: connectionOptions.port
3135
}
32-
})
36+
}, ctx)
37+
38+
return ctx.currentStore
3339
}
3440

3541
configure (config) {

‎packages/datadog-plugin-redis/test/client.spec.js

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,15 @@ describe('Plugin', () => {
132132
async () => client.get('foo'),
133133
rawExpectedSchema.outbound
134134
)
135+
136+
it('should restore the parent context in the callback', async () => {
137+
const span = {}
138+
tracer.scope().activate(span, () => {
139+
client.get('foo', () => {
140+
expect(span.context().active()).to.equal(span)
141+
})
142+
})
143+
})
135144
})
136145

137146
describe('with configuration', () => {

0 commit comments

Comments
 (0)