Skip to content

fix: Properly assign async keyword only to wrapped functions that had async keyword on original #50

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 26, 2025
Merged

fix: Properly assign async keyword only to wrapped functions that had async keyword on original #50

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
"build:wrapper": "tsc index.ts --declaration --module commonjs",
"build:inline-binary": "node inline-binary.js",
"test": "vitest run",
"test:update-snapshots": "vitest -u run",
"test:watch": "vitest"
},
"devDependencies": {
Expand Down
26 changes: 13 additions & 13 deletions src/instrumentation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ impl Instrumentation {
define_channel
}

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

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

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

let wrapped_fn = new_fn(original_body, original_params);
let wrapped_fn = new_fn(original_body, original_params, is_async);

let traced_body = BlockStmt {
span: Span::default(),
Expand All @@ -110,7 +110,7 @@ impl Instrumentation {
],
};

let traced_fn = new_fn(traced_body, vec![]);
let traced_fn = new_fn(traced_body, vec![], is_async);

let id_name = self.config.get_identifier_name();
let ch_ident = ident!(format!("tr_ch_apm${}", &id_name));
Expand Down Expand Up @@ -208,7 +208,11 @@ impl Instrumentation {
&& func_expr.function.body.is_some()
{
if let Some(body) = func_expr.function.body.as_mut() {
self.insert_tracing(body, &func_expr.function.params);
self.insert_tracing(
body,
&func_expr.function.params,
func_expr.function.is_async,
);
}
true
} else {
Expand Down Expand Up @@ -246,7 +250,7 @@ impl Instrumentation {
&& node.function.body.is_some()
{
if let Some(body) = node.function.body.as_mut() {
self.insert_tracing(body, &node.function.params);
self.insert_tracing(body, &node.function.params, node.function.is_async);
}
}
true
Expand Down Expand Up @@ -302,7 +306,7 @@ impl Instrumentation {
&& node.function.body.is_some()
{
if let Some(body) = node.function.body.as_mut() {
self.insert_tracing(body, &node.function.params);
self.insert_tracing(body, &node.function.params, node.function.is_async);
}
}
true
Expand Down Expand Up @@ -335,7 +339,7 @@ impl Instrumentation {
&& node.function.body.is_some()
{
if let Some(body) = node.function.body.as_mut() {
self.insert_tracing(body, &node.function.params);
self.insert_tracing(body, &node.function.params, node.function.is_async);
}
}
false
Expand Down Expand Up @@ -387,15 +391,11 @@ pub fn get_script_start_index(script: &Script) -> usize {
}

#[must_use]
pub fn new_fn(body: BlockStmt, params: Vec<Pat>) -> ArrowExpr {
pub fn new_fn(body: BlockStmt, params: Vec<Pat>, is_async: bool) -> ArrowExpr {
ArrowExpr {
params,
body: Box::new(body.into()),
// if we set this to `true` and it's an async function,
// it will wrap the original promise in a new native one
// which may not be semantically correct in cases where custom
// promises are returned
is_async: false,
is_async,
is_generator: false,
type_params: None,
return_type: None,
Expand Down
86 changes: 85 additions & 1 deletion tests/wasm/__snapshots__/tests.test.mjs.snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
exports[`Orchestrion JS Transformer > should transform CommonJS module correctly 1`] = `
{
"code": "const { tracingChannel: tr_ch_apm_tracingChannel } = require("diagnostics_channel");
const tr_ch_apm$up_asyncGet = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncGet");
const tr_ch_apm$up_asyncFetch = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncFetch");
const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
const tr_ch_apm$up_constructor = tr_ch_apm_tracingChannel("orchestrion:one:up:constructor");
module.exports = class Up {
Expand Down Expand Up @@ -47,6 +49,32 @@ module.exports = class Up {
moduleVersion: "1.0.0"
});
}
async asyncFetch() {
const __apm$original_args = arguments;
const __apm$traced = async ()=>{
const __apm$wrapped = async ()=>{};
return __apm$wrapped.apply(null, __apm$original_args);
};
if (!tr_ch_apm$up_asyncFetch.hasSubscribers) return __apm$traced();
return tr_ch_apm$up_asyncFetch.tracePromise(__apm$traced, {
arguments,
self: this,
moduleVersion: "1.0.0"
});
}
get() {
const __apm$original_args = arguments;
const __apm$traced = ()=>{
const __apm$wrapped = ()=>{};
return __apm$wrapped.apply(null, __apm$original_args);
};
if (!tr_ch_apm$up_asyncGet.hasSubscribers) return __apm$traced();
return tr_ch_apm$up_asyncGet.tracePromise(__apm$traced, {
arguments,
self: this,
moduleVersion: "1.0.0"
});
}
};
",
"map": undefined,
Expand All @@ -56,6 +84,8 @@ module.exports = class Up {
exports[`Orchestrion JS Transformer > should transform ESM module correctly 1`] = `
{
"code": "import { tracingChannel as tr_ch_apm_tracingChannel } from "diagnostics_channel";
const tr_ch_apm$up_asyncGet = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncGet");
const tr_ch_apm$up_asyncFetch = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncFetch");
const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
const tr_ch_apm$up_constructor = tr_ch_apm_tracingChannel("orchestrion:one:up:constructor");
export class Up {
Expand Down Expand Up @@ -100,6 +130,32 @@ export class Up {
moduleVersion: "1.0.0"
});
}
async asyncFetch() {
const __apm$original_args = arguments;
const __apm$traced = async ()=>{
const __apm$wrapped = async ()=>{};
return __apm$wrapped.apply(null, __apm$original_args);
};
if (!tr_ch_apm$up_asyncFetch.hasSubscribers) return __apm$traced();
return tr_ch_apm$up_asyncFetch.tracePromise(__apm$traced, {
arguments,
self: this,
moduleVersion: "1.0.0"
});
}
get() {
const __apm$original_args = arguments;
const __apm$traced = ()=>{
const __apm$wrapped = ()=>{};
return __apm$wrapped.apply(null, __apm$original_args);
};
if (!tr_ch_apm$up_asyncGet.hasSubscribers) return __apm$traced();
return tr_ch_apm$up_asyncGet.tracePromise(__apm$traced, {
arguments,
self: this,
moduleVersion: "1.0.0"
});
}
}
",
"map": undefined,
Expand All @@ -109,6 +165,8 @@ export class Up {
exports[`Orchestrion JS Transformer > should transform TypeScript with source map correctly 1`] = `
{
"code": "import { tracingChannel as tr_ch_apm_tracingChannel } from "diagnostics_channel";
const tr_ch_apm$up_asyncGet = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncGet");
const tr_ch_apm$up_asyncFetch = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncFetch");
const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
const tr_ch_apm$up_constructor = tr_ch_apm_tracingChannel("orchestrion:one:up:constructor");
export class Up {
Expand Down Expand Up @@ -153,8 +211,34 @@ export class Up {
moduleVersion: "1.0.0"
});
}
async asyncFetch() {
const __apm$original_args = arguments;
const __apm$traced = async ()=>{
const __apm$wrapped = async ()=>{};
return __apm$wrapped.apply(null, __apm$original_args);
};
if (!tr_ch_apm$up_asyncFetch.hasSubscribers) return __apm$traced();
return tr_ch_apm$up_asyncFetch.tracePromise(__apm$traced, {
arguments,
self: this,
moduleVersion: "1.0.0"
});
}
get() {
const __apm$original_args = arguments;
const __apm$traced = ()=>{
const __apm$wrapped = ()=>{};
return __apm$wrapped.apply(null, __apm$original_args);
};
if (!tr_ch_apm$up_asyncGet.hasSubscribers) return __apm$traced();
return tr_ch_apm$up_asyncGet.tracePromise(__apm$traced, {
arguments,
self: this,
moduleVersion: "1.0.0"
});
}
}
",
"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"}",
"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"}",
}
`;
34 changes: 31 additions & 3 deletions tests/wasm/tests.test.mjs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,24 @@ describe('Orchestrion JS Transformer', () => {
kind: "Sync",
},
},
{
channelName: "up:asyncFetch",
module: { name: "one", versionRange: ">=1", filePath: "index.js" },
functionQuery: {
className: "Up",
methodName: "asyncFetch",
kind: "Async",
},
},
{
channelName: "up:asyncGet",
module: { name: "one", versionRange: ">=1", filePath: "index.js" },
functionQuery: {
className: "Up",
methodName: "get",
kind: "Async",
},
},
]);

const matchedTransforms = instrumentor.getTransformer(
Expand All @@ -40,6 +58,9 @@ describe('Orchestrion JS Transformer', () => {
fetch() {
console.log('fetch')
}

async asyncFetch() {}
get() {}
}`;

const output = matchedTransforms.transform(originalEsm, 'esm');
Expand All @@ -55,13 +76,16 @@ describe('Orchestrion JS Transformer', () => {
fetch() {
console.log('fetch')
}

async asyncFetch() {}
get() {}
}

`;
const outputCjs = matchedTransforms.transform(originalCjs, 'cjs');
expect(outputCjs).toMatchSnapshot();
});

test('should transform TypeScript with source map correctly', async () => {
const originalTypescript = `type Url = { href: string };

Expand All @@ -72,6 +96,8 @@ export class Up {
fetch(url: Url): void {
console.log('fetch');
}
async asyncFetch(): void {}
get(): void {}
}`;

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

const originalPosition = sourceMapConsumer.originalPositionFor({
// This is the position of the fetch function in the transformed JavaScript
line: 31,
line: 33,
column: 4,
});

Expand All @@ -114,10 +140,12 @@ export class Up {
fetch() {
console.log('fetch')
}
async asyncFetch() {}
get() {}
}`;

expect(() => {
matchedTransforms.transform(noMatchSource, 'unknown');
}).toThrow('Failed to find injection points for: ["constructor", "fetch"]');
}).toThrow('Failed to find injection points for: ["constructor", "fetch", "asyncFetch", "get"]');
});
});
Loading