Skip to content

Commit b7a7392

Browse files
authored
fix super support for traceCallback (#7327)
1 parent 081a148 commit b7a7392

File tree

9 files changed

+237
-22
lines changed

9 files changed

+237
-22
lines changed

packages/datadog-instrumentations/src/helpers/rewriter/compiler.js

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,4 +24,10 @@ module.exports = {
2424

2525
return esquery.traverse(ast, selector, visitor)
2626
},
27+
28+
query: (ast, query) => {
29+
esquery ??= require('../../../../../vendor/dist/esquery').default
30+
31+
return esquery.query(ast, query)
32+
},
2733
}

packages/datadog-instrumentations/src/helpers/rewriter/transforms.js

Lines changed: 73 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
'use strict'
22

3-
const { parse } = require('./compiler')
3+
const { parse, query } = require('./compiler')
44

55
const tracingChannelPredicate = (node) => (
66
node.specifiers?.[0]?.local?.name === 'tr_ch_apm_tracingChannel' ||
@@ -106,34 +106,91 @@ function traceInstanceMethod (state, node, program) {
106106
}
107107

108108
function wrap (state, node) {
109-
const { channelName, operator, functionQuery: { index = -1 } } = state
109+
const { channelName, operator } = state
110+
111+
if (operator === 'traceCallback') return wrapCallback(state, node)
112+
110113
const async = operator === 'tracePromise' ? 'async' : ''
111114
const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
112-
const tracedArgs = operator === 'traceCallback'
113-
? `__apm$original_args.splice(${index}, 1, arguments[${index >= 0 ? index : `arguments.length + ${index}`}])`
114-
: '__apm$original_args'
115-
const traceParams = operator === 'traceCallback'
116-
? `__apm$traced, ${index}`
117-
: '__apm$traced'
118115
const wrapper = parse(`
119116
function wrapper () {
120-
const __apm$original_args = arguments;
121-
const __apm$traced = ${async} function () {
117+
const __apm$traced = ${async} () => {
122118
const __apm$wrapped = () => {};
123-
const __apm$traced_args = ${tracedArgs};
124-
return __apm$wrapped.apply(this, __apm$traced_args);
119+
return __apm$wrapped.apply(this, arguments);
125120
};
126-
if (!${channelVariable}.hasSubscribers) return __apm$traced.apply(this, arguments);
127-
return ${channelVariable}.${operator}(${traceParams}, {
121+
if (!${channelVariable}.hasSubscribers) return __apm$traced();
122+
return ${channelVariable}.${operator}(__apm$traced, {
123+
arguments,
124+
self: this,
125+
moduleVersion: "1.0.0"
126+
});
127+
}
128+
`).body[0].body // Extract only block statement of function body.
129+
130+
// Replace the right-hand side assignment of `const __apm$wrapped = () => {}`.
131+
query(wrapper, '[id.name=__apm$wrapped]')[0].init = node
132+
133+
return wrapper
134+
}
135+
136+
function wrapCallback (state, node) {
137+
const { channelName, functionQuery: { index = -1 } } = state
138+
const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
139+
const wrapper = parse(`
140+
function wrapper () {
141+
const __apm$cb = Array.prototype.at.call(arguments, ${index});
142+
const __apm$ctx = {
128143
arguments,
129144
self: this,
130145
moduleVersion: "1.0.0"
131-
}, this, ...arguments);
146+
};
147+
const __apm$traced = () => {
148+
const __apm$wrapped = () => {};
149+
return __apm$wrapped.apply(this, arguments);
150+
};
151+
152+
if (!${channelVariable}.start.hasSubscribers) return __apm$traced();
153+
154+
function __apm$wrappedCb(err, res) {
155+
if (err) {
156+
__apm$ctx.error = err;
157+
${channelVariable}.error.publish(__apm$ctx);
158+
} else {
159+
__apm$ctx.result = res;
160+
}
161+
162+
${channelVariable}.asyncStart.runStores(__apm$ctx, () => {
163+
try {
164+
if (__apm$cb) {
165+
return __apm$cb.apply(this, arguments);
166+
}
167+
} finally {
168+
${channelVariable}.asyncEnd.publish(__apm$ctx);
169+
}
170+
});
171+
}
172+
173+
if (typeof __apm$cb !== 'function') {
174+
return __apm$traced();
175+
}
176+
Array.prototype.splice.call(arguments, ${index}, 1, __apm$wrappedCb);
177+
178+
return ${channelVariable}.start.runStores(__apm$ctx, () => {
179+
try {
180+
return __apm$traced();
181+
} catch (err) {
182+
__apm$ctx.error = err;
183+
${channelVariable}.error.publish(__apm$ctx);
184+
throw err;
185+
} finally {
186+
${channelVariable}.end.publish(__apm$ctx);
187+
}
188+
});
132189
}
133190
`).body[0].body // Extract only block statement of function body.
134191

135192
// Replace the right-hand side assignment of `const __apm$wrapped = () => {}`.
136-
wrapper.body[1].declarations[0].init.body.body[0].declarations[0].init = node
193+
query(wrapper, '[id.name=__apm$wrapped]')[0].init = node
137194

138195
return wrapper
139196
}

packages/datadog-instrumentations/test/helpers/rewriter/index.spec.js

Lines changed: 90 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,19 @@ describe('check-require-cache', () => {
4242
},
4343
channelName: 'test_invoke'
4444
},
45+
{
46+
module: {
47+
name: 'test-trace-sync-super',
48+
versionRange: '>=0.1',
49+
filePath: 'index.js'
50+
},
51+
functionQuery: {
52+
methodName: 'test',
53+
kind: 'Sync',
54+
className: 'B'
55+
},
56+
channelName: 'test_invoke'
57+
},
4558
{
4659
module: {
4760
name: 'test-trace-async',
@@ -54,6 +67,19 @@ describe('check-require-cache', () => {
5467
},
5568
channelName: 'test_invoke'
5669
},
70+
{
71+
module: {
72+
name: 'test-trace-async-super',
73+
versionRange: '>=0.1',
74+
filePath: 'index.js'
75+
},
76+
functionQuery: {
77+
methodName: 'test',
78+
kind: 'Async',
79+
className: 'B'
80+
},
81+
channelName: 'test_invoke'
82+
},
5783
{
5884
module: {
5985
name: 'test-trace-callback',
@@ -66,6 +92,19 @@ describe('check-require-cache', () => {
6692
},
6793
channelName: 'test_invoke'
6894
},
95+
{
96+
module: {
97+
name: 'test-trace-callback-super',
98+
versionRange: '>=0.1',
99+
filePath: 'index.js'
100+
},
101+
functionQuery: {
102+
methodName: 'test',
103+
kind: 'Callback',
104+
className: 'B'
105+
},
106+
channelName: 'test_invoke'
107+
},
69108
{
70109
module: {
71110
name: 'test-trace-class-instance-method',
@@ -101,7 +140,7 @@ describe('check-require-cache', () => {
101140
})
102141

103142
it('should auto instrument sync functions', done => {
104-
const test = compile('test-trace-sync')
143+
const { test } = compile('test-trace-sync')
105144

106145
subs = {
107146
start () {
@@ -112,11 +151,26 @@ describe('check-require-cache', () => {
112151
ch = tracingChannel('orchestrion:test-trace-sync:test_invoke')
113152
ch.subscribe(subs)
114153

115-
test.test()
154+
test()
155+
})
156+
157+
it('should auto instrument sync functions with super', done => {
158+
const { test } = compile('test-trace-sync-super')
159+
160+
subs = {
161+
start () {
162+
done()
163+
}
164+
}
165+
166+
ch = tracingChannel('orchestrion:test-trace-sync-super:test_invoke')
167+
ch.subscribe(subs)
168+
169+
test(() => {})
116170
})
117171

118172
it('should auto instrument async functions', done => {
119-
const test = compile('test-trace-async')
173+
const { test } = compile('test-trace-async')
120174

121175
subs = {
122176
start () {
@@ -127,11 +181,26 @@ describe('check-require-cache', () => {
127181
ch = tracingChannel('orchestrion:test-trace-async:test_invoke')
128182
ch.subscribe(subs)
129183

130-
test.test()
184+
test()
185+
})
186+
187+
it('should auto instrument async functions using super', done => {
188+
const { test } = compile('test-trace-async-super')
189+
190+
subs = {
191+
start () {
192+
done()
193+
}
194+
}
195+
196+
ch = tracingChannel('orchestrion:test-trace-async-super:test_invoke')
197+
ch.subscribe(subs)
198+
199+
test(() => {})
131200
})
132201

133202
it('should auto instrument callback functions', done => {
134-
const test = compile('test-trace-callback')
203+
const { test } = compile('test-trace-callback')
135204

136205
subs = {
137206
start () {
@@ -142,7 +211,22 @@ describe('check-require-cache', () => {
142211
ch = tracingChannel('orchestrion:test-trace-callback:test_invoke')
143212
ch.subscribe(subs)
144213

145-
test.test(() => {})
214+
test(() => {})
215+
})
216+
217+
it('should auto instrument callback functions using super', done => {
218+
const { test } = compile('test-trace-callback-super')
219+
220+
subs = {
221+
start () {
222+
done()
223+
}
224+
}
225+
226+
ch = tracingChannel('orchestrion:test-trace-callback-super:test_invoke')
227+
ch.subscribe(subs)
228+
229+
test(() => {})
146230
})
147231

148232
it('should auto instrument class instance methods', done => {

packages/datadog-instrumentations/test/helpers/rewriter/node_modules/test-trace-async-super/index.js

Lines changed: 19 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/datadog-instrumentations/test/helpers/rewriter/node_modules/test-trace-async-super/package.json

Lines changed: 4 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/datadog-instrumentations/test/helpers/rewriter/node_modules/test-trace-callback-super/index.js

Lines changed: 19 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/datadog-instrumentations/test/helpers/rewriter/node_modules/test-trace-callback-super/package.json

Lines changed: 4 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/datadog-instrumentations/test/helpers/rewriter/node_modules/test-trace-sync-super/index.js

Lines changed: 18 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/datadog-instrumentations/test/helpers/rewriter/node_modules/test-trace-sync-super/package.json

Lines changed: 4 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)