durableobject.ts

/* eslint-disable @typescript-eslint/unbound-method */
import { captureException } from '@sentry/core';
import type { DurableObject } from 'cloudflare:workers';
import { setAsyncLocalStorageAsyncContextStrategy } from './async';
import type { CloudflareOptions } from './client';
import { ensureInstrumented, getInstrumented, markAsInstrumented } from './instrument';
import { instrumentEnv } from './instrumentations/worker/instrumentEnv';
import { getFinalOptions } from './options';
import { wrapRequestHandler } from './request';
import { instrumentContext } from './utils/instrumentContext';
import type { UncheckedMethod } from './wrapMethodWithSentry';
import { wrapMethodWithSentry } from './wrapMethodWithSentry';

/**
 * Instruments a Durable Object class to capture errors and performance data.
 *
 * Instruments the following methods:
 * - fetch
 * - alarm
 * - webSocketMessage
 * - webSocketClose
 * - webSocketError
 *
 * as well as any other public RPC methods on the Durable Object instance.
 *
 * @param optionsCallback Function that returns the options for the SDK initialization.
 * @param DurableObjectClass The Durable Object class to instrument.
 * @returns The instrumented Durable Object class.
 *
 * @example
 * ```ts
 * class MyDurableObjectBase extends DurableObject {
 *   constructor(ctx: DurableObjectState, env: Env) {
 *     super(ctx, env);
 *   }
 * }
 *
 * export const MyDurableObject = instrumentDurableObjectWithSentry(
 *   env => ({
 *     dsn: env.SENTRY_DSN,
 *     tracesSampleRate: 1.0,
 *   }),
 *   MyDurableObjectBase,
 * );
 * ```
 */
export function instrumentDurableObjectWithSentry<
  E,
  T extends DurableObject<E>,
  C extends new (state: DurableObjectState, env: E) => T,
>(optionsCallback: (env: E) => CloudflareOptions, DurableObjectClass: C): C {
  return new Proxy(DurableObjectClass, {
    construct(target, [ctx, env]) {
      setAsyncLocalStorageAsyncContextStrategy();
      const context = instrumentContext(ctx);
      const options = getFinalOptions(optionsCallback(env), env);
      const instrumentedEnv = instrumentEnv(env);

      const obj = new target(context, instrumentedEnv);

      // These are the methods that are available on a Durable Object
      // ref: https://developers.cloudflare.com/durable-objects/api/base/
      // obj.alarm
      // obj.fetch
      // obj.webSocketError
      // obj.webSocketClose
      // obj.webSocketMessage

      // Any other public methods on the Durable Object instance are RPC calls.

      if (obj.fetch && typeof obj.fetch === 'function') {
        obj.fetch = ensureInstrumented(
          obj.fetch,
          original =>
            new Proxy(original, {
              apply(target, thisArg, args) {
                return wrapRequestHandler({ options, request: args[0], context }, () => {
                  return Reflect.apply(target, thisArg, args);
                });
              },
            }),
        );
      }

      if (obj.alarm && typeof obj.alarm === 'function') {
        // Alarms are independent invocations, so we start a new trace and link to the previous alarm
        obj.alarm = wrapMethodWithSentry(
          { options, context, spanName: 'alarm', spanOp: 'function', startNewTrace: true },
          obj.alarm,
        );
      }

      if (obj.webSocketMessage && typeof obj.webSocketMessage === 'function') {
        obj.webSocketMessage = wrapMethodWithSentry(
          { options, context, spanName: 'webSocketMessage' },
          obj.webSocketMessage,
        );
      }

      if (obj.webSocketClose && typeof obj.webSocketClose === 'function') {
        obj.webSocketClose = wrapMethodWithSentry({ options, context, spanName: 'webSocketClose' }, obj.webSocketClose);
      }

      if (obj.webSocketError && typeof obj.webSocketError === 'function') {
        obj.webSocketError = wrapMethodWithSentry(
          { options, context, spanName: 'webSocketError' },
          obj.webSocketError,
          (_, error) =>
            captureException(error, {
              mechanism: {
                type: 'auto.faas.cloudflare.durable_object_websocket',
                handled: false,
              },
            }),
        );
      }

      for (const method of Object.getOwnPropertyNames(obj)) {
        if (
          method === 'fetch' ||
          method === 'alarm' ||
          method === 'webSocketError' ||
          method === 'webSocketClose' ||
          method === 'webSocketMessage'
        ) {
          continue;
        }

        // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
        const value = (obj as any)[method] as unknown;
        if (typeof value === 'function') {
          // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
          (obj as any)[method] = wrapMethodWithSentry(
            { options, context, spanName: method, spanOp: 'rpc' },
            value as UncheckedMethod,
          );
        }
      }

      // Store context and options on the instance for prototype methods to access
      Object.defineProperty(obj, '__SENTRY_CONTEXT__', {
        value: context,
        enumerable: false,
        writable: false,
        configurable: false,
      });

      Object.defineProperty(obj, '__SENTRY_OPTIONS__', {
        value: options,
        enumerable: false,
        writable: false,
        configurable: false,
      });

      if (options?.instrumentPrototypeMethods) {
        instrumentPrototype(target, options.instrumentPrototypeMethods);
      }

      return obj;
    },
  });
}

function instrumentPrototype<T extends NewableFunction>(target: T, methodsToInstrument: boolean | string[]): void {
  const proto = target.prototype;

  // Get all methods from the prototype chain
  const methodNames = new Set<string>();
  let current = proto;

  while (current && current !== Object.prototype) {
    Object.getOwnPropertyNames(current).forEach(name => {
      if (name !== 'constructor' && typeof (current as Record<string, unknown>)[name] === 'function') {
        methodNames.add(name);
      }
    });
    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];

    if (!originalMethod) {
      return;
    }

    const existingInstrumented = getInstrumented(originalMethod);
    if (existingInstrumented) {
      Object.defineProperty(proto, methodName, {
        value: existingInstrumented,
        enumerable: false,
        writable: true,
        configurable: true,
      });
      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: unknown, ...args: unknown[]): unknown {
      const thisWithSentry = this as {
        __SENTRY_CONTEXT__: DurableObjectState;
        __SENTRY_OPTIONS__: CloudflareOptions;
      };
      const instanceContext = thisWithSentry.__SENTRY_CONTEXT__;
      const instanceOptions = thisWithSentry.__SENTRY_OPTIONS__;

      if (!instanceOptions) {
        // Fallback to original method if no Sentry data found
        return (originalMethod as UncheckedMethod).apply(this, args);
      }

      // Use the existing wrapper but with instance-specific context/options
      const wrapper = wrapMethodWithSentry(
        {
          options: instanceOptions,
          context: instanceContext,
          spanName: methodName,
          spanOp: 'rpc',
        },
        originalMethod as UncheckedMethod,
        undefined,
        true, // noMark = true since we'll mark the prototype method
      );

      return wrapper.apply(this, args);
    };

    // Only mark wrappedMethod as instrumented (not originalMethod → wrappedMethod).
    // originalMethod must stay unmapped because wrappedMethod calls
    // wrapMethodWithSentry(options, originalMethod) on each invocation to create
    // a per-instance proxy. If originalMethod mapped to wrappedMethod, that call
    // would return wrappedMethod itself, causing infinite recursion.
    markAsInstrumented(wrappedMethod);

    // Replace the prototype method
    Object.defineProperty(proto, methodName, {
      value: wrappedMethod,
      enumerable: false,
      writable: true,
      configurable: true,
    });
  });
}