feat(cloudflare): Instrument async KV API

JPeer264 · JPeer264 · commit 0eb536eb0d39 · 2026-02-19T09:49:51.000-01:00
diff --git a/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts b/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts
@@ -20,17 +20,26 @@ class MyDurableObjectBase extends DurableObject<Env> {
   }
 
   async fetch(request: Request) {
-    const { pathname } = new URL(request.url);
-    switch (pathname) {
+    const url = new URL(request.url);
+    switch (url.pathname) {
       case '/throwException': {
         await this.throwException();
         break;
       }
-      case '/ws':
+      case '/ws': {
         const webSocketPair = new WebSocketPair();
         const [client, server] = Object.values(webSocketPair);
         this.ctx.acceptWebSocket(server);
         return new Response(null, { status: 101, webSocket: client });
+      }
+      case '/storage/put': {
+        await this.ctx.storage.put('test-key', 'test-value');
+        return new Response('Stored');
+      }
+      case '/storage/get': {
+        const value = await this.ctx.storage.get('test-key');
+        return new Response(`Got: ${value}`);
+      }
     }
     return new Response('DO is fine');
   }
diff --git a/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts b/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts
@@ -1,5 +1,5 @@
 import { expect, test } from '@playwright/test';
-import { waitForError, waitForRequest } from '@sentry-internal/test-utils';
+import { waitForError, waitForRequest, waitForTransaction } from '@sentry-internal/test-utils';
 import { SDK_VERSION } from '@sentry/cloudflare';
 import { WebSocket } from 'ws';
 
@@ -82,3 +82,20 @@ test('sends user-agent header with SDK name and version in envelope requests', a
     'user-agent': `sentry.javascript.cloudflare/${SDK_VERSION}`,
   });
 });
+
+test('Storage operations create spans in Durable Object transactions', async ({ baseURL }) => {
+  const transactionWaiter = waitForTransaction('cloudflare-workers', event => {
+    return event.spans?.some(span => span.op === 'db' && span.description === 'durable_object_storage_put') ?? false;
+  });
+
+  const response = await fetch(`${baseURL}/pass-to-object/storage/put`);
+  expect(response.status).toBe(200);
+
+  const transaction = await transactionWaiter;
+  const putSpan = transaction.spans?.find(span => span.description === 'durable_object_storage_put');
+
+  expect(putSpan).toBeDefined();
+  expect(putSpan?.op).toBe('db');
+  expect(putSpan?.data?.['db.system.name']).toBe('cloudflare.durable_object.storage');
+  expect(putSpan?.data?.['db.operation.name']).toBe('put');
+});
diff --git a/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts
@@ -0,0 +1,50 @@
+import type { DurableObjectStorage } from '@cloudflare/workers-types';
+import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, startSpan } from '@sentry/core';
+
+const STORAGE_METHODS_TO_INSTRUMENT = ['get', 'put', 'delete', 'list'] as const;
+
+type StorageMethod = (typeof STORAGE_METHODS_TO_INSTRUMENT)[number];
+
+/**
+ * Instruments DurableObjectStorage methods with Sentry spans.
+ *
+ * Wraps the following async methods:
+ * - get, put, delete, list (KV API)
+ *
+ * @param storage - The DurableObjectStorage instance to instrument
+ * @returns An instrumented DurableObjectStorage instance
+ */
+export function instrumentDurableObjectStorage(storage: DurableObjectStorage): DurableObjectStorage {
+  return new Proxy(storage, {
+    get(target, prop, receiver) {
+      const original = Reflect.get(target, prop, receiver);
+
+      if (typeof original !== 'function') {
+        return original;
+      }
+
+      const methodName = prop as string;
+      if (!STORAGE_METHODS_TO_INSTRUMENT.includes(methodName as StorageMethod)) {
+        return (original as (...args: unknown[]) => unknown).bind(target);
+      }
+
+      return function (this: unknown, ...args: unknown[]) {
+        return startSpan(
+          {
+            // Use underscore naming to match Cloudflare's native instrumentation (e.g., "durable_object_storage_get")
+            name: `durable_object_storage_${methodName}`,
+            op: 'db',
+            attributes: {
+              [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.durable_object',
+              'db.system.name': 'cloudflare.durable_object.storage',
+              'db.operation.name': methodName,
+            },
+          },
+          () => {
+            return (original as (...args: unknown[]) => unknown).apply(target, args);
+          },
+        );
+      };
+    },
+  });
+}
diff --git a/packages/cloudflare/src/utils/instrumentContext.ts b/packages/cloudflare/src/utils/instrumentContext.ts
@@ -1,4 +1,5 @@
-import { type DurableObjectState, type ExecutionContext } from '@cloudflare/workers-types';
+import { type DurableObjectState, type DurableObjectStorage, type ExecutionContext } from '@cloudflare/workers-types';
+import { instrumentDurableObjectStorage } from '../instrumentations/instrumentDurableObjectStorage';
 
 type ContextType = ExecutionContext | DurableObjectState;
 type OverridesStore<T extends ContextType> = Map<keyof T, (...args: unknown[]) => unknown>;
@@ -8,6 +9,8 @@ type OverridesStore<T extends ContextType> = Map<keyof T, (...args: unknown[]) =
  *
  * Creates a copy of the context that:
  * - Allows overriding of methods (e.g., waitUntil)
+ * - For DurableObjectState: instruments storage operations (get, put, delete, list, etc.)
+ *   to create Sentry spans automatically
  *
  * @param ctx - The execution context or DurableObjectState to instrument
  * @returns An instrumented copy of the context
@@ -34,6 +37,30 @@ export function instrumentContext<T extends ContextType>(ctx: T): T {
     {} as PropertyDescriptorMap,
   );
 
+  // Check if this is a DurableObjectState context with a storage property
+  // If so, wrap the storage with instrumentation
+  if ('storage' in ctx && ctx.storage) {
+    const originalStorage = ctx.storage;
+    let instrumentedStorage: DurableObjectStorage | undefined;
+    descriptors.storage = {
+      configurable: true,
+      enumerable: true,
+      get: () => {
+        if (!instrumentedStorage) {
+          instrumentedStorage = instrumentDurableObjectStorage(originalStorage);
+        }
+        return instrumentedStorage;
+      },
+    };
+    // Expose the original uninstrumented storage for internal Sentry operations
+    // This avoids creating spans for internal storage operations
+    descriptors.originalStorage = {
+      configurable: true,
+      enumerable: false,
+      get: () => originalStorage,
+    };
+  }
+
   return Object.create(ctx, descriptors);
 }
 
diff --git a/packages/cloudflare/src/wrapMethodWithSentry.ts b/packages/cloudflare/src/wrapMethodWithSentry.ts
@@ -1,3 +1,4 @@
+import type { DurableObjectStorage } from '@cloudflare/workers-types';
 import {
   captureException,
   flush,
@@ -14,6 +15,11 @@ import type { CloudflareOptions } from './client';
 import { isInstrumented, markAsInstrumented } from './instrument';
 import { init } from './sdk';
 
+/** Extended DurableObjectState with originalStorage exposed by instrumentContext */
+interface InstrumentedDurableObjectState extends DurableObjectState {
+  originalStorage?: DurableObjectStorage;
+}
+
 type MethodWrapperOptions = {
   spanName?: string;
   spanOp?: string;
@@ -58,13 +64,13 @@ export function wrapMethodWithSentry<T extends OriginalMethod>(
         // In certain situations, the passed context can become undefined.
         // For example, for Astro while prerendering pages at build time.
         // see: https://github.com/getsentry/sentry-javascript/issues/13217
-        const context = wrapperOptions.context as ExecutionContext | undefined;
+        const context = wrapperOptions.context as InstrumentedDurableObjectState | undefined;
 
         const waitUntil = context?.waitUntil?.bind?.(context);
 
         const currentClient = scope.getClient();
         if (!currentClient) {
-          const client = init({ ...wrapperOptions.options, ctx: context });
+          const client = init({ ...wrapperOptions.options, ctx: context as unknown as ExecutionContext | undefined });
           scope.setClient(client);
         }
 
diff --git a/packages/cloudflare/test/instrumentContext.test.ts b/packages/cloudflare/test/instrumentContext.test.ts
@@ -46,6 +46,57 @@ describe('instrumentContext', () => {
     const instrumented = instrumentContext(context);
     expect(instrumented[s]).toBe(context[s]);
   });
+
+  describe('DurableObjectState storage instrumentation', () => {
+    it('instruments storage property', () => {
+      const mockStorage = createMockStorage();
+      const context = makeDurableObjectStateMock(mockStorage);
+      const instrumented = instrumentContext(context);
+
+      // The storage property should be instrumented (wrapped)
+      expect(instrumented.storage).toBeDefined();
+      // The instrumented storage should not be the same reference
+      expect(instrumented.storage).not.toBe(mockStorage);
+    });
+
+    it('exposes originalStorage as the uninstrumented storage', () => {
+      const mockStorage = createMockStorage();
+      const context = makeDurableObjectStateMock(mockStorage);
+      const instrumented = instrumentContext(context);
+
+      // originalStorage should be the original uninstrumented storage
+      expect(instrumented.originalStorage).toBe(mockStorage);
+    });
+
+    it('originalStorage is not enumerable', () => {
+      const mockStorage = createMockStorage();
+      const context = makeDurableObjectStateMock(mockStorage);
+      const instrumented = instrumentContext(context);
+
+      // originalStorage should not appear in Object.keys
+      expect(Object.keys(instrumented)).not.toContain('originalStorage');
+    });
+
+    it('returns instrumented storage lazily', () => {
+      const mockStorage = createMockStorage();
+      const context = makeDurableObjectStateMock(mockStorage);
+      const instrumented = instrumentContext(context);
+
+      // Access storage twice to ensure memoization
+      const storage1 = instrumented.storage;
+      const storage2 = instrumented.storage;
+
+      expect(storage1).toBe(storage2);
+    });
+
+    it('handles context without storage property', () => {
+      const context = makeExecutionContextMock();
+      const instrumented = instrumentContext(context) as any;
+
+      // Should not have originalStorage if no storage property
+      expect(instrumented.originalStorage).toBeUndefined();
+    });
+  });
 });
 
 function makeExecutionContextMock<T extends ExecutionContext>() {
@@ -54,3 +105,39 @@ function makeExecutionContextMock<T extends ExecutionContext>() {
     passThroughOnException: vi.fn(),
   } as unknown as Mocked<T>;
 }
+
+function makeDurableObjectStateMock(storage?: any) {
+  return {
+    waitUntil: vi.fn(),
+    blockConcurrencyWhile: vi.fn(),
+    id: { toString: () => 'test-id', equals: vi.fn(), name: 'test' },
+    storage: storage || createMockStorage(),
+    acceptWebSocket: vi.fn(),
+    getWebSockets: vi.fn().mockReturnValue([]),
+    setWebSocketAutoResponse: vi.fn(),
+    getWebSocketAutoResponse: vi.fn(),
+    getWebSocketAutoResponseTimestamp: vi.fn(),
+    setHibernatableWebSocketEventTimeout: vi.fn(),
+    getHibernatableWebSocketEventTimeout: vi.fn(),
+    getTags: vi.fn().mockReturnValue([]),
+    abort: vi.fn(),
+  } as any;
+}
+
+function createMockStorage(): any {
+  return {
+    get: vi.fn().mockResolvedValue(undefined),
+    put: vi.fn().mockResolvedValue(undefined),
+    delete: vi.fn().mockResolvedValue(false),
+    list: vi.fn().mockResolvedValue(new Map()),
+    getAlarm: vi.fn().mockResolvedValue(null),
+    setAlarm: vi.fn().mockResolvedValue(undefined),
+    deleteAlarm: vi.fn().mockResolvedValue(undefined),
+    deleteAll: vi.fn().mockResolvedValue(undefined),
+    sync: vi.fn().mockResolvedValue(undefined),
+    transaction: vi.fn().mockImplementation(async (cb: () => unknown) => cb()),
+    sql: {
+      exec: vi.fn(),
+    },
+  };
+}
diff --git a/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts b/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts
diff --git a/packages/cloudflare/test/wrapMethodWithSentry.test.ts b/packages/cloudflare/test/wrapMethodWithSentry.test.ts
