Add support for class constructor instrumentation (#3)

Author: bengl

src/function_query.rs

@@ -65,6 +65,7 @@ impl FunctionKind {
 
 pub struct FunctionQuery {
     pub name: String,
+    pub class: Option<String>,
     pub typ: FunctionType,
     pub kind: FunctionKind,
     pub index: usize,
@@ -113,10 +114,14 @@ impl TryFrom<&Yaml> for FunctionQuery {
         let typ = get_str!(query, "type");
         let kind = get_str!(query, "kind");
         let name = get_str!(query, "name");
+        let class = query["class"]
+            .as_str()
+            .map(std::string::ToString::to_string);
         let index: usize = query["index"].as_i64().unwrap_or(0).try_into().unwrap_or(0);
 
         Ok(FunctionQuery {
             name: name.to_string(),
+            class,
             typ: FunctionType::from_str(typ).ok_or(format!(
                 "Invalid config: 'type' must be one of 'decl', 'expr', or 'method', got '{typ}'"
             ))?,

src/instrumentation.rs

@@ -3,9 +3,9 @@ use std::path::PathBuf;
 use swc_core::common::{Span, SyntaxContext};
 use swc_core::ecma::{
     ast::{
-        ArrowExpr, AssignExpr, AssignTarget, BlockStmt, ClassMethod, Expr, FnDecl, FnExpr, Ident,
-        Lit, MemberProp, MethodProp, Module, ModuleItem, Pat, PropName, Script, SimpleAssignTarget,
-        Stmt, Str, VarDecl,
+        ArrowExpr, AssignExpr, AssignTarget, BlockStmt, ClassDecl, ClassMethod, Constructor, Expr,
+        FnDecl, FnExpr, Ident, Lit, MemberProp, MethodProp, Module, ModuleItem, Pat, PropName,
+        Script, SimpleAssignTarget, Stmt, Str, VarDecl,
     },
     atoms::Atom,
 };
@@ -26,11 +26,16 @@ macro_rules! ident {
 pub struct Instrumentation {
     config: InstrumentationConfig,
     count: usize,
+    is_correct_class: bool,
 }
 
 impl Instrumentation {
     pub(crate) fn new(config: InstrumentationConfig) -> Self {
-        Self { config, count: 0 }
+        Self {
+            config,
+            count: 0,
+            is_correct_class: false,
+        }
     }
 
     fn new_fn(&self, body: BlockStmt) -> ArrowExpr {
@@ -97,6 +102,49 @@ impl Instrumentation {
         ];
     }
 
+    fn insert_constructor_tracing(&self, body: &mut BlockStmt) {
+        let original_stmts = std::mem::take(&mut body.stmts);
+
+        let ch_ident = ident!(format!("tr_ch_apm${}", &self.config.channel_name));
+        let ctx_ident = ident!(format!("tr_ch_apm_ctx${}", &self.config.channel_name));
+        let mut try_catch = quote!(
+            "try {
+                if ($ch.hasSubscribers) {
+                    $ch.start.publish($ctx);
+                }
+            } catch (tr_ch_err) { 
+                if ($ch.hasSubscribers) {
+                    $ctx.error = tr_ch_err;
+                    try {
+                        $ctx.self = this;
+                    } catch (refErr) {
+                        // This can only error out if super hasn't been called yet.
+                        // Safe to ignore, but note that self/this won't get into the context.
+                    }
+                    $ch.error.publish($ctx);
+                }
+                throw tr_ch_err;
+            } finally {
+                if ($ch.hasSubscribers) {
+                    $ctx.self = this;
+                    $ch.end.publish($ctx);
+                }
+            }" as Stmt,
+            ch = ch_ident,
+            ctx = ctx_ident.clone(),
+        );
+        if let Some(try_catch_stmt) = try_catch.as_mut_try_stmt() {
+            for stmt in &original_stmts {
+                try_catch_stmt.block.stmts.push(stmt.clone());
+            }
+        }
+
+        body.stmts = vec![
+            quote!("const $ctx = { arguments };" as Stmt, ctx = ctx_ident,),
+            try_catch,
+        ];
+    }
+
     fn trace_expr_or_count(&mut self, func_expr: &mut FnExpr, name: &Atom) -> bool {
         if self
             .config
@@ -163,11 +211,27 @@ impl Instrumentation {
         !traced
     }
 
+    pub fn visit_mut_class_decl(&mut self, node: &mut ClassDecl) -> bool {
+        self.is_correct_class = self
+            .config
+            .function_query
+            .class
+            .as_ref()
+            .map_or(true, |class| node.ident.sym.as_ref() == class);
+        true
+    }
+
     pub fn visit_mut_class_method(&mut self, node: &mut ClassMethod) -> bool {
         let name = match &node.key {
             PropName::Ident(ident) => ident.sym.clone(),
             _ => return false,
         };
+
+        // Only increment count when class matches
+        if !self.is_correct_class {
+            return true;
+        }
+
         if self
             .config
             .function_query
@@ -178,6 +242,21 @@ impl Instrumentation {
                 self.insert_tracing(body);
             }
         }
+        true
+    }
+
+    pub fn visit_mut_constructor(&mut self, node: &mut Constructor) -> bool {
+        if !self.is_correct_class || self.config.function_query.name != "constructor" {
+            return false;
+        }
+
+        if self.count == self.config.function_query.index && node.body.is_some() {
+            if let Some(body) = node.body.as_mut() {
+                self.insert_constructor_tracing(body);
+            }
+        } else {
+            self.count += 1;
+        }
         false
     }
 

src/lib.rs

@@ -18,7 +18,10 @@ use std::path::PathBuf;
 use std::str::FromStr;
 
 use swc_core::ecma::{
-    ast::{AssignExpr, ClassMethod, FnDecl, MethodProp, Module, Script, Str, VarDecl},
+    ast::{
+        AssignExpr, ClassDecl, ClassMethod, Constructor, FnDecl, MethodProp, Module, Script, Str,
+        VarDecl,
+    },
     visit::{VisitMut, VisitMutWith},
 };
 use swc_core::quote;
@@ -142,7 +145,9 @@ impl VisitMut for InstrumentationVisitor<'_> {
 
     visit_with_all_fn!(visit_mut_fn_decl, FnDecl);
     visit_with_all_fn!(visit_mut_var_decl, VarDecl);
-    visit_with_all_fn!(visit_mut_class_method, ClassMethod);
     visit_with_all_fn!(visit_mut_method_prop, MethodProp);
     visit_with_all_fn!(visit_mut_assign_expr, AssignExpr);
+    visit_with_all_fn!(visit_mut_class_decl, ClassDecl);
+    visit_with_all_fn!(visit_mut_class_method, ClassMethod);
+    visit_with_all_fn!(visit_mut_constructor, Constructor);
 }

tests/class_method_cjs/instrumentations.yml

@@ -4,9 +4,9 @@ instrumentations:
     version_range: ">=0.0.1"
     file_path: index.mjs
     function_query:
+      class: Undici
       name: fetch
       type: method
       kind: async
-      index: 0
     operator: tracePromise
     channel_name: Undici_fetch

tests/class_method_cjs/mod.js

@@ -1,7 +1,12 @@
-class Undici {
-    async fetch (lmao) {
+class UndiciBase {
+    async fetch (url) {
         return 42;
     }
 }
+class Undici extends UndiciBase {
+    async fetch (url) {
+        return super.fetch(url);
+    }
+}
 
 module.exports = Undici;

tests/constructor_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: constructor
+      type: method
+      kind: sync
+      class: Undici
+    operator: traceSync
+    channel_name: Undici_constructor

tests/constructor_cjs/mod.js

@@ -0,0 +1,13 @@
+class UndiciBase {
+    constructor () {
+    }
+}
+
+class Undici extends UndiciBase {
+    constructor (val) {
+        super();
+        this.val = val;
+    }
+}
+
+module.exports = Undici;

tests/constructor_cjs/test.js

@@ -0,0 +1,11 @@
+const Undici = require('./instrumented.js');
+const { assert, getContext } = require('../common/preamble.js');
+const context = getContext('orchestrion:undici:Undici_constructor');
+(() => {
+  const undici = new Undici(42);
+  assert.deepEqual(undici.val, 42);
+  assert.deepStrictEqual(context, {
+    start: true,
+    end: true
+  });
+})();

tests/constructor_mjs/instrumentations.yml

@@ -0,0 +1,12 @@
+version: 1
+instrumentations:
+  - module_name: undici
+    version_range: ">=0.0.1"
+    file_path: index.mjs
+    function_query:
+      name: constructor
+      type: method
+      kind: sync
+      class: Undici
+    operator: traceSync
+    channel_name: Undici_constructor

tests/constructor_mjs/mod.mjs

@@ -0,0 +1,13 @@
+class UndiciBase {
+  constructor() {
+  }
+}
+
+class Undici extends UndiciBase {
+  constructor(val) {
+    super();
+    this.val = val;
+  }
+}
+
+export { Undici };