Skip to content

Commit af50fef

Browse files
rochdevrochdev
authored
add rewriter support for generator functions (#7472)
1 parent 07c6d02 commit af50fef

File tree

12 files changed

+480
-41
lines changed

12 files changed

+480
-41
lines changed

‎packages/datadog-instrumentations/src/helpers/rewriter/index.js‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,15 @@ Orchestrion-JS that will need to be backported:
2929
is not a function. We'll see over time if something like this is needed to be
3030
backported or if it can be replaced by simpler queries.
3131
- Supports replacing methods of child class instances in the base constructor.
32+
- Supports tracing iterator (sync/async) returning functions (sync/async).
3233
*/
3334

3435
const { readFileSync } = require('fs')
3536
const { join } = require('path')
3637
const semifies = require('../../../../../vendor/dist/semifies')
3738
const log = require('../../../../dd-trace/src/log')
3839
const { getEnvironmentVariable } = require('../../../../dd-trace/src/config/helper')
39-
const transforms = require('./transforms')
40+
const { transform } = require('./transformer')
4041
const { generate, parse, traverse } = require('./compiler')
4142
const instrumentations = require('./instrumentations')
4243

@@ -61,19 +62,15 @@ function rewrite (content, filename, format) {
6162

6263
for (const inst of instrumentations) {
6364
const { astQuery, functionQuery = {}, module: { name, versionRange, filePath } } = inst
64-
const { kind } = functionQuery
65-
const operator = kind === 'Async' ? 'tracePromise' : kind === 'Callback' ? 'traceCallback' : 'traceSync'
66-
const transform = transforms[operator]
6765

6866
if (disabled.has(name)) continue
6967
if (!filename.endsWith(`${name}/${filePath}`)) continue
70-
if (!transform) continue
7168
if (!satisfies(filename, filePath, versionRange)) continue
7269

7370
ast ??= parse(content.toString(), { loc: true, ranges: true, module: format === 'module' })
7471

7572
const query = astQuery || fromFunctionQuery(functionQuery)
76-
const state = { ...inst, format, functionQuery, operator }
73+
const state = { ...inst, format, functionQuery }
7774

7875
traverse(ast, query, (...args) => transform(state, ...args))
7976
}

‎packages/datadog-instrumentations/src/helpers/rewriter/transformer.js‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
'use strict'
2+
3+
const transforms = require('./transforms')
4+
5+
function transform (state, ...args) {
6+
const operator = state.operator = getOperator(state)
7+
8+
transforms[operator](state, ...args)
9+
}
10+
11+
function getOperator ({ functionQuery: { kind } }) {
12+
switch (kind) {
13+
case 'Async': return 'tracePromise'
14+
case 'AsyncIterator': return 'traceAsyncIterator'
15+
case 'Callback': return 'traceCallback'
16+
case 'Iterator': return 'traceIterator'
17+
case 'Sync': return 'traceSync'
18+
}
19+
}
20+
21+
module.exports = { transform }

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

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

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

55
const tracingChannelPredicate = (node) => (
66
node.specifiers?.[0]?.local?.name === 'tr_ch_apm_tracingChannel' ||
@@ -35,7 +35,9 @@ const transforms = module.exports = {
3535
node.body.splice(index + 1, 0, parse(code).body[0])
3636
},
3737

38+
traceAsyncIterator: traceAny,
3839
traceCallback: traceAny,
40+
traceIterator: traceAny,
3941
tracePromise: traceAny,
4042
traceSync: traceAny,
4143
}
@@ -51,18 +53,25 @@ function traceAny (state, node, _parent, ancestry) {
5153
}
5254

5355
function traceFunction (state, node, program) {
54-
const { operator } = state
55-
5656
transforms.tracingChannelDeclaration(state, program)
5757

5858
node.body = wrap(state, {
59-
type: 'ArrowFunctionExpression',
59+
type: 'FunctionExpression',
6060
params: node.params,
6161
body: node.body,
62-
async: operator === 'tracePromise',
62+
async: node.async,
6363
expression: false,
64-
generator: false,
65-
})
64+
generator: node.generator,
65+
}, program)
66+
67+
// The original function no longer contains any calls to `await` or `yield` as
68+
// the function body is copied to the internal wrapped function, so we set
69+
// these to false to avoid altering the return value of the wrapper. The old
70+
// values are instead copied to the new AST node above.
71+
node.generator = false
72+
node.async = false
73+
74+
wrapSuper(state, node)
6675
}
6776

6877
function traceInstanceMethod (state, node, program) {
@@ -100,15 +109,19 @@ function traceInstanceMethod (state, node, program) {
100109
const fn = ctorBody[1].expression.right
101110

102111
fn.async = operator === 'tracePromise'
103-
fn.body = wrap(state, { type: 'Identifier', name: `__apm$${methodName}` })
112+
fn.body = wrap(state, { type: 'Identifier', name: `__apm$${methodName}` }, program)
113+
114+
wrapSuper(state, fn)
104115

105116
ctor.value.body.body.push(...ctorBody)
106117
}
107118

108-
function wrap (state, node) {
119+
function wrap (state, node, program) {
109120
const { channelName, operator } = state
110121

122+
if (operator === 'traceAsyncIterator') return wrapIterator(state, node, program)
111123
if (operator === 'traceCallback') return wrapCallback(state, node)
124+
if (operator === 'traceIterator') return wrapIterator(state, node, program)
112125

113126
const async = operator === 'tracePromise' ? 'async' : ''
114127
const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
@@ -133,6 +146,55 @@ function wrap (state, node) {
133146
return wrapper
134147
}
135148

149+
function wrapSuper (_state, node) {
150+
const members = new Set()
151+
152+
traverse(
153+
node.body,
154+
'[object.type=Super]',
155+
(node, parent) => {
156+
const { name } = node.property
157+
158+
let child
159+
160+
if (parent.callee) {
161+
// This is needed because for generator functions we have to move the
162+
// original function to a nested wrapped function, but we can't use an
163+
// arrow function because arrow function cannot be generator functions,
164+
// and `super` cannot be called from a nested function, so we have to
165+
// rewrite any `super` call to not use the keyword.
166+
const { expression } = parse(`__apm$super['${name}'].call(this)`).body[0]
167+
168+
parent.callee = child = expression.callee
169+
parent.arguments.unshift(...expression.arguments)
170+
} else {
171+
parent.expression = child = parse(`__apm$super['${name}']`).body[0]
172+
}
173+
174+
child.computed = parent.callee.computed
175+
child.optional = parent.callee.optional
176+
177+
members.add(name)
178+
}
179+
)
180+
181+
for (const name of members) {
182+
const member = parse(`
183+
class Wrapper {
184+
wrapper () {
185+
__apm$super['${name}'] = super['${name}']
186+
}
187+
}
188+
`).body[0].body.body[0].value.body.body[0]
189+
190+
node.body.body.unshift(member)
191+
}
192+
193+
if (members.size > 0) {
194+
node.body.body.unshift(parse('const __apm$super = {}').body[0])
195+
}
196+
}
197+
136198
function wrapCallback (state, node) {
137199
const { channelName, functionQuery: { index = -1 } } = state
138200
const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
@@ -194,3 +256,67 @@ function wrapCallback (state, node) {
194256

195257
return wrapper
196258
}
259+
260+
function wrapIterator (state, node, program) {
261+
const { channelName, operator } = state
262+
const baseChannel = channelName.replaceAll(':', '_')
263+
const channelVariable = 'tr_ch_apm$' + baseChannel
264+
const nextChannel = baseChannel + '_next'
265+
const traceMethod = operator === 'traceAsyncIterator' ? 'tracePromise' : 'traceSync'
266+
const traceNext = `tr_ch_apm$${nextChannel}.${traceMethod}`
267+
268+
transforms.tracingChannelDeclaration({ ...state, channelName: nextChannel }, program)
269+
270+
const wrapper = parse(`
271+
function wrapper () {
272+
const __apm$traced = () => {
273+
const __apm$wrapped = () => {};
274+
return __apm$wrapped.apply(this, arguments);
275+
};
276+
277+
if (!${channelVariable}.start.hasSubscribers) return __apm$traced();
278+
279+
{
280+
const wrap = iter => {
281+
const { next: iterNext, return: iterReturn, throw: iterThrow } = iter;
282+
283+
iter.next = (...args) => ${traceNext}(iterNext, ctx, iter, ...args);
284+
iter.return = (...args) => ${traceNext}(iterReturn, ctx, iter, ...args);
285+
iter.throw = (...args) => ${traceNext}(iterThrow, ctx, iter, ...args);
286+
287+
return iter;
288+
};
289+
const ctx = {
290+
arguments,
291+
self: this,
292+
moduleVersion: "1.0.0"
293+
};
294+
const iter = ${channelVariable}.traceSync(__apm$traced, ctx);
295+
296+
if (typeof iter.then !== 'function') return wrap(iter);
297+
298+
return iter.then(result => {
299+
ctx.result = result;
300+
301+
${channelVariable}.asyncStart.publish(ctx);
302+
${channelVariable}.asyncEnd.publish(ctx);
303+
304+
return wrap(result);
305+
}, err => {
306+
ctx.error = err;
307+
308+
${channelVariable}.error.publish(ctx);
309+
${channelVariable}.asyncStart.publish(ctx);
310+
${channelVariable}.asyncEnd.publish(ctx);
311+
312+
return Promise.reject(err);
313+
});
314+
};
315+
}
316+
`).body[0].body // Extract only block statement of function body.
317+
318+
// Replace the right-hand side assignment of `const __apm$wrapped = () => {}`.
319+
query(wrapper, '[id.name=__apm$wrapped]')[0].init = node
320+
321+
return wrapper
322+
}

0 commit comments

Comments
 (0)