diff --git a/src/instrumentation.rs b/src/instrumentation.rs index f2d0247..cdfee49 100644 --- a/src/instrumentation.rs +++ b/src/instrumentation.rs @@ -186,15 +186,19 @@ impl Instrumentation { // `visit_mut_children_with`. pub fn visit_mut_module(&mut self, node: &mut Module) -> bool { - node.body - .insert(1, ModuleItem::Stmt(self.create_tracing_channel())); + let channel_element = ModuleItem::Stmt(self.create_tracing_channel()); + if !node.body.iter().any(|item| item == &channel_element) { + node.body.insert(1, channel_element); + } true } pub fn visit_mut_script(&mut self, node: &mut Script) -> bool { let start_index = get_script_start_index(node); - node.body - .insert(start_index + 1, self.create_tracing_channel()); + let channel_element = self.create_tracing_channel(); + if !node.body.iter().any(|item| item == &channel_element) { + node.body.insert(start_index + 1, channel_element); + } true } diff --git a/tests/instrumentor_test.rs b/tests/instrumentor_test.rs index ec5b9fd..840f092 100644 --- a/tests/instrumentor_test.rs +++ b/tests/instrumentor_test.rs @@ -17,3 +17,5 @@ mod multiple_load_cjs; mod object_method_cjs; mod polyfill_cjs; mod polyfill_mjs; +mod single_channel_multiple_usages_cjs; +mod single_channel_multiple_usages_mjs; diff --git a/tests/single_channel_multiple_usages_cjs/mod.js b/tests/single_channel_multiple_usages_cjs/mod.js new file mode 100644 index 0000000..0d4b23b --- /dev/null +++ b/tests/single_channel_multiple_usages_cjs/mod.js @@ -0,0 +1,9 @@ +async function foo () { + return 'foo' +} + +async function bar () { + return 'bar' +} + +module.exports = { foo, bar }; \ No newline at end of file diff --git a/tests/single_channel_multiple_usages_cjs/mod.rs b/tests/single_channel_multiple_usages_cjs/mod.rs new file mode 100644 index 0000000..62dc5f2 --- /dev/null +++ b/tests/single_channel_multiple_usages_cjs/mod.rs @@ -0,0 +1,25 @@ +use crate::common::*; +use orchestrion_js::*; + +#[test] +fn same_channel_multiple_usages_cjs() { + transpile_and_test( + file!(), + false, + Config::new( + vec![ + InstrumentationConfig::new( + "method_call", + test_module_matcher(), + FunctionQuery::function_declaration("foo", FunctionKind::Async), + ), + InstrumentationConfig::new( + "method_call", + test_module_matcher(), + FunctionQuery::function_declaration("bar", FunctionKind::Async), + ), + ], + None, + ), + ); +} diff --git a/tests/single_channel_multiple_usages_cjs/test.js b/tests/single_channel_multiple_usages_cjs/test.js new file mode 100644 index 0000000..0012a47 --- /dev/null +++ b/tests/single_channel_multiple_usages_cjs/test.js @@ -0,0 +1,26 @@ +/** + * 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 { foo, bar } = require('./instrumented.js'); +const { assert, getContext } = require('../common/preamble.js'); +const context = getContext('orchestrion:undici:method_call'); +(async () => { + const result = await foo(); + assert.strictEqual(result, 'foo'); + assert.deepStrictEqual(context, { + start: true, + end: true, + asyncStart: 'foo', + asyncEnd: 'foo' + }); + + const result2 = await bar(); + assert.strictEqual(result2, 'bar'); + assert.deepStrictEqual(context, { + start: true, + end: true, + asyncStart: 'bar', + asyncEnd: 'bar' + }); +})(); \ No newline at end of file diff --git a/tests/single_channel_multiple_usages_mjs/mod.mjs b/tests/single_channel_multiple_usages_mjs/mod.mjs new file mode 100644 index 0000000..ced4f48 --- /dev/null +++ b/tests/single_channel_multiple_usages_mjs/mod.mjs @@ -0,0 +1,9 @@ +async function foo () { + return 'foo' +} + +async function bar () { + return 'bar' +} + +export { foo, bar }; \ No newline at end of file diff --git a/tests/single_channel_multiple_usages_mjs/mod.rs b/tests/single_channel_multiple_usages_mjs/mod.rs new file mode 100644 index 0000000..55d6373 --- /dev/null +++ b/tests/single_channel_multiple_usages_mjs/mod.rs @@ -0,0 +1,25 @@ +use crate::common::*; +use orchestrion_js::*; + +#[test] +fn same_channel_multiple_usages_mjs() { + transpile_and_test( + file!(), + true, + Config::new( + vec![ + InstrumentationConfig::new( + "method_call", + test_module_matcher(), + FunctionQuery::function_declaration("foo", FunctionKind::Async), + ), + InstrumentationConfig::new( + "method_call", + test_module_matcher(), + FunctionQuery::function_declaration("bar", FunctionKind::Async), + ), + ], + None, + ), + ); +} diff --git a/tests/single_channel_multiple_usages_mjs/test.mjs b/tests/single_channel_multiple_usages_mjs/test.mjs new file mode 100644 index 0000000..9c98fec --- /dev/null +++ b/tests/single_channel_multiple_usages_mjs/test.mjs @@ -0,0 +1,24 @@ +/** + * 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 { foo, bar } from './instrumented.mjs'; +import { assert, getContext } from '../common/preamble.js'; +const context = getContext('orchestrion:undici:method_call'); +const result = await foo(); +assert.strictEqual(result, 'foo'); +assert.deepStrictEqual(context, { + start: true, + end: true, + asyncStart: 'foo', + asyncEnd: 'foo' +}); + +const result2 = await bar(); +assert.strictEqual(result2, 'bar'); +assert.deepStrictEqual(context, { + start: true, + end: true, + asyncStart: 'bar', + asyncEnd: 'bar' +}); \ No newline at end of file