Added support for traceCallback

src/config.rs

@@ -47,6 +47,23 @@ impl ModuleMatcher {
     }
 }
 
+#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Serialize, serde::Deserialize),
+    serde(rename_all = "camelCase")
+)]
+#[derive(Debug, Clone)]
+/// `CallbackConfig` represents details needed when wrapping a function that
+/// accepts a callback. These details will be used to construct the
+/// `tracingChannel.traceCallback` invocation.
+pub struct CallbackConfig {
+    /// `position` is the ordinal of the callback function within the wrapped
+    /// function's parameter list. A value of `-1` indicates that the callback
+    /// is the last parameter in the list. The value is zero based.
+    pub position: i32,
+}
+
 #[cfg_attr(
     feature = "serde",
     derive(serde::Serialize, serde::Deserialize),
@@ -62,16 +79,35 @@ pub struct InstrumentationConfig {
     pub channel_name: String,
     pub module: ModuleMatcher,
     pub function_query: FunctionQuery,
+    #[tsify(optional)]
+    #[serde(default = "InstrumentationConfig::empty_callback_config")]
+    pub callback_config: CallbackConfig,
 }
 
 impl InstrumentationConfig {
     #[must_use]
     pub fn new(channel_name: &str, module: ModuleMatcher, function_query: FunctionQuery) -> Self {
-        Self {
+        return Self::new_with_callback_config(
+            channel_name,
+            module,
+            function_query,
+            CallbackConfig { position: -1 },
+        );
+    }
+
+    #[must_use]
+    pub fn new_with_callback_config(
+        channel_name: &str,
+        module: ModuleMatcher,
+        function_query: FunctionQuery,
+        callback_config: CallbackConfig,
+    ) -> Self {
+        return Self {
             channel_name: channel_name.to_string(),
             module,
             function_query,
-        }
+            callback_config,
+        };
     }
 
     #[must_use]
@@ -114,4 +150,9 @@ impl InstrumentationConfig {
     pub fn matches(&self, module_name: &str, version: &str, file_path: &PathBuf) -> bool {
         self.module.matches(module_name, version, file_path)
     }
+
+    #[must_use]
+    pub fn empty_callback_config() -> CallbackConfig {
+        return CallbackConfig { position: -1 };
+    }
 }

src/function_query.rs

@@ -17,6 +17,7 @@ pub(crate) enum FunctionType {
 pub enum FunctionKind {
     Sync,
     Async,
+    Callback,
 }
 
 impl FunctionKind {
@@ -25,11 +26,17 @@ impl FunctionKind {
         matches!(self, FunctionKind::Async)
     }
 
+    #[must_use]
+    pub fn is_callback(&self) -> bool {
+        matches!(self, FunctionKind::Callback)
+    }
+
     #[must_use]
     pub fn tracing_operator(&self) -> &'static str {
         match self {
             FunctionKind::Sync => "traceSync",
             FunctionKind::Async => "tracePromise",
+            FunctionKind::Callback => "traceCallback",
         }
     }
 }
@@ -43,7 +50,7 @@ impl FunctionKind {
 #[derive(Debug, Clone)]
 pub enum FunctionQuery {
     // The order here matters because this enum is untagged, serde will try
-    // choose the first variant that matches the data.
+    // to choose the first variant that matches the data.
     ClassMethod {
         class_name: String,
         method_name: String,

src/instrumentation.rs

@@ -119,19 +119,20 @@ impl Instrumentation {
             ctxt: SyntaxContext::empty(),
             stmts: vec![
                 quote!("const __apm$wrapped = $wrapped;" as Stmt, wrapped: Expr = wrapped_fn.into()),
-                quote!("return __apm$wrapped.apply(null, __apm$original_args);" as Stmt),
+                // TODO: instead of a static `this` we should be able to declare the value
+                quote!("return __apm$wrapped.apply(this, __apm$original_args);" as Stmt),
             ],
         };
 
         let traced_fn = self.new_fn(traced_body, vec![]);
 
         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${}.{}",
+        let trace_statement = construct_trace_statement(
+            &self.config,
             &id_name,
-            self.config.function_query.kind().tracing_operator()
-        ));
+            self.module_version.clone().unwrap_or_default().as_str(),
+        );
 
         body.stmts = vec![
             quote!("const __apm$original_args = arguments" as Stmt),
@@ -140,18 +141,7 @@ impl Instrumentation {
                 "if (!$ch.hasSubscribers) return __apm$traced();" as Stmt,
                 ch = ch_ident
             ),
-            match &self.module_version {
-                Some(version) => quote!(
-                    "return $trace(__apm$traced, { arguments, self: this, moduleVersion: $version } );"
-                        as Stmt,
-                    trace = trace_ident,
-                    version: Expr = version.as_str().into(),
-                ),
-                None => quote!(
-                    "return $trace(__apm$traced, { arguments, self: this } );" as Stmt,
-                    trace = trace_ident,
-                ),
-            },
+            trace_statement,
         ];
 
         self.has_injected = true;
@@ -235,7 +225,7 @@ impl Instrumentation {
     }
 
     // The rest of these functions are from `VisitMut`, except they return a boolean to indicate
-    // whether recusrsing through the tree is necessary, rather than calling
+    // whether recursing through the tree is necessary, rather than calling
     // `visit_mut_children_with`.
 
     pub fn visit_mut_module(&mut self, node: &mut Module) -> bool {
@@ -356,7 +346,7 @@ impl Instrumentation {
 
     pub fn visit_mut_assign_expr(&mut self, node: &mut AssignExpr) -> bool {
         // TODO(bengl) This is by far the hardest bit. We're trying to infer a name for this
-        // function expresion using the surrounding code, but it's not always possible, and even
+        // function expression using the surrounding code, but it's not always possible, and even
         // where it is, there are so many ways to give a function expression a "name", that the
         // code paths here can get pretty hairy. Right now this is only covering some basic cases.
         // The following cases are missing:
@@ -398,3 +388,33 @@ pub fn get_script_start_index(script: &Script) -> usize {
     }
     0
 }
+
+/// Builds a tracing channel trace invocation. For example, the result may
+/// look like:
+///
+/// ```js
+/// return tr_ch_amp$myChannel.traceSync(__apm$traced, { arguments, self: this, moduleVersion: "1.0.0" })
+/// ```
+#[rustfmt::skip]
+fn construct_trace_statement(config: &InstrumentationConfig, channel_name: &str, mod_version: &str) -> Stmt {
+    let mut ctx = "{ arguments, self: this }".to_string();
+    if mod_version.is_empty() == false {
+        ctx = ["{ arguments, self: this, moduleVersion: \"", mod_version, "\" }"].join("");
+    }
+
+    let operator = ["tr_ch_apm$", channel_name, ".", config.function_query.kind().tracing_operator()].join("");
+    #[allow(clippy::needless_late_init)]
+    let stmt_str;
+    if config.function_query.kind().is_callback() == true {
+        // TODO: figure out how we can pass a `this` argument
+        let pos = config.callback_config.position.to_string();
+        stmt_str = ["return ", &*operator, "(__apm$traced, ", pos.as_str(), ", ", ctx.as_str(), ", this, ...arguments)"].join("");
+    } else {
+        stmt_str = ["return ", &*operator, "(__apm$traced, ", ctx.as_str(), ")"].join("");
+    }
+
+    return quote!(
+        "$stmt" as Stmt,
+        stmt = Ident::new(Atom::from(stmt_str), Span::default(), SyntaxContext::empty())
+    )
+}

src/lib.rs

@@ -13,6 +13,8 @@
 #![deny(clippy::complexity)]
 #![deny(clippy::correctness)]
 #![deny(clippy::unwrap_used)]
+#![allow(clippy::bool_comparison)]
+#![allow(clippy::needless_return)]
 
 /**
  * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.

tests/Readme.md

@@ -0,0 +1,47 @@
+# Unit Tests
+
+Unit tests in this project follow a specific structure:
+
+1. A directory named the same as the unit test. For example,
+the `index_cjs` test will be within the `index_cjs/` directory.
+2. A `mod.rs` file that contains the unit test.
+3. A `mod.{js,mjs}` file that contains a JavaScript script to apply
+the test against, i.e. this replicates a potential module from
+npmjs.com.
+4. A `test.js` file that exercises the patched code.
+5. A new `mod` line in the [`instrumentor_test.rs`](./instrumentor_test.rs)
+file.
+
+To run a specific test from the command line:
+
+```sh
+# To run the index_cjs test only:
+cargo test index_cjs
+```
+
+Each unit test should utilize the `transpile_and_test` function
+exported by [`common/mod.rs`](./common/mod.rs). This function:
+
+1. Reads in the unit test file to get the base path to the test.
+2. Reads in the JavaScript module file.
+3. Transforms the read-in JavaScript according to the defined configuration.
+4. Generates a new `instrumented.js` file within the test directory. This
+file contains the `mod.{js,mjs}` code patched by our tool.
+5. Runs `node test.js` within the test directory. The result of this test
+should be a `0` exit code for the process.
+
+## Running In A Debugger
+
+### RustRover
+
+To run a specific test through the [RustRover](https://www.jetbrains.com/rust/)
+debugger:
+
+1. Create a new run configuration of type "Cargo".
+2. For the "command", enter (e.g. for the `index_cjs` test):
+
+        test --package orchestrion-js --test instrumentor_test index_cjs::index_cjs -- --exact
+3. Set breakpoints and run the profile.
+
+For a shortcut, open the desired `mod.rs` test file and click the "run"
+button in the gutter.

tests/callback_cjs/mod.js

@@ -0,0 +1,12 @@
+'use strict'
+
+module.exports = Foo
+
+function Foo() {}
+
+Foo.prototype.doWork = function doWork(input, callback) {
+  if (input === 'fail') {
+    return callback(Error('boom'))
+  }
+  return callback(null, 'done')
+}

tests/callback_cjs/mod.rs

@@ -0,0 +1,25 @@
+use crate::common::*;
+use orchestrion_js::*;
+
+#[test]
+fn callback_cjs() {
+    transpile_and_test_with_imports(
+        file!(),
+        false,
+        Config::new_single(InstrumentationConfig::new_with_callback_config(
+            "foo_ch",
+            ModuleMatcher::new("foo", ">=1.0.0", "index.js").unwrap(),
+            FunctionQuery::FunctionExpression {
+                expression_name: "doWork".to_string(),
+                kind: FunctionKind::Callback,
+                index: 0,
+            },
+            CallbackConfig { position: 1 },
+        )),
+        &[PackageImport {
+            module_name: "foo".to_string(),
+            module_version: "1.1.1".to_string(),
+            file: "index.js".to_string(),
+        }],
+    )
+}

tests/callback_cjs/test.js

@@ -0,0 +1,46 @@
+'use strict'
+
+const assert = require('node:assert')
+const { getContext } = require('../common/preamble.js');
+const context = getContext('orchestrion:foo:foo_ch');
+const Foo = require('./instrumented.js')
+
+const start = process.hrtime.bigint()
+const timer1 = setInterval(timeout, 1_000)
+const timer2 = setInterval(timeout, 1_000)
+const foo = new Foo()
+
+foo.doWork('success', (error, result) => {
+  assert.equal(error, undefined)
+  assert.equal(result, 'done')
+
+  assert.deepStrictEqual(context, {
+    start: true
+  })
+  delete context.start
+
+  clearInterval(timer1)
+})
+
+foo.doWork('fail', (error, result) => {
+  assert.equal(error.message, 'boom')
+  assert.equal(undefined, result)
+
+  assert.deepStrictEqual(context, {
+    end: true,
+    start: true
+  })
+  delete context.end
+  delete context.start
+
+  clearInterval(timer2)
+})
+
+function timeout() {
+  const current = process.hrtime.bigint()
+  if ((current - start) > 3e+10) {
+    // Exceeded 30 seconds, terminate it all.
+    clearInterval(timer1)
+    clearInterval(timer2)
+  }
+}

tests/common/mod.rs

@@ -12,14 +12,43 @@ use swc::config::IsModule;
 static TEST_MODULE_NAME: &str = "undici";
 static TEST_MODULE_PATH: &str = "index.mjs";
 
+/// PackageImport represents metadata around a module that we are pretending
+/// was imported from a npm module. So, if we are pretending to test
+/// instrumentation of a module `@foo/bar`, the `module_name` would be
+/// `@foo/bar`, the `module_version` would be a simple version string like
+/// `0.1.0`, and the `file` would be some filename within that package like
+/// `index.js` or `lib/utils.js`.
+///
+/// This information will be used to match a unit test to the instrumentation.
+pub struct PackageImport {
+    pub module_name: String,
+    pub module_version: String,
+    pub file: String,
+}
+
 pub fn transpile_and_test(test_file: &str, mjs: bool, config: Config) {
+    transpile_and_test_with_imports(test_file, mjs, config, &[]);
+}
+
+#[rustfmt::skip]
+pub fn transpile_and_test_with_imports(test_file: &str, mjs: bool, config: Config, imports: &[PackageImport]) {
     let test_file = PathBuf::from(test_file);
     let test_dir = test_file.parent().expect("Couldn't find test directory");
 
     let file_path = PathBuf::from("index.mjs");
     let instrumentor = Instrumentor::new(config);
-    let mut instrumentations =
+    let mut instrumentations;
+    if imports.len() > 0 {
+        let import = &imports[0];
+         instrumentations = instrumentor.get_matching_instrumentations(
+            import.module_name.as_str(),
+            import.module_version.as_str(),
+            &PathBuf::from(&import.file)
+        );
+    } else {
+       instrumentations =
         instrumentor.get_matching_instrumentations(TEST_MODULE_NAME, "0.0.1", &file_path);
+    }
 
     let extension = if mjs { "mjs" } else { "js" };
     let instrumentable = test_dir.join(format!("mod.{extension}"));