Add option to instrument rpc methods

Author: andreiborza

packages/cloudflare/src/client.ts

@@ -68,6 +68,29 @@ interface BaseCloudflareOptions {
    * @default false
    */
   skipOpenTelemetrySetup?: boolean;
+
+  /**
+   * Enable instrumentation of prototype methods for DurableObjects.
+   *
+   * When `true`, the SDK will wrap all methods on the DurableObject prototype chain
+   * to automatically create spans and capture errors for RPC method calls.
+   *
+   * When an array of strings is provided, only the specified method names will be instrumented.
+   *
+   * This feature adds runtime overhead as it wraps methods at the prototype level.
+   * Only enable this if you need automatic instrumentation of prototype methods.
+   *
+   * @default false
+   * @example
+   * ```ts
+   * // Instrument all prototype methods
+   * instrumentPrototypeMethods: true
+   *
+   * // Instrument only specific methods
+   * instrumentPrototypeMethods: ['myMethod', 'anotherMethod']
+   * ```
+   */
+  instrumentPrototypeMethods?: boolean | string[];
 }
 
 /**

packages/cloudflare/src/durableobject.ts

@@ -288,14 +288,16 @@ export function instrumentDurableObjectWithSentry<
         configurable: false,
       });
 
-      instrumentPrototype(target);
+      if (options?.instrumentPrototypeMethods) {
+        instrumentPrototype(target, options.instrumentPrototypeMethods);
+      }
 
       return obj;
     },
   });
 }
 
-function instrumentPrototype<T extends NewableFunction>(target: T): void {
+function instrumentPrototype<T extends NewableFunction>(target: T, methodsToInstrument: boolean | string[]): void {
   const proto = target.prototype;
 
   // Get all methods from the prototype chain
@@ -311,6 +313,9 @@ function instrumentPrototype<T extends NewableFunction>(target: T): void {
     current = Object.getPrototypeOf(current);
   }
 
+  // Create a set for efficient lookups when methodsToInstrument is an array
+  const methodsToInstrumentSet = Array.isArray(methodsToInstrument) ? new Set(methodsToInstrument) : null;
+
   // Instrument each method on the prototype
   methodNames.forEach(methodName => {
     const originalMethod = (proto as Record<string, unknown>)[methodName];
@@ -319,6 +324,11 @@ function instrumentPrototype<T extends NewableFunction>(target: T): void {
       return;
     }
 
+    // If methodsToInstrument is an array, only instrument methods in that set
+    if (methodsToInstrumentSet && !methodsToInstrumentSet.has(methodName)) {
+      return;
+    }
+
     // Create a wrapper that gets context/options from the instance at runtime
     const wrappedMethod = function (this: any, ...args: any[]): unknown {
       const thisWithSentry = this as {

packages/cloudflare/test/durableobject.test.ts

@@ -10,7 +10,7 @@ describe('instrumentDurableObjectWithSentry', () => {
   });
 
   it('Generic functionality', () => {
-    const options = vi.fn();
+    const options = vi.fn().mockReturnValue({});
     const instrumented = instrumentDurableObjectWithSentry(options, vi.fn());
     expect(instrumented).toBeTypeOf('function');
     expect(() => Reflect.construct(instrumented, [])).not.toThrow();
@@ -23,7 +23,7 @@ describe('instrumentDurableObjectWithSentry', () => {
         return 'sync-result';
       }
     };
-    const obj = Reflect.construct(instrumentDurableObjectWithSentry(vi.fn(), testClass as any), []) as any;
+    const obj = Reflect.construct(instrumentDurableObjectWithSentry(vi.fn().mockReturnValue({}), testClass as any), []) as any;
     expect(obj.method).toBe(obj.method);
 
     const result = obj.method();
@@ -37,7 +37,7 @@ describe('instrumentDurableObjectWithSentry', () => {
         return 'async-result';
       }
     };
-    const obj = Reflect.construct(instrumentDurableObjectWithSentry(vi.fn(), testClass as any), []) as any;
+    const obj = Reflect.construct(instrumentDurableObjectWithSentry(vi.fn().mockReturnValue({}), testClass as any), []) as any;
     expect(obj.asyncMethod).toBe(obj.asyncMethod);
 
     const result = obj.asyncMethod();
@@ -56,9 +56,11 @@ describe('instrumentDurableObjectWithSentry', () => {
       .fn()
       .mockReturnValueOnce({
         orgId: 1,
+        instrumentPrototypeMethods: true,
       })
       .mockReturnValueOnce({
         orgId: 2,
+        instrumentPrototypeMethods: true,
       });
     const testClass = class {
       method() {}
@@ -79,7 +81,7 @@ describe('instrumentDurableObjectWithSentry', () => {
     expect(initCore).nthCalledWith(2, expect.any(Function), expect.objectContaining({ orgId: 2 }));
   });
 
-  it('All available durable object methods are instrumented', () => {
+  it('All available durable object methods are instrumented when instrumentPrototypeMethods is enabled', () => {
     const testClass = class {
       propertyFunction = vi.fn();
 
@@ -95,7 +97,7 @@ describe('instrumentDurableObjectWithSentry', () => {
 
       webSocketError() {}
     };
-    const instrumented = instrumentDurableObjectWithSentry(vi.fn(), testClass as any);
+    const instrumented = instrumentDurableObjectWithSentry(vi.fn().mockReturnValue({ instrumentPrototypeMethods: true }), testClass as any);
     const obj = Reflect.construct(instrumented, []);
     for (const method_name of [
       'propertyFunction',
@@ -135,4 +137,93 @@ describe('instrumentDurableObjectWithSentry', () => {
     await Promise.all(waitUntil.mock.calls.map(([p]) => p));
     expect(flush).toBeCalled();
   });
+
+  describe('instrumentPrototypeMethods option', () => {
+    it('does not instrument prototype methods when option is not set', () => {
+      const testClass = class {
+        prototypeMethod() {
+          return 'prototype-result';
+        }
+      };
+      const options = vi.fn().mockReturnValue({});
+      const instrumented = instrumentDurableObjectWithSentry(options, testClass as any);
+      const obj = Reflect.construct(instrumented, []) as any;
+
+      expect(isInstrumented(obj.prototypeMethod)).toBeFalsy();
+    });
+
+    it('does not instrument prototype methods when option is false', () => {
+      const testClass = class {
+        prototypeMethod() {
+          return 'prototype-result';
+        }
+      };
+      const options = vi.fn().mockReturnValue({ instrumentPrototypeMethods: false });
+      const instrumented = instrumentDurableObjectWithSentry(options, testClass as any);
+      const obj = Reflect.construct(instrumented, []) as any;
+
+      expect(isInstrumented(obj.prototypeMethod)).toBeFalsy();
+    });
+
+    it('instruments all prototype methods when option is true', () => {
+      const testClass = class {
+        methodOne() {
+          return 'one';
+        }
+        methodTwo() {
+          return 'two';
+        }
+      };
+      const options = vi.fn().mockReturnValue({ instrumentPrototypeMethods: true });
+      const instrumented = instrumentDurableObjectWithSentry(options, testClass as any);
+      const obj = Reflect.construct(instrumented, []) as any;
+
+      expect(isInstrumented(obj.methodOne)).toBeTruthy();
+      expect(isInstrumented(obj.methodTwo)).toBeTruthy();
+    });
+
+    it('instruments only specified methods when option is array', () => {
+      const testClass = class {
+        methodOne() {
+          return 'one';
+        }
+        methodTwo() {
+          return 'two';
+        }
+        methodThree() {
+          return 'three';
+        }
+      };
+      const options = vi.fn().mockReturnValue({ instrumentPrototypeMethods: ['methodOne', 'methodThree'] });
+      const instrumented = instrumentDurableObjectWithSentry(options, testClass as any);
+      const obj = Reflect.construct(instrumented, []) as any;
+
+      expect(isInstrumented(obj.methodOne)).toBeTruthy();
+      expect(isInstrumented(obj.methodTwo)).toBeFalsy();
+      expect(isInstrumented(obj.methodThree)).toBeTruthy();
+    });
+
+    it('still instruments instance methods regardless of prototype option', () => {
+      const testClass = class {
+        propertyFunction = vi.fn();
+
+        fetch() {}
+        alarm() {}
+        webSocketMessage() {}
+        webSocketClose() {}
+        webSocketError() {}
+      };
+      const options = vi.fn().mockReturnValue({ instrumentPrototypeMethods: false });
+      const instrumented = instrumentDurableObjectWithSentry(options, testClass as any);
+      const obj = Reflect.construct(instrumented, []) as any;
+
+      // Instance methods should still be instrumented
+      expect(isInstrumented(obj.propertyFunction)).toBeTruthy();
+      expect(isInstrumented(obj.fetch)).toBeTruthy();
+      expect(isInstrumented(obj.alarm)).toBeTruthy();
+      expect(isInstrumented(obj.webSocketMessage)).toBeTruthy();
+      expect(isInstrumented(obj.webSocketClose)).toBeTruthy();
+      expect(isInstrumented(obj.webSocketError)).toBeTruthy();
+    });
+  });
 });