diff --git a/src/config.rs b/src/config.rs index 3c8f44e..b8e0ac9 100644 --- a/src/config.rs +++ b/src/config.rs @@ -23,7 +23,7 @@ macro_rules! get_arr { }; } -#[derive(Clone)] +#[derive(Clone, Debug)] pub enum InstrumentationOperator { Callback, Promise, @@ -52,6 +52,7 @@ impl InstrumentationOperator { } } +#[derive(Debug)] pub struct InstrumentationConfig { pub module_name: String, pub version_range: Range, diff --git a/src/function_query.rs b/src/function_query.rs index d74ab81..a920052 100644 --- a/src/function_query.rs +++ b/src/function_query.rs @@ -10,6 +10,7 @@ macro_rules! get_str { }; } +#[derive(Debug)] pub enum FunctionType { FunctionDeclaration, FunctionExpression, @@ -27,6 +28,7 @@ impl FunctionType { } } +#[derive(Debug)] pub enum FunctionKind { Sync, Async, @@ -63,6 +65,7 @@ impl FunctionKind { } } +#[derive(Debug)] pub struct FunctionQuery { pub name: String, pub class: Option, diff --git a/src/instrumentation.rs b/src/instrumentation.rs index 74df0d1..13e34b8 100644 --- a/src/instrumentation.rs +++ b/src/instrumentation.rs @@ -38,6 +38,11 @@ impl Instrumentation { } } + pub(crate) fn reset(&mut self) { + self.count = 0; + self.is_correct_class = false; + } + fn new_fn(&self, body: BlockStmt) -> ArrowExpr { ArrowExpr { params: vec![], @@ -70,7 +75,9 @@ impl Instrumentation { define_channel } - fn insert_tracing(&self, body: &mut BlockStmt) { + fn insert_tracing(&mut self, body: &mut BlockStmt) { + self.count += 1; + let original_stmts = std::mem::take(&mut body.stmts); // Create a new BlockStmt with the original statements @@ -102,7 +109,9 @@ impl Instrumentation { ]; } - fn insert_constructor_tracing(&self, body: &mut BlockStmt) { + fn insert_constructor_tracing(&mut self, body: &mut BlockStmt) { + self.count += 1; + let original_stmts = std::mem::take(&mut body.stmts); let ch_ident = ident!(format!("tr_ch_apm${}", &self.config.channel_name)); diff --git a/src/lib.rs b/src/lib.rs index 75ec30d..aa82f56 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -104,7 +104,8 @@ macro_rules! visit_with_all { ($self:expr, $method:ident, $item:expr) => { let mut recurse = false; for instr in &mut $self.instrumentations { - recurse = recurse || instr.$method($item); + let needs_recurse = instr.$method($item); + recurse = recurse || needs_recurse; } if recurse { $item.visit_mut_children_with($self); @@ -132,6 +133,9 @@ impl VisitMut for InstrumentationVisitor<'_> { } } visit_with_all!(self, visit_mut_module, item); + for instr in &mut self.instrumentations { + instr.reset(); + } } fn visit_mut_script(&mut self, item: &mut Script) { @@ -141,6 +145,9 @@ impl VisitMut for InstrumentationVisitor<'_> { ); item.body.insert(get_script_start_index(item), import); visit_with_all!(self, visit_mut_script, item); + for instr in &mut self.instrumentations { + instr.reset(); + } } visit_with_all_fn!(visit_mut_fn_decl, FnDecl); diff --git a/tests/index_cjs/instrumentations.yml b/tests/index_cjs/instrumentations.yml new file mode 100644 index 0000000..725dda5 --- /dev/null +++ b/tests/index_cjs/instrumentations.yml @@ -0,0 +1,13 @@ +version: 1 +instrumentations: + - module_name: undici + version_range: ">=0.0.1" + file_path: index.mjs + function_query: + class: Undici + name: fetch + type: method + kind: async + index: 2 + operator: tracePromise + channel_name: Undici_fetch diff --git a/tests/index_cjs/mod.js b/tests/index_cjs/mod.js new file mode 100644 index 0000000..437d9c6 --- /dev/null +++ b/tests/index_cjs/mod.js @@ -0,0 +1,45 @@ +{ + class Undici { + async fetch (url) { + return 0; + } + } + + exports.Undici0 = Undici; +} +{ + class Undici { + async fetch (url) { + return 1; + } + } + + exports.Undici1 = Undici; +} +{ + class Undici { + async fetch (url) { + return 2; + } + } + + exports.Undici2 = Undici; +} +{ + class Undici { + async fetch (url) { + return 3; + } + } + + exports.Undici3 = Undici; +} +{ + class Undici { + async fetch (url) { + return 4; + } + } + + exports.Undici4 = Undici; +} diff --git a/tests/index_cjs/test.js b/tests/index_cjs/test.js new file mode 100644 index 0000000..525abe7 --- /dev/null +++ b/tests/index_cjs/test.js @@ -0,0 +1,27 @@ +const undicis = require('./instrumented.js'); +const { assert, getContext } = require('../common/preamble.js'); +const context = getContext('orchestrion:undici:Undici_fetch'); + +async function testOne(Undici, num, expectedCtx) { + const undici = new Undici; + const result = await undici.fetch('https://example.com'); + assert.strictEqual(result, num); + assert.deepStrictEqual(context, expectedCtx); + delete context.start; + delete context.end; + delete context.asyncStart; + delete context.asyncEnd; +} + +(async () => { + await testOne(undicis.Undici0, 0, {}); + await testOne(undicis.Undici1, 1, {}); + await testOne(undicis.Undici2, 2, { + start: true, + end: true, + asyncStart: 2, + asyncEnd: 2 + }); + await testOne(undicis.Undici3, 3, {}); + await testOne(undicis.Undici4, 4, {}); +})(); diff --git a/tests/instrumentor_test.rs b/tests/instrumentor_test.rs index 0aa853f..b7926f5 100644 --- a/tests/instrumentor_test.rs +++ b/tests/instrumentor_test.rs @@ -13,6 +13,9 @@ macro_rules! make_test { instrumentor.get_matching_instrumentations("undici", "0.0.1", &file_path); transpile_and_test(stringify!($name), $mjs, &mut instrumentations); + + // It has to work twice, since we might use the same instrumentor on multiple files + transpile_and_test(stringify!($name), $mjs, &mut instrumentations); } }; } @@ -27,6 +30,8 @@ make_test!(expr_cjs, false); make_test!(class_method_cjs, false); +make_test!(multiple_class_method_cjs, false); + make_test!(object_method_cjs, false); make_test!(constructor_cjs, false); @@ -36,3 +41,7 @@ make_test!(constructor_mjs, true); make_test!(polyfill_mjs, true); make_test!(polyfill_cjs, false); + +make_test!(index_cjs, false); + +make_test!(no_index_cjs, false); diff --git a/tests/multiple_class_method_cjs/instrumentations.yml b/tests/multiple_class_method_cjs/instrumentations.yml new file mode 100644 index 0000000..c68c482 --- /dev/null +++ b/tests/multiple_class_method_cjs/instrumentations.yml @@ -0,0 +1,22 @@ +version: 1 +instrumentations: + - module_name: undici + version_range: ">=0.0.1" + file_path: index.mjs + function_query: + class: Undici + name: fetch1 + type: method + kind: async + operator: tracePromise + channel_name: Undici_fetch1 + - module_name: undici + version_range: ">=0.0.1" + file_path: index.mjs + function_query: + class: Undici + name: fetch2 + type: method + kind: async + operator: tracePromise + channel_name: Undici_fetch2 diff --git a/tests/multiple_class_method_cjs/mod.js b/tests/multiple_class_method_cjs/mod.js new file mode 100644 index 0000000..d5d622b --- /dev/null +++ b/tests/multiple_class_method_cjs/mod.js @@ -0,0 +1,11 @@ +class Undici { + async fetch1 (url) { + return 42; + } + + async fetch2 (url) { + return 43; + } +} + +module.exports = Undici; diff --git a/tests/multiple_class_method_cjs/test.js b/tests/multiple_class_method_cjs/test.js new file mode 100644 index 0000000..75568f3 --- /dev/null +++ b/tests/multiple_class_method_cjs/test.js @@ -0,0 +1,24 @@ +const Undici = require('./instrumented.js'); +const { assert, getContext } = require('../common/preamble.js'); +const context1 = getContext('orchestrion:undici:Undici_fetch1'); +const context2 = getContext('orchestrion:undici:Undici_fetch2'); + +(async () => { + const undici = new Undici; + const result1 = await undici.fetch1('https://example.com'); + assert.strictEqual(result1, 42); + assert.deepStrictEqual(context1, { + start: true, + end: true, + asyncStart: 42, + asyncEnd: 42 + }); + const result2 = await undici.fetch2('https://example.com'); + assert.strictEqual(result2, 43); + assert.deepStrictEqual(context2, { + start: true, + end: true, + asyncStart: 43, + asyncEnd: 43 + }); +})(); diff --git a/tests/no_index_cjs/instrumentations.yml b/tests/no_index_cjs/instrumentations.yml new file mode 100644 index 0000000..acabc37 --- /dev/null +++ b/tests/no_index_cjs/instrumentations.yml @@ -0,0 +1,12 @@ +version: 1 +instrumentations: + - module_name: undici + version_range: ">=0.0.1" + file_path: index.mjs + function_query: + name: fetch + type: decl + kind: async + # intentionally omitting index + operator: tracePromise + channel_name: fetch_no_index \ No newline at end of file diff --git a/tests/no_index_cjs/mod.js b/tests/no_index_cjs/mod.js new file mode 100644 index 0000000..4c130bf --- /dev/null +++ b/tests/no_index_cjs/mod.js @@ -0,0 +1,5 @@ +async function fetch(url) { + return 42; +} + +module.exports = { fetch }; \ No newline at end of file diff --git a/tests/no_index_cjs/test.js b/tests/no_index_cjs/test.js new file mode 100644 index 0000000..96f5124 --- /dev/null +++ b/tests/no_index_cjs/test.js @@ -0,0 +1,13 @@ +const { fetch } = require('./instrumented.js'); +const { assert, getContext } = require('../common/preamble.js'); +const context = getContext('orchestrion:undici:fetch_no_index'); +(async () => { + const result = await fetch('https://example.com'); + assert.strictEqual(result, 42); + assert.deepStrictEqual(context, { + start: true, + end: true, + asyncStart: 42, + asyncEnd: 42 + }); +})(); \ No newline at end of file