Skip to content

Commit 6fbf2f9

Browse files
committed
fix: Properly assign async keyword only to wrapped functions that had async keyword on original
1 parent 0489576 commit 6fbf2f9

File tree

5 files changed

+520
-389
lines changed

5 files changed

+520
-389
lines changed

package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
"build:wrapper": "tsc index.ts --declaration --module commonjs",
2222
"build:inline-binary": "node inline-binary.js",
2323
"test": "vitest run",
24+
"test:update-snapshots": "vitest -u run",
2425
"test:watch": "vitest"
2526
},
2627
"devDependencies": {

src/instrumentation.rs

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ impl Instrumentation {
8585
define_channel
8686
}
8787

88-
fn insert_tracing(&mut self, body: &mut BlockStmt, params: &[Param]) {
88+
fn insert_tracing(&mut self, body: &mut BlockStmt, params: &[Param], is_async: bool) {
8989
self.count += 1;
9090

9191
let original_stmts = std::mem::take(&mut body.stmts);
@@ -99,7 +99,7 @@ impl Instrumentation {
9999

100100
let original_params: Vec<Pat> = params.iter().map(|p| p.pat.clone()).collect();
101101

102-
let wrapped_fn = new_fn(original_body, original_params);
102+
let wrapped_fn = new_fn(original_body, original_params, is_async);
103103

104104
let traced_body = BlockStmt {
105105
span: Span::default(),
@@ -110,7 +110,7 @@ impl Instrumentation {
110110
],
111111
};
112112

113-
let traced_fn = new_fn(traced_body, vec![]);
113+
let traced_fn = new_fn(traced_body, vec![], is_async);
114114

115115
let id_name = self.config.get_identifier_name();
116116
let ch_ident = ident!(format!("tr_ch_apm${}", &id_name));
@@ -208,7 +208,11 @@ impl Instrumentation {
208208
&& func_expr.function.body.is_some()
209209
{
210210
if let Some(body) = func_expr.function.body.as_mut() {
211-
self.insert_tracing(body, &func_expr.function.params);
211+
self.insert_tracing(
212+
body,
213+
&func_expr.function.params,
214+
func_expr.function.is_async,
215+
);
212216
}
213217
true
214218
} else {
@@ -246,7 +250,7 @@ impl Instrumentation {
246250
&& node.function.body.is_some()
247251
{
248252
if let Some(body) = node.function.body.as_mut() {
249-
self.insert_tracing(body, &node.function.params);
253+
self.insert_tracing(body, &node.function.params, node.function.is_async);
250254
}
251255
}
252256
true
@@ -302,7 +306,7 @@ impl Instrumentation {
302306
&& node.function.body.is_some()
303307
{
304308
if let Some(body) = node.function.body.as_mut() {
305-
self.insert_tracing(body, &node.function.params);
309+
self.insert_tracing(body, &node.function.params, node.function.is_async);
306310
}
307311
}
308312
true
@@ -335,7 +339,7 @@ impl Instrumentation {
335339
&& node.function.body.is_some()
336340
{
337341
if let Some(body) = node.function.body.as_mut() {
338-
self.insert_tracing(body, &node.function.params);
342+
self.insert_tracing(body, &node.function.params, node.function.is_async);
339343
}
340344
}
341345
false
@@ -373,6 +377,20 @@ impl Instrumentation {
373377
}
374378
}
375379

380+
#[must_use]
381+
pub fn new_fn(body: BlockStmt, params: Vec<Pat>, is_async: bool) -> ArrowExpr {
382+
ArrowExpr {
383+
params,
384+
body: Box::new(body.into()),
385+
is_async,
386+
is_generator: false,
387+
type_params: None,
388+
return_type: None,
389+
span: Span::default(),
390+
ctxt: SyntaxContext::empty(),
391+
}
392+
}
393+
376394
/// If the script starts with a "use strict" directive, we need to skip it when inserting there
377395
#[must_use]
378396
pub fn get_script_start_index(script: &Script) -> usize {

tests/wasm/__snapshots__/tests.test.mjs.snap

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
exports[`Orchestrion JS Transformer > should transform CommonJS module correctly 1`] = `
44
{
55
"code": "const { tracingChannel: tr_ch_apm_tracingChannel } = require("diagnostics_channel");
6+
const tr_ch_apm$up_asyncGet = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncGet");
7+
const tr_ch_apm$up_asyncFetch = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncFetch");
68
const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
79
const tr_ch_apm$up_constructor = tr_ch_apm_tracingChannel("orchestrion:one:up:constructor");
810
module.exports = class Up {
@@ -47,6 +49,32 @@ module.exports = class Up {
4749
moduleVersion: "1.0.0"
4850
});
4951
}
52+
async asyncFetch() {
53+
const __apm$original_args = arguments;
54+
const __apm$traced = async ()=>{
55+
const __apm$wrapped = async ()=>{};
56+
return __apm$wrapped.apply(null, __apm$original_args);
57+
};
58+
if (!tr_ch_apm$up_asyncFetch.hasSubscribers) return __apm$traced();
59+
return tr_ch_apm$up_asyncFetch.tracePromise(__apm$traced, {
60+
arguments,
61+
self: this,
62+
moduleVersion: "1.0.0"
63+
});
64+
}
65+
get() {
66+
const __apm$original_args = arguments;
67+
const __apm$traced = ()=>{
68+
const __apm$wrapped = ()=>{};
69+
return __apm$wrapped.apply(null, __apm$original_args);
70+
};
71+
if (!tr_ch_apm$up_asyncGet.hasSubscribers) return __apm$traced();
72+
return tr_ch_apm$up_asyncGet.tracePromise(__apm$traced, {
73+
arguments,
74+
self: this,
75+
moduleVersion: "1.0.0"
76+
});
77+
}
5078
};
5179
",
5280
"map": undefined,
@@ -56,6 +84,8 @@ module.exports = class Up {
5684
exports[`Orchestrion JS Transformer > should transform ESM module correctly 1`] = `
5785
{
5886
"code": "import { tracingChannel as tr_ch_apm_tracingChannel } from "diagnostics_channel";
87+
const tr_ch_apm$up_asyncGet = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncGet");
88+
const tr_ch_apm$up_asyncFetch = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncFetch");
5989
const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
6090
const tr_ch_apm$up_constructor = tr_ch_apm_tracingChannel("orchestrion:one:up:constructor");
6191
export class Up {
@@ -100,6 +130,32 @@ export class Up {
100130
moduleVersion: "1.0.0"
101131
});
102132
}
133+
async asyncFetch() {
134+
const __apm$original_args = arguments;
135+
const __apm$traced = async ()=>{
136+
const __apm$wrapped = async ()=>{};
137+
return __apm$wrapped.apply(null, __apm$original_args);
138+
};
139+
if (!tr_ch_apm$up_asyncFetch.hasSubscribers) return __apm$traced();
140+
return tr_ch_apm$up_asyncFetch.tracePromise(__apm$traced, {
141+
arguments,
142+
self: this,
143+
moduleVersion: "1.0.0"
144+
});
145+
}
146+
get() {
147+
const __apm$original_args = arguments;
148+
const __apm$traced = ()=>{
149+
const __apm$wrapped = ()=>{};
150+
return __apm$wrapped.apply(null, __apm$original_args);
151+
};
152+
if (!tr_ch_apm$up_asyncGet.hasSubscribers) return __apm$traced();
153+
return tr_ch_apm$up_asyncGet.tracePromise(__apm$traced, {
154+
arguments,
155+
self: this,
156+
moduleVersion: "1.0.0"
157+
});
158+
}
103159
}
104160
",
105161
"map": undefined,
@@ -109,6 +165,8 @@ export class Up {
109165
exports[`Orchestrion JS Transformer > should transform TypeScript with source map correctly 1`] = `
110166
{
111167
"code": "import { tracingChannel as tr_ch_apm_tracingChannel } from "diagnostics_channel";
168+
const tr_ch_apm$up_asyncGet = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncGet");
169+
const tr_ch_apm$up_asyncFetch = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncFetch");
112170
const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
113171
const tr_ch_apm$up_constructor = tr_ch_apm_tracingChannel("orchestrion:one:up:constructor");
114172
export class Up {
@@ -153,8 +211,34 @@ export class Up {
153211
moduleVersion: "1.0.0"
154212
});
155213
}
214+
async asyncFetch() {
215+
const __apm$original_args = arguments;
216+
const __apm$traced = async ()=>{
217+
const __apm$wrapped = async ()=>{};
218+
return __apm$wrapped.apply(null, __apm$original_args);
219+
};
220+
if (!tr_ch_apm$up_asyncFetch.hasSubscribers) return __apm$traced();
221+
return tr_ch_apm$up_asyncFetch.tracePromise(__apm$traced, {
222+
arguments,
223+
self: this,
224+
moduleVersion: "1.0.0"
225+
});
226+
}
227+
get() {
228+
const __apm$original_args = arguments;
229+
const __apm$traced = ()=>{
230+
const __apm$wrapped = ()=>{};
231+
return __apm$wrapped.apply(null, __apm$original_args);
232+
};
233+
if (!tr_ch_apm$up_asyncGet.hasSubscribers) return __apm$traced();
234+
return tr_ch_apm$up_asyncGet.tracePromise(__apm$traced, {
235+
arguments,
236+
self: this,
237+
moduleVersion: "1.0.0"
238+
});
239+
}
156240
}
157241
",
158-
"map": "{"version":3,"file":"module.js","sources":["module.ts"],"sourceRoot":"","names":[],"mappings":";;;AAEA,MAAM,CAAA,MAAO,EAAE;IACX,aAAA;;;;;;;;;YACI,OAAO,CAAC,GAAG,CAAC,aAAa,CAAC,CAAC;;;;;;;;;;;;;;;;IAC/B,CAAC;IACD,KAAK,IAAS,EAAA;;;mCAAR;gBACF,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,CAAC;;;;;;;;;;IACzB,CAAC;CACJ"}",
242+
"map": "{"version":3,"file":"module.js","sources":["module.ts"],"sourceRoot":"","names":[],"mappings":";;;;;AAEA,MAAM,CAAA,MAAO,EAAE;IACX,aAAA;;;;;;;;;YACI,OAAO,CAAC,GAAG,CAAC,aAAa,CAAC,CAAC;;;;;;;;;;;;;;;;IAC/B,CAAC;IACD,KAAK,IAAS,EAAA;;;mCAAR;gBACF,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,CAAC;;;;;;;;;;IACzB,CAAC;IACD,KAAK,CAAC,UAAU,GAAA;;;;;;;;;;;;IAAU,CAAC;IAC3B,GAAG,GAAA;;;;;;;;;;;;IAAU,CAAC;CACjB"}",
159243
}
160244
`;

tests/wasm/tests.test.mjs

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,24 @@ describe('Orchestrion JS Transformer', () => {
1919
kind: "Sync",
2020
},
2121
},
22+
{
23+
channelName: "up:asyncFetch",
24+
module: { name: "one", versionRange: ">=1", filePath: "index.js" },
25+
functionQuery: {
26+
className: "Up",
27+
methodName: "asyncFetch",
28+
kind: "Async",
29+
},
30+
},
31+
{
32+
channelName: "up:asyncGet",
33+
module: { name: "one", versionRange: ">=1", filePath: "index.js" },
34+
functionQuery: {
35+
className: "Up",
36+
methodName: "get",
37+
kind: "Async",
38+
},
39+
},
2240
]);
2341

2442
const matchedTransforms = instrumentor.getTransformer(
@@ -40,6 +58,9 @@ describe('Orchestrion JS Transformer', () => {
4058
fetch() {
4159
console.log('fetch')
4260
}
61+
62+
async asyncFetch() {}
63+
get() {}
4364
}`;
4465

4566
const output = matchedTransforms.transform(originalEsm, 'esm');
@@ -55,13 +76,16 @@ describe('Orchestrion JS Transformer', () => {
5576
fetch() {
5677
console.log('fetch')
5778
}
79+
80+
async asyncFetch() {}
81+
get() {}
5882
}
5983
6084
`;
6185
const outputCjs = matchedTransforms.transform(originalCjs, 'cjs');
6286
expect(outputCjs).toMatchSnapshot();
6387
});
64-
88+
6589
test('should transform TypeScript with source map correctly', async () => {
6690
const originalTypescript = `type Url = { href: string };
6791
@@ -72,6 +96,8 @@ export class Up {
7296
fetch(url: Url): void {
7397
console.log('fetch');
7498
}
99+
async asyncFetch(): void {}
100+
get(): void {}
75101
}`;
76102

77103
const { outputText: outputJavaScript, sourceMapText: originalTypescriptSourceMap } = tsc.transpileModule(originalTypescript, {
@@ -94,7 +120,7 @@ export class Up {
94120

95121
const originalPosition = sourceMapConsumer.originalPositionFor({
96122
// This is the position of the fetch function in the transformed JavaScript
97-
line: 31,
123+
line: 33,
98124
column: 4,
99125
});
100126

@@ -114,10 +140,12 @@ export class Up {
114140
fetch() {
115141
console.log('fetch')
116142
}
143+
async asyncFetch() {}
144+
get() {}
117145
}`;
118146

119147
expect(() => {
120148
matchedTransforms.transform(noMatchSource, 'unknown');
121-
}).toThrow('Failed to find injection points for: ["constructor", "fetch"]');
149+
}).toThrow('Failed to find injection points for: ["constructor", "fetch", "asyncFetch", "get"]');
122150
});
123151
});

0 commit comments

Comments
 (0)