feat: Added support for wrapping private class methods

timfish · web-flow · commit 22d069d23ea5 · 2026-02-26T21:40:57.000Z
diff --git a/README.md b/README.md
@@ -129,6 +129,8 @@ type FunctionQuery =
     { functionName: string; kind: FunctionKind; index?: number; isExportAlias?: boolean }
     | // Match arrow function or function expression
     { expressionName: string; kind: FunctionKind; index?: number; isExportAlias?: boolean };
+    | // Match private class methods
+    { className: string; privateMethodName: string; kind: FunctionKind; index?: number };
 ```
 
 #### **`ModuleMatcher`**
diff --git a/src/function_query.rs b/src/function_query.rs
@@ -10,6 +10,7 @@ pub(crate) enum FunctionType {
     FunctionDeclaration,
     FunctionExpression,
     Method,
+    PrivateMethod,
 }
 
 /// The kind of function - Sync or returns a promise
@@ -58,6 +59,14 @@ pub enum FunctionQuery {
         #[cfg_attr(feature = "wasm", tsify(optional))]
         is_export_alias: bool,
     },
+    PrivateMethod {
+        class_name: String,
+        private_method_name: String,
+        kind: FunctionKind,
+        #[cfg_attr(feature = "serde", serde(default))]
+        #[cfg_attr(feature = "wasm", tsify(optional))]
+        index: usize,
+    },
     ClassConstructor {
         class_name: String,
         #[cfg_attr(feature = "serde", serde(default))]
@@ -146,6 +155,16 @@ impl FunctionQuery {
         }
     }
 
+    #[must_use]
+    pub fn private_method(class_name: &str, private_method_name: &str, kind: FunctionKind) -> Self {
+        FunctionQuery::PrivateMethod {
+            class_name: class_name.to_string(),
+            private_method_name: private_method_name.to_string(),
+            kind,
+            index: 0,
+        }
+    }
+
     #[must_use]
     pub fn as_export_alias(mut self) -> Self {
         match &mut self {
@@ -161,7 +180,7 @@ impl FunctionQuery {
             | FunctionQuery::FunctionExpression {
                 is_export_alias, ..
             } => *is_export_alias = true,
-            FunctionQuery::ObjectMethod { .. } => {}
+            FunctionQuery::ObjectMethod { .. } | FunctionQuery::PrivateMethod { .. } => {}
         }
         self
     }
@@ -172,7 +191,8 @@ impl FunctionQuery {
             FunctionQuery::ClassMethod { kind, .. }
             | FunctionQuery::ObjectMethod { kind, .. }
             | FunctionQuery::FunctionDeclaration { kind, .. }
-            | FunctionQuery::FunctionExpression { kind, .. } => kind,
+            | FunctionQuery::FunctionExpression { kind, .. }
+            | FunctionQuery::PrivateMethod { kind, .. } => kind,
         }
     }
 
@@ -181,6 +201,10 @@ impl FunctionQuery {
             FunctionQuery::ClassConstructor { .. } => "constructor",
             FunctionQuery::ClassMethod { method_name, .. }
             | FunctionQuery::ObjectMethod { method_name, .. } => method_name,
+            FunctionQuery::PrivateMethod {
+                private_method_name,
+                ..
+            } => private_method_name,
             FunctionQuery::FunctionDeclaration { function_name, .. } => function_name,
             FunctionQuery::FunctionExpression {
                 expression_name, ..
@@ -195,6 +219,7 @@ impl FunctionQuery {
             | FunctionQuery::ObjectMethod { .. } => FunctionType::Method,
             FunctionQuery::FunctionDeclaration { .. } => FunctionType::FunctionDeclaration,
             FunctionQuery::FunctionExpression { .. } => FunctionType::FunctionExpression,
+            FunctionQuery::PrivateMethod { .. } => FunctionType::PrivateMethod,
         }
     }
 
@@ -205,15 +230,17 @@ impl FunctionQuery {
             | FunctionQuery::ClassMethod { index, .. }
             | FunctionQuery::ObjectMethod { index, .. }
             | FunctionQuery::FunctionDeclaration { index, .. }
-            | FunctionQuery::FunctionExpression { index, .. } => *index,
+            | FunctionQuery::FunctionExpression { index, .. }
+            | FunctionQuery::PrivateMethod { index, .. } => *index,
         }
     }
 
     #[must_use]
     pub(crate) fn class_name(&self) -> Option<&str> {
         match self {
             FunctionQuery::ClassConstructor { class_name, .. }
-            | FunctionQuery::ClassMethod { class_name, .. } => Some(class_name),
+            | FunctionQuery::ClassMethod { class_name, .. }
+            | FunctionQuery::PrivateMethod { class_name, .. } => Some(class_name),
             _ => None,
         }
     }
@@ -240,7 +267,7 @@ impl FunctionQuery {
                 is_export_alias,
                 ..
             } => Some((expression_name, *is_export_alias)),
-            FunctionQuery::ObjectMethod { .. } => None,
+            FunctionQuery::ObjectMethod { .. } | FunctionQuery::PrivateMethod { .. } => None,
         }
     }
 
@@ -282,4 +309,10 @@ impl FunctionQuery {
             matches!(self.typ(), FunctionType::Method) && name == self.name();
         self.maybe_increment_count(matches_except_count, count)
     }
+
+    pub fn matches_private_method(&self, count: &mut usize, name: &str) -> bool {
+        let matches_except_count =
+            matches!(self.typ(), FunctionType::PrivateMethod) && name == self.name();
+        self.maybe_increment_count(matches_except_count, count)
+    }
 }
diff --git a/src/instrumentation.rs b/src/instrumentation.rs
@@ -10,7 +10,7 @@ use swc_core::ecma::{
     ast::{
         ArrowExpr, AssignExpr, AssignTarget, BlockStmt, ClassDecl, ClassExpr, ClassMethod,
         Constructor, Expr, FnDecl, FnExpr, Ident, Lit, MemberProp, MethodProp, Module, ModuleItem,
-        Param, Pat, PropName, Script, SimpleAssignTarget, Stmt, Str, VarDecl,
+        Param, Pat, PrivateMethod, PropName, Script, SimpleAssignTarget, Stmt, Str, VarDecl,
     },
     atoms::Atom,
 };
@@ -335,6 +335,27 @@ impl Instrumentation {
         true
     }
 
+    pub fn visit_mut_private_method(&mut self, node: &mut PrivateMethod) -> bool {
+        let name = node.key.name.clone();
+
+        // Only increment count when class matches
+        if !self.is_correct_class {
+            return true;
+        }
+
+        if self
+            .config
+            .function_query
+            .matches_private_method(&mut self.count, name.as_ref())
+            && node.function.body.is_some()
+        {
+            if let Some(body) = node.function.body.as_mut() {
+                self.insert_tracing(body, &node.function.params, node.function.is_async);
+            }
+        }
+        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;
diff --git a/src/lib.rs b/src/lib.rs
@@ -29,8 +29,8 @@ use swc_core::{
     ecma::{
         ast::{
             AssignExpr, ClassDecl, ClassExpr, ClassMethod, Constructor, EsVersion, ExportSpecifier,
-            FnDecl, MethodProp, Module, ModuleDecl, ModuleExportName, ModuleItem, Script, Str,
-            VarDecl,
+            FnDecl, MethodProp, Module, ModuleDecl, ModuleExportName, ModuleItem, PrivateMethod,
+            Script, Str, VarDecl,
         },
         visit::{VisitMut, VisitMutWith},
     },
@@ -352,5 +352,6 @@ impl VisitMut for InstrumentationVisitor {
     visit_with_all_fn!(visit_mut_class_decl, ClassDecl);
     visit_with_all_fn!(visit_mut_class_expr, ClassExpr);
     visit_with_all_fn!(visit_mut_class_method, ClassMethod);
+    visit_with_all_fn!(visit_mut_private_method, PrivateMethod);
     visit_with_all_fn!(visit_mut_constructor, Constructor);
 }
diff --git a/tests/instrumentor_test.rs b/tests/instrumentor_test.rs
@@ -24,3 +24,4 @@ mod nested_functions;
 mod object_method_cjs;
 mod polyfill_cjs;
 mod polyfill_mjs;
+mod private_method_cjs;
diff --git a/tests/private_method_cjs/mod.js b/tests/private_method_cjs/mod.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.
+ **/
+class TestClass {
+  constructor() {
+    this.prop = true
+  }
+
+   async #testMe() {
+     return 42
+   }
+
+  async testMe() {
+    return this.#testMe()
+  }
+}
+
+module.exports = TestClass 
diff --git a/tests/private_method_cjs/mod.rs b/tests/private_method_cjs/mod.rs
@@ -0,0 +1,15 @@
+use crate::common::*;
+use orchestrion_js::*;
+
+#[test]
+fn private_method_cjs() {
+    transpile_and_test(
+        file!(),
+        false,
+        Config::new_single(InstrumentationConfig::new(
+            "TestClass:testMe",
+            test_module_matcher(),
+            FunctionQuery::private_method("TestClass", "testMe", FunctionKind::Async),
+        )),
+    );
+}
diff --git a/tests/private_method_cjs/test.js b/tests/private_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 TestClass = require('./instrumented.js');
+const { assert, getContext } = require('../common/preamble.js');
+const context = getContext('orchestrion:undici:TestClass:testMe');
+(async () => {
+  const test = new TestClass;
+  const result = await test.testMe() 
+  assert.strictEqual(result, 42);
+  assert.deepStrictEqual(context, {
+    start: true,
+    end: true,
+    asyncStart: 42,
+    asyncEnd: 42
+  });
+})();
diff --git a/tests/wasm/__snapshots__/tests.test.mjs.snap b/tests/wasm/__snapshots__/tests.test.mjs.snap
@@ -3,6 +3,7 @@
 exports[`Orchestrion JS Transformer > should transform CommonJS module correctly 1`] = `
 {
   "code": "const { tracingChannel: tr_ch_apm_tracingChannel } = require("diagnostics_channel");
+const tr_ch_apm$up_privateSend = tr_ch_apm_tracingChannel("orchestrion:one:up:privateSend");
 const tr_ch_apm$up_asyncGet = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncGet");
 const tr_ch_apm$up_asyncFetch = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncFetch");
 const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
@@ -75,6 +76,20 @@ module.exports = class Up {
             moduleVersion: "1.0.0"
         });
     }
+    async #send() {
+        const __apm$original_args = arguments;
+        const __apm$traced = async ()=>{
+            const __apm$wrapped = async ()=>{};
+            return __apm$wrapped.apply(null, __apm$original_args);
+        };
+        if (!tr_ch_apm$up_privateSend.hasSubscribers) return __apm$traced();
+        return tr_ch_apm$up_privateSend.tracePromise(__apm$traced, {
+            arguments,
+            self: this,
+            moduleVersion: "1.0.0"
+        });
+    }
+    async send() {}
 };
 ",
   "map": undefined,
@@ -84,6 +99,7 @@ module.exports = class Up {
 exports[`Orchestrion JS Transformer > should transform ESM module correctly 1`] = `
 {
   "code": "import { tracingChannel as tr_ch_apm_tracingChannel } from "diagnostics_channel";
+const tr_ch_apm$up_privateSend = tr_ch_apm_tracingChannel("orchestrion:one:up:privateSend");
 const tr_ch_apm$up_asyncGet = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncGet");
 const tr_ch_apm$up_asyncFetch = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncFetch");
 const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
@@ -156,6 +172,20 @@ export class Up {
             moduleVersion: "1.0.0"
         });
     }
+    async #send() {
+        const __apm$original_args = arguments;
+        const __apm$traced = async ()=>{
+            const __apm$wrapped = async ()=>{};
+            return __apm$wrapped.apply(null, __apm$original_args);
+        };
+        if (!tr_ch_apm$up_privateSend.hasSubscribers) return __apm$traced();
+        return tr_ch_apm$up_privateSend.tracePromise(__apm$traced, {
+            arguments,
+            self: this,
+            moduleVersion: "1.0.0"
+        });
+    }
+    async send() {}
 }
 ",
   "map": undefined,
@@ -165,6 +195,7 @@ export class Up {
 exports[`Orchestrion JS Transformer > should transform TypeScript with source map correctly 1`] = `
 {
   "code": "import { tracingChannel as tr_ch_apm_tracingChannel } from "diagnostics_channel";
+const tr_ch_apm$up_privateSend = tr_ch_apm_tracingChannel("orchestrion:one:up:privateSend");
 const tr_ch_apm$up_asyncGet = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncGet");
 const tr_ch_apm$up_asyncFetch = tr_ch_apm_tracingChannel("orchestrion:one:up:asyncFetch");
 const tr_ch_apm$up_fetch = tr_ch_apm_tracingChannel("orchestrion:one:up:fetch");
@@ -237,8 +268,22 @@ export class Up {
             moduleVersion: "1.0.0"
         });
     }
+    async #send() {
+        const __apm$original_args = arguments;
+        const __apm$traced = async ()=>{
+            const __apm$wrapped = async ()=>{};
+            return __apm$wrapped.apply(null, __apm$original_args);
+        };
+        if (!tr_ch_apm$up_privateSend.hasSubscribers) return __apm$traced();
+        return tr_ch_apm$up_privateSend.tracePromise(__apm$traced, {
+            arguments,
+            self: this,
+            moduleVersion: "1.0.0"
+        });
+    }
+    async send() {}
 }
 ",
-  "map": "{"version":3,"file":"module.js","sources":["module.ts"],"sourceRoot":"","names":[],"mappings":";;;;;AAEA,MAAM,CAAA,MAAO,EAAE;IACX,aAAA;;;;;;;;;YACI,OAAO,CAAC,GAAG,CAAC,aAAa,CAAC,CAAC;;;;;;;;;;;;;;;;IAC/B,CAAC;IACD,KAAK,IAAS,EAAA;;;mCAAR;gBACF,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,CAAC;;;;;;;;;;IACzB,CAAC;IACD,KAAK,CAAC,UAAU,GAAA;;;;;;;;;;;;IAAU,CAAC;IAC3B,GAAG,GAAA;;;;;;;;;;;;IAAU,CAAC;CACjB"}",
+  "map": "{"version":3,"file":"module.js","sources":["module.ts"],"sourceRoot":"","names":[],"mappings":";;;;;;AAEA,MAAM,CAAA,MAAO,EAAE;IACX,aAAA;;;;;;;;;YACI,OAAO,CAAC,GAAG,CAAC,aAAa,CAAC,CAAC;;;;;;;;;;;;;;;;IAC/B,CAAC;IACD,KAAK,IAAS,EAAA;;;mCAAR;gBACF,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,CAAC;YACzB,CAAC;;;;;;;;;;IACD,KAAK,CAAC,UAAU,GAAA;;;;;;;;;;;;IAAU,CAAC;IAC3B,GAAG,GAAA;;;;;;;;;;;;IAAU,CAAC;IACd,KAAK,EAAC,IAAK;;;;;;;;;;;;IAAU,CAAC;IACtB,KAAK,CAAC,IAAI,GAAA,CAAU,CAAC;CACxB"}",
 }
 `;
diff --git a/tests/wasm/tests.test.mjs b/tests/wasm/tests.test.mjs
@@ -37,6 +37,15 @@ describe('Orchestrion JS Transformer', () => {
                 kind: "Async",
             },
         },
+        {
+            channelName: "up:privateSend",
+            module: { name: "one", versionRange: ">=1", filePath: "index.js" },
+            functionQuery: {
+                className: "Up",
+                privateMethodName: "send",
+                kind: "Async",
+            },
+        },
     ]);
 
     const matchedTransforms = instrumentor.getTransformer(
@@ -61,6 +70,8 @@ describe('Orchestrion JS Transformer', () => {
 
   async asyncFetch() {}
   get() {}
+  async #send() {}
+  async send() {}
 }`;
 
         const output = matchedTransforms.transform(originalEsm, 'esm');
@@ -79,6 +90,8 @@ describe('Orchestrion JS Transformer', () => {
 
   async asyncFetch() {}
   get() {}
+  async #send() {}
+  async send() {}
 }
 
 `;
@@ -98,6 +111,8 @@ export class Up {
     }
     async asyncFetch(): void {}
     get(): void {}
+    async #send(): void {}
+    async send(): void {}
 }`;
 
         const { outputText: outputJavaScript, sourceMapText: originalTypescriptSourceMap } = tsc.transpileModule(originalTypescript, {
@@ -125,7 +140,7 @@ export class Up {
         });
 
         // This is the position of the fetch function in the original TypeScript
-        expect(originalPosition.line).toEqual(7);
+        expect(originalPosition.line).toEqual(6);
         expect(originalPosition.column).toEqual(4);
 
         sourceMapConsumer.destroy();
@@ -142,10 +157,12 @@ export class Up {
 	}
   async asyncFetch() {}
   get() {}
+  async #send() {}
+  async send() {}
 }`;
 
         expect(() => {
             matchedTransforms.transform(noMatchSource, 'unknown');
-        }).toThrow('Failed to find injection points for: ["constructor", "fetch", "asyncFetch", "get"]');
+        }).toThrow('Failed to find injection points for: ["constructor", "fetch", "asyncFetch", "get", "send"]');
     });
 });
