From 40d26d2216dfd1c16b0157f9b34905cc2d0cd076 Mon Sep 17 00:00:00 2001
From: Abhijeet Prasad <aprasad@sentry.io>
Date: Tue, 11 Apr 2023 12:03:40 +0200
Subject: [PATCH] feat(browser): Add `setBrowserErrorFrameAsyncContextStrategy`

---
 packages/browser/src/async.ts            |  53 ++++++++++
 packages/browser/test/unit/async.test.ts | 124 +++++++++++++++++++++++
 2 files changed, 177 insertions(+)
 create mode 100644 packages/browser/src/async.ts
 create mode 100644 packages/browser/test/unit/async.test.ts

diff --git a/packages/browser/src/async.ts b/packages/browser/src/async.ts
new file mode 100644
index 000000000000..bd65ac811275
--- /dev/null
+++ b/packages/browser/src/async.ts
@@ -0,0 +1,53 @@
+import type { Carrier, Hub, RunWithAsyncContextOptions } from '@sentry/core';
+import { ensureHubOnCarrier, getHubFromCarrier, setAsyncContextStrategy } from '@sentry/core';
+
+/** */
+export function setBrowserErrorFrameAsyncContextStrategy(): void {
+  let id = 0;
+  const hubs = new Map<number, Hub>();
+
+  /** */
+  function getCurrentHub(): Hub | undefined {
+    const stackId = _getHubIdFromStack();
+    return stackId === undefined ? undefined : hubs.get(stackId);
+  }
+
+  /** */
+  function createNewHub(parent: Hub | undefined): Hub {
+    const carrier: Carrier = {};
+    ensureHubOnCarrier(carrier, parent);
+    return getHubFromCarrier(carrier);
+  }
+
+  /** */
+  function runWithAsyncContext<T>(callback: (hub: Hub) => T, options: RunWithAsyncContextOptions): T {
+    const existingHub = getCurrentHub();
+
+    if (existingHub && options && options.reuseExisting) {
+      // We're already in an async context, so we don't need to create a new one
+      // just call the callback with the current hub
+      return callback(existingHub);
+    }
+
+    const newHub = createNewHub(existingHub);
+
+    const hubId = id++;
+    const fnName = `SENTRY_HUB_ID_${hubId}`;
+
+    return {
+      [fnName]: (cb: (hub: Hub) => T) => {
+        hubs.set(hubId, newHub);
+        return cb(newHub);
+      },
+    }[fnName](callback);
+  }
+
+  setAsyncContextStrategy({ getCurrentHub, runWithAsyncContext });
+}
+
+function _getHubIdFromStack(): number | undefined {
+  const e = new Error();
+  const key = (e.stack && e.stack.match(/(?<=SENTRY_HUB_ID_)(?:\d+)/)) || [];
+  const value = Number.parseInt(key[0], 10);
+  return Number.isNaN(value) ? undefined : value;
+}
diff --git a/packages/browser/test/unit/async.test.ts b/packages/browser/test/unit/async.test.ts
new file mode 100644
index 000000000000..1234a14729f7
--- /dev/null
+++ b/packages/browser/test/unit/async.test.ts
@@ -0,0 +1,124 @@
+import { getCurrentHub, Hub, runWithAsyncContext, setAsyncContextStrategy } from '@sentry/core';
+
+import { setBrowserErrorFrameAsyncContextStrategy } from '../../src/async';
+
+describe('async browser context', () => {
+  afterAll(() => {
+    // clear the strategy
+    setAsyncContextStrategy(undefined);
+  });
+
+  test('without context', () => {
+    const hub = getCurrentHub();
+    expect(hub).toEqual(new Hub());
+  });
+
+  test('without strategy hubs should be equal', () => {
+    runWithAsyncContext(hub1 => {
+      runWithAsyncContext(hub2 => {
+        expect(hub1).toBe(hub2);
+      });
+    });
+  });
+
+  test('hub scope inheritance', () => {
+    setBrowserErrorFrameAsyncContextStrategy();
+
+    const globalHub = getCurrentHub();
+    globalHub.setExtra('a', 'b');
+
+    runWithAsyncContext(hub1 => {
+      expect(hub1).toEqual(globalHub);
+
+      hub1.setExtra('c', 'd');
+      expect(hub1).not.toEqual(globalHub);
+
+      runWithAsyncContext(hub2 => {
+        expect(hub2).toEqual(hub1);
+        expect(hub2).not.toEqual(globalHub);
+
+        hub2.setExtra('e', 'f');
+        expect(hub2).not.toEqual(hub1);
+      });
+    });
+  });
+
+  test('hub scope getCurrentHub', () => {
+    setBrowserErrorFrameAsyncContextStrategy();
+
+    const globalHub = getCurrentHub();
+    globalHub.setExtra('a', 'b');
+
+    runWithAsyncContext(hub1 => {
+      expect(getCurrentHub()).toBe(hub1);
+      runWithAsyncContext(hub2 => {
+        expect(getCurrentHub()).toBe(hub2);
+        runWithAsyncContext(hub3 => {
+          expect(getCurrentHub()).toBe(hub3);
+        });
+      });
+    });
+  });
+
+  test('context single instance', () => {
+    setBrowserErrorFrameAsyncContextStrategy();
+
+    runWithAsyncContext(hub => {
+      expect(hub).toBe(getCurrentHub());
+    });
+  });
+
+  test('context within a context not reused', () => {
+    setBrowserErrorFrameAsyncContextStrategy();
+
+    runWithAsyncContext(hub1 => {
+      runWithAsyncContext(hub2 => {
+        expect(hub1).not.toBe(hub2);
+      });
+    });
+  });
+
+  test('context within a context reused when requested', () => {
+    setBrowserErrorFrameAsyncContextStrategy();
+
+    runWithAsyncContext(hub1 => {
+      runWithAsyncContext(
+        hub2 => {
+          expect(hub1).toBe(hub2);
+        },
+        { reuseExisting: true },
+      );
+    });
+  });
+
+  test('concurrent hub contexts', done => {
+    setBrowserErrorFrameAsyncContextStrategy();
+
+    let d1done = false;
+    let d2done = false;
+
+    runWithAsyncContext(hub => {
+      hub.getStack().push({ client: 'process' } as any);
+      expect(hub.getStack()[1]).toEqual({ client: 'process' });
+      // Just in case so we don't have to worry which one finishes first
+      // (although it always should be d2)
+      setTimeout(() => {
+        d1done = true;
+        if (d2done) {
+          done();
+        }
+      });
+    });
+
+    runWithAsyncContext(hub => {
+      hub.getStack().push({ client: 'local' } as any);
+      expect(hub.getStack()[1]).toEqual({ client: 'local' });
+      setTimeout(() => {
+        d2done = true;
+        if (d1done) {
+          done();
+        }
+      });
+    });
+  });
+});