fix(core): Ensure withScope sets current scope correctly with async callbacks (#9974)

Oops, I noticed that our current `withScope` implementation does not
actually wait for async callbacks for setting the current scope 😬
Related to #9958,
which I copied there now!

Author: mydea

packages/core/src/hub.ts

@@ -20,7 +20,15 @@ import type {
   TransactionContext,
   User,
 } from '@sentry/types';
-import { GLOBAL_OBJ, consoleSandbox, dateTimestampInSeconds, getGlobalSingleton, logger, uuid4 } from '@sentry/utils';
+import {
+  GLOBAL_OBJ,
+  consoleSandbox,
+  dateTimestampInSeconds,
+  getGlobalSingleton,
+  isThenable,
+  logger,
+  uuid4,
+} from '@sentry/utils';
 
 import { DEFAULT_ENVIRONMENT } from './constants';
 import { DEBUG_BUILD } from './debug-build';
@@ -177,12 +185,35 @@ export class Hub implements HubInterface {
   public withScope<T>(callback: (scope: Scope) => T): T {
     // eslint-disable-next-line deprecation/deprecation
     const scope = this.pushScope();
+
+    let maybePromiseResult: T;
     try {
-      return callback(scope);
-    } finally {
+      maybePromiseResult = callback(scope);
+    } catch (e) {
       // eslint-disable-next-line deprecation/deprecation
       this.popScope();
+      throw e;
     }
+
+    if (isThenable(maybePromiseResult)) {
+      // @ts-expect-error - isThenable returns the wrong type
+      return maybePromiseResult.then(
+        res => {
+          // eslint-disable-next-line deprecation/deprecation
+          this.popScope();
+          return res;
+        },
+        e => {
+          // eslint-disable-next-line deprecation/deprecation
+          this.popScope();
+          throw e;
+        },
+      );
+    }
+
+    // eslint-disable-next-line deprecation/deprecation
+    this.popScope();
+    return maybePromiseResult;
   }
 
   /**

packages/core/test/lib/exports.test.ts

@@ -42,12 +42,86 @@ describe('withScope', () => {
     expect(res).toBe('foo');
   });
 
-  it('works with an async function', async () => {
+  it('works with an async function return value', async () => {
     const res = withScope(async scope => {
       return 'foo';
     });
 
     expect(res).toBeInstanceOf(Promise);
     expect(await res).toBe('foo');
   });
+
+  it('correctly sets & resets the current scope', () => {
+    const scope1 = getCurrentScope();
+
+    withScope(scope2 => {
+      expect(scope2).not.toBe(scope1);
+      expect(getCurrentScope()).toBe(scope2);
+    });
+
+    withScope(scope3 => {
+      expect(scope3).not.toBe(scope1);
+      expect(getCurrentScope()).toBe(scope3);
+    });
+
+    expect(getCurrentScope()).toBe(scope1);
+  });
+
+  it('correctly sets & resets the current scope with async functions', async () => {
+    const scope1 = getCurrentScope();
+
+    await withScope(async scope2 => {
+      expect(scope2).not.toBe(scope1);
+      expect(getCurrentScope()).toBe(scope2);
+
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      expect(getCurrentScope()).toBe(scope2);
+    });
+
+    await withScope(async scope3 => {
+      expect(scope3).not.toBe(scope1);
+      expect(getCurrentScope()).toBe(scope3);
+
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      expect(getCurrentScope()).toBe(scope3);
+    });
+
+    expect(getCurrentScope()).toBe(scope1);
+  });
+
+  it('correctly sets & resets the current scope when an error happens', () => {
+    const scope1 = getCurrentScope();
+
+    const error = new Error('foo');
+
+    expect(() =>
+      withScope(scope2 => {
+        expect(scope2).not.toBe(scope1);
+        expect(getCurrentScope()).toBe(scope2);
+
+        throw error;
+      }),
+    ).toThrow(error);
+
+    expect(getCurrentScope()).toBe(scope1);
+  });
+
+  it('correctly sets & resets the current scope with async functions & errors', async () => {
+    const scope1 = getCurrentScope();
+
+    const error = new Error('foo');
+
+    await expect(
+      withScope(async scope2 => {
+        expect(scope2).not.toBe(scope1);
+        expect(getCurrentScope()).toBe(scope2);
+
+        throw error;
+      }),
+    ).rejects.toBe(error);
+
+    expect(getCurrentScope()).toBe(scope1);
+  });
 });