From 8b7d6073ea475dc63df5a2e8f02f002815341679 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 27 Feb 2026 13:19:53 -0500 Subject: [PATCH] feat: async generator support --- src/function_query.rs | 5 +- src/instrumentation.rs | 137 ++++++++++++++---- tests/async_generator_class_method_cjs/mod.js | 11 ++ tests/async_generator_class_method_cjs/mod.rs | 15 ++ .../async_generator_class_method_cjs/test.js | 19 +++ tests/async_generator_decl_cjs/mod.js | 9 ++ tests/async_generator_decl_cjs/mod.rs | 15 ++ tests/async_generator_decl_cjs/test.js | 18 +++ tests/async_generator_decl_mjs/mod.mjs | 10 ++ tests/async_generator_decl_mjs/mod.rs | 15 ++ tests/async_generator_decl_mjs/test.mjs | 16 ++ tests/async_generator_expr_cjs/mod.js | 9 ++ tests/async_generator_expr_cjs/mod.rs | 15 ++ tests/async_generator_expr_cjs/test.js | 18 +++ .../async_generator_object_method_cjs/mod.js | 11 ++ .../async_generator_object_method_cjs/mod.rs | 15 ++ .../async_generator_object_method_cjs/test.js | 18 +++ .../async_generator_private_method_cjs/mod.js | 15 ++ .../async_generator_private_method_cjs/mod.rs | 15 ++ .../test.js | 19 +++ tests/instrumentor_test.rs | 6 + 21 files changed, 384 insertions(+), 27 deletions(-) create mode 100644 tests/async_generator_class_method_cjs/mod.js create mode 100644 tests/async_generator_class_method_cjs/mod.rs create mode 100644 tests/async_generator_class_method_cjs/test.js create mode 100644 tests/async_generator_decl_cjs/mod.js create mode 100644 tests/async_generator_decl_cjs/mod.rs create mode 100644 tests/async_generator_decl_cjs/test.js create mode 100644 tests/async_generator_decl_mjs/mod.mjs create mode 100644 tests/async_generator_decl_mjs/mod.rs create mode 100644 tests/async_generator_decl_mjs/test.mjs create mode 100644 tests/async_generator_expr_cjs/mod.js create mode 100644 tests/async_generator_expr_cjs/mod.rs create mode 100644 tests/async_generator_expr_cjs/test.js create mode 100644 tests/async_generator_object_method_cjs/mod.js create mode 100644 tests/async_generator_object_method_cjs/mod.rs create mode 100644 tests/async_generator_object_method_cjs/test.js create mode 100644 tests/async_generator_private_method_cjs/mod.js create mode 100644 tests/async_generator_private_method_cjs/mod.rs create mode 100644 tests/async_generator_private_method_cjs/test.js diff --git a/src/function_query.rs b/src/function_query.rs index 1c7904b..3280f91 100644 --- a/src/function_query.rs +++ b/src/function_query.rs @@ -20,18 +20,19 @@ pub(crate) enum FunctionType { pub enum FunctionKind { Sync, Async, + AsyncGenerator, } impl FunctionKind { #[must_use] pub fn is_async(&self) -> bool { - matches!(self, FunctionKind::Async) + matches!(self, FunctionKind::Async | FunctionKind::AsyncGenerator) } #[must_use] pub fn tracing_operator(&self) -> &'static str { match self { - FunctionKind::Sync => "traceSync", + FunctionKind::Sync | FunctionKind::AsyncGenerator => "traceSync", FunctionKind::Async => "tracePromise", } } diff --git a/src/instrumentation.rs b/src/instrumentation.rs index 58d5917..3ca8481 100644 --- a/src/instrumentation.rs +++ b/src/instrumentation.rs @@ -9,8 +9,9 @@ use swc_core::common::{Span, SyntaxContext}; use swc_core::ecma::{ ast::{ ArrowExpr, AssignExpr, AssignTarget, BlockStmt, ClassDecl, ClassExpr, ClassMethod, - Constructor, Expr, FnDecl, FnExpr, Ident, Lit, MemberProp, MethodProp, Module, ModuleItem, - Param, Pat, PrivateMethod, PropName, Script, SimpleAssignTarget, Stmt, Str, VarDecl, + Constructor, Expr, FnDecl, FnExpr, Function, Ident, Lit, MemberProp, MethodProp, Module, + ModuleItem, Param, Pat, PrivateMethod, PropName, Script, SimpleAssignTarget, Stmt, Str, + VarDecl, }, atoms::Atom, }; @@ -108,7 +109,13 @@ impl Instrumentation { define_channel } - fn insert_tracing(&mut self, body: &mut BlockStmt, params: &[Param], is_async: bool) { + fn insert_tracing( + &mut self, + body: &mut BlockStmt, + params: &[Param], + is_async: bool, + is_generator: bool, + ) { self.count += 1; let original_stmts = std::mem::take(&mut body.stmts); @@ -122,30 +129,40 @@ impl Instrumentation { let original_params: Vec = params.iter().map(|p| p.pat.clone()).collect(); - let wrapped_fn = new_fn(original_body, original_params, is_async); + let wrapped_fn = new_fn(original_body, original_params, is_async, is_generator); let traced_body = BlockStmt { span: Span::default(), ctxt: SyntaxContext::empty(), stmts: vec![ - quote!("const __apm$wrapped = $wrapped;" as Stmt, wrapped: Expr = wrapped_fn.into()), + quote!("const __apm$wrapped = $wrapped;" as Stmt, wrapped: Expr = wrapped_fn), quote!("return __apm$wrapped.apply(null, __apm$original_args);" as Stmt), ], }; - let traced_fn = new_fn(traced_body, vec![], is_async); + // When is_generator, __apm$traced is a regular (non-async) arrow since generator + // creation is synchronous. Otherwise, it mirrors the original function's async-ness. + let traced_fn = new_fn( + traced_body, + vec![], + if is_generator { false } else { is_async }, + false, + ); let id_name = self.config.get_identifier_name(); let ch_ident = ident!(format!("tr_ch_apm${}", &id_name)); - let trace_ident = ident!(format!( - "tr_ch_apm${}.{}", - &id_name, + + // For generators, always use traceSync regardless of configured kind + let tracing_op = if is_generator { + "traceSync" + } else { self.config.function_query.kind().tracing_operator() - )); + }; + let trace_ident = ident!(format!("tr_ch_apm${}.{}", &id_name, tracing_op)); body.stmts = vec![ quote!("const __apm$original_args = arguments" as Stmt), - quote!("const __apm$traced = $traced;" as Stmt, traced: Expr = traced_fn.into()), + quote!("const __apm$traced = $traced;" as Stmt, traced: Expr = traced_fn), quote!( "if (!$ch.hasSubscribers) return __apm$traced();" as Stmt, ch = ch_ident @@ -230,13 +247,19 @@ impl Instrumentation { .matches_expr(&mut self.count, name.as_ref()) && func_expr.function.body.is_some() { + let is_generator = func_expr.function.is_generator; if let Some(body) = func_expr.function.body.as_mut() { self.insert_tracing( body, &func_expr.function.params, func_expr.function.is_async, + is_generator, ); } + if is_generator { + func_expr.function.is_async = false; + func_expr.function.is_generator = false; + } true } else { false @@ -272,8 +295,18 @@ impl Instrumentation { .matches_decl(node, &mut self.count) && node.function.body.is_some() { + let is_generator = node.function.is_generator; if let Some(body) = node.function.body.as_mut() { - self.insert_tracing(body, &node.function.params, node.function.is_async); + self.insert_tracing( + body, + &node.function.params, + node.function.is_async, + is_generator, + ); + } + if is_generator { + node.function.is_async = false; + node.function.is_generator = false; } } true @@ -328,8 +361,18 @@ impl Instrumentation { .matches_method(&mut self.count, name.as_ref()) && node.function.body.is_some() { + let is_generator = node.function.is_generator; if let Some(body) = node.function.body.as_mut() { - self.insert_tracing(body, &node.function.params, node.function.is_async); + self.insert_tracing( + body, + &node.function.params, + node.function.is_async, + is_generator, + ); + } + if is_generator { + node.function.is_async = false; + node.function.is_generator = false; } } true @@ -349,8 +392,18 @@ impl Instrumentation { .matches_private_method(&mut self.count, name.as_ref()) && node.function.body.is_some() { + let is_generator = node.function.is_generator; if let Some(body) = node.function.body.as_mut() { - self.insert_tracing(body, &node.function.params, node.function.is_async); + self.insert_tracing( + body, + &node.function.params, + node.function.is_async, + is_generator, + ); + } + if is_generator { + node.function.is_async = false; + node.function.is_generator = false; } } true @@ -382,8 +435,18 @@ impl Instrumentation { .matches_method(&mut self.count, name.as_ref()) && node.function.body.is_some() { + let is_generator = node.function.is_generator; if let Some(body) = node.function.body.as_mut() { - self.insert_tracing(body, &node.function.params, node.function.is_async); + self.insert_tracing( + body, + &node.function.params, + node.function.is_async, + is_generator, + ); + } + if is_generator { + node.function.is_async = false; + node.function.is_generator = false; } } false @@ -435,15 +498,39 @@ pub fn get_script_start_index(script: &Script) -> usize { } #[must_use] -pub fn new_fn(body: BlockStmt, params: Vec, is_async: bool) -> ArrowExpr { - ArrowExpr { - params, - body: Box::new(body.into()), - is_async, - is_generator: false, - type_params: None, - return_type: None, - span: Span::default(), - ctxt: SyntaxContext::empty(), +pub fn new_fn(body: BlockStmt, params: Vec, is_async: bool, is_generator: bool) -> Expr { + if is_generator { + Expr::Fn(FnExpr { + ident: None, + function: Box::new(Function { + params: params + .into_iter() + .map(|p| Param { + span: Span::default(), + decorators: vec![], + pat: p, + }) + .collect(), + decorators: vec![], + span: Span::default(), + ctxt: SyntaxContext::empty(), + body: Some(body), + is_generator: true, + is_async, + type_params: None, + return_type: None, + }), + }) + } else { + Expr::Arrow(ArrowExpr { + params, + body: Box::new(body.into()), + is_async, + is_generator: false, + type_params: None, + return_type: None, + span: Span::default(), + ctxt: SyntaxContext::empty(), + }) } } diff --git a/tests/async_generator_class_method_cjs/mod.js b/tests/async_generator_class_method_cjs/mod.js new file mode 100644 index 0000000..96219a6 --- /dev/null +++ b/tests/async_generator_class_method_cjs/mod.js @@ -0,0 +1,11 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +class Streamer { + async *generate(n) { + for (let i = 0; i < n; i++) yield i; + } +} + +module.exports = Streamer; diff --git a/tests/async_generator_class_method_cjs/mod.rs b/tests/async_generator_class_method_cjs/mod.rs new file mode 100644 index 0000000..a1150e0 --- /dev/null +++ b/tests/async_generator_class_method_cjs/mod.rs @@ -0,0 +1,15 @@ +use crate::common::*; +use orchestrion_js::*; + +#[test] +fn async_generator_class_method_cjs() { + transpile_and_test( + file!(), + false, + Config::new_single(InstrumentationConfig::new( + "Streamer:generate", + test_module_matcher(), + FunctionQuery::class_method("Streamer", "generate", FunctionKind::AsyncGenerator), + )), + ); +} diff --git a/tests/async_generator_class_method_cjs/test.js b/tests/async_generator_class_method_cjs/test.js new file mode 100644 index 0000000..8990a74 --- /dev/null +++ b/tests/async_generator_class_method_cjs/test.js @@ -0,0 +1,19 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +const Streamer = require('./instrumented.js'); +const { assert, getContext } = require('../common/preamble.js'); +const context = getContext('orchestrion:undici:Streamer:generate'); +(async () => { + const streamer = new Streamer; + const values = []; + for await (const val of streamer.generate(3)) { + values.push(val); + } + assert.deepStrictEqual(values, [0, 1, 2]); + assert.deepStrictEqual(context, { + start: true, + end: true, + }); +})(); diff --git a/tests/async_generator_decl_cjs/mod.js b/tests/async_generator_decl_cjs/mod.js new file mode 100644 index 0000000..45ba85e --- /dev/null +++ b/tests/async_generator_decl_cjs/mod.js @@ -0,0 +1,9 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +async function* generate(n) { + for (let i = 0; i < n; i++) yield i; +} + +module.exports = { generate }; diff --git a/tests/async_generator_decl_cjs/mod.rs b/tests/async_generator_decl_cjs/mod.rs new file mode 100644 index 0000000..d6a28fb --- /dev/null +++ b/tests/async_generator_decl_cjs/mod.rs @@ -0,0 +1,15 @@ +use crate::common::*; +use orchestrion_js::*; + +#[test] +fn async_generator_decl_cjs() { + transpile_and_test( + file!(), + false, + Config::new_single(InstrumentationConfig::new( + "generate_decl", + test_module_matcher(), + FunctionQuery::function_declaration("generate", FunctionKind::AsyncGenerator), + )), + ); +} diff --git a/tests/async_generator_decl_cjs/test.js b/tests/async_generator_decl_cjs/test.js new file mode 100644 index 0000000..46351c5 --- /dev/null +++ b/tests/async_generator_decl_cjs/test.js @@ -0,0 +1,18 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +const { generate } = require('./instrumented.js'); +const { assert, getContext } = require('../common/preamble.js'); +const context = getContext('orchestrion:undici:generate_decl'); +(async () => { + const values = []; + for await (const val of generate(3)) { + values.push(val); + } + assert.deepStrictEqual(values, [0, 1, 2]); + assert.deepStrictEqual(context, { + start: true, + end: true, + }); +})(); diff --git a/tests/async_generator_decl_mjs/mod.mjs b/tests/async_generator_decl_mjs/mod.mjs new file mode 100644 index 0000000..558053d --- /dev/null +++ b/tests/async_generator_decl_mjs/mod.mjs @@ -0,0 +1,10 @@ + +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +async function* generate(n) { + for (let i = 0; i < n; i++) yield i; +} + +export { generate }; diff --git a/tests/async_generator_decl_mjs/mod.rs b/tests/async_generator_decl_mjs/mod.rs new file mode 100644 index 0000000..637597d --- /dev/null +++ b/tests/async_generator_decl_mjs/mod.rs @@ -0,0 +1,15 @@ +use crate::common::*; +use orchestrion_js::*; + +#[test] +fn async_generator_decl_mjs() { + transpile_and_test( + file!(), + true, + Config::new_single(InstrumentationConfig::new( + "generate_decl", + test_module_matcher(), + FunctionQuery::function_declaration("generate", FunctionKind::AsyncGenerator), + )), + ); +} diff --git a/tests/async_generator_decl_mjs/test.mjs b/tests/async_generator_decl_mjs/test.mjs new file mode 100644 index 0000000..eee9741 --- /dev/null +++ b/tests/async_generator_decl_mjs/test.mjs @@ -0,0 +1,16 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +import { generate } from './instrumented.mjs'; +import { assert, getContext } from '../common/preamble.js'; +const context = getContext('orchestrion:undici:generate_decl'); +const values = []; +for await (const val of generate(3)) { + values.push(val); +} +assert.deepStrictEqual(values, [0, 1, 2]); +assert.deepStrictEqual(context, { + start: true, + end: true, +}); diff --git a/tests/async_generator_expr_cjs/mod.js b/tests/async_generator_expr_cjs/mod.js new file mode 100644 index 0000000..96e035c --- /dev/null +++ b/tests/async_generator_expr_cjs/mod.js @@ -0,0 +1,9 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +'use strict'; + +exports.generate = async function* (n) { + for (let i = 0; i < n; i++) yield i; +} diff --git a/tests/async_generator_expr_cjs/mod.rs b/tests/async_generator_expr_cjs/mod.rs new file mode 100644 index 0000000..72ee3bd --- /dev/null +++ b/tests/async_generator_expr_cjs/mod.rs @@ -0,0 +1,15 @@ +use crate::common::*; +use orchestrion_js::*; + +#[test] +fn async_generator_expr_cjs() { + transpile_and_test( + file!(), + false, + Config::new_single(InstrumentationConfig::new( + "generate_expr", + test_module_matcher(), + FunctionQuery::function_expression("generate", FunctionKind::AsyncGenerator), + )), + ); +} diff --git a/tests/async_generator_expr_cjs/test.js b/tests/async_generator_expr_cjs/test.js new file mode 100644 index 0000000..f8b562b --- /dev/null +++ b/tests/async_generator_expr_cjs/test.js @@ -0,0 +1,18 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +const { generate } = require('./instrumented.js'); +const { assert, getContext } = require('../common/preamble.js'); +const context = getContext('orchestrion:undici:generate_expr'); +(async () => { + const values = []; + for await (const val of generate(3)) { + values.push(val); + } + assert.deepStrictEqual(values, [0, 1, 2]); + assert.deepStrictEqual(context, { + start: true, + end: true, + }); +})(); diff --git a/tests/async_generator_object_method_cjs/mod.js b/tests/async_generator_object_method_cjs/mod.js new file mode 100644 index 0000000..67cb7f4 --- /dev/null +++ b/tests/async_generator_object_method_cjs/mod.js @@ -0,0 +1,11 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +const streamer = { + async *generate(n) { + for (let i = 0; i < n; i++) yield i; + } +} + +module.exports = streamer; diff --git a/tests/async_generator_object_method_cjs/mod.rs b/tests/async_generator_object_method_cjs/mod.rs new file mode 100644 index 0000000..422e7f8 --- /dev/null +++ b/tests/async_generator_object_method_cjs/mod.rs @@ -0,0 +1,15 @@ +use crate::common::*; +use orchestrion_js::*; + +#[test] +fn async_generator_object_method_cjs() { + transpile_and_test( + file!(), + false, + Config::new_single(InstrumentationConfig::new( + "streamer_generate", + test_module_matcher(), + FunctionQuery::object_method("generate", FunctionKind::AsyncGenerator), + )), + ); +} diff --git a/tests/async_generator_object_method_cjs/test.js b/tests/async_generator_object_method_cjs/test.js new file mode 100644 index 0000000..b40b2d2 --- /dev/null +++ b/tests/async_generator_object_method_cjs/test.js @@ -0,0 +1,18 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +const { generate } = require('./instrumented.js'); +const { assert, getContext } = require('../common/preamble.js'); +const context = getContext('orchestrion:undici:streamer_generate'); +(async () => { + const values = []; + for await (const val of generate(3)) { + values.push(val); + } + assert.deepStrictEqual(values, [0, 1, 2]); + assert.deepStrictEqual(context, { + start: true, + end: true, + }); +})(); diff --git a/tests/async_generator_private_method_cjs/mod.js b/tests/async_generator_private_method_cjs/mod.js new file mode 100644 index 0000000..c215bee --- /dev/null +++ b/tests/async_generator_private_method_cjs/mod.js @@ -0,0 +1,15 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +class Streamer { + async *#generate(n) { + for (let i = 0; i < n; i++) yield i; + } + + generate(n) { + return this.#generate(n); + } +} + +module.exports = Streamer; diff --git a/tests/async_generator_private_method_cjs/mod.rs b/tests/async_generator_private_method_cjs/mod.rs new file mode 100644 index 0000000..ec664fe --- /dev/null +++ b/tests/async_generator_private_method_cjs/mod.rs @@ -0,0 +1,15 @@ +use crate::common::*; +use orchestrion_js::*; + +#[test] +fn async_generator_private_method_cjs() { + transpile_and_test( + file!(), + false, + Config::new_single(InstrumentationConfig::new( + "Streamer:generate", + test_module_matcher(), + FunctionQuery::private_method("Streamer", "generate", FunctionKind::AsyncGenerator), + )), + ); +} diff --git a/tests/async_generator_private_method_cjs/test.js b/tests/async_generator_private_method_cjs/test.js new file mode 100644 index 0000000..8990a74 --- /dev/null +++ b/tests/async_generator_private_method_cjs/test.js @@ -0,0 +1,19 @@ +/** + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. + **/ +const Streamer = require('./instrumented.js'); +const { assert, getContext } = require('../common/preamble.js'); +const context = getContext('orchestrion:undici:Streamer:generate'); +(async () => { + const streamer = new Streamer; + const values = []; + for await (const val of streamer.generate(3)) { + values.push(val); + } + assert.deepStrictEqual(values, [0, 1, 2]); + assert.deepStrictEqual(context, { + start: true, + end: true, + }); +})(); diff --git a/tests/instrumentor_test.rs b/tests/instrumentor_test.rs index 6c50772..dc140fb 100644 --- a/tests/instrumentor_test.rs +++ b/tests/instrumentor_test.rs @@ -5,6 +5,12 @@ mod common; mod arguments_mutation; +mod async_generator_class_method_cjs; +mod async_generator_decl_cjs; +mod async_generator_decl_mjs; +mod async_generator_expr_cjs; +mod async_generator_object_method_cjs; +mod async_generator_private_method_cjs; mod class_expression_cjs; mod class_method_cjs; mod constructor_cjs;