ref(core): Streamline how active spans are set in core

packages/core/src/tracing/trace.ts

@@ -88,22 +88,13 @@ export function startSpan<T>(options: StartSpanOptions, callback: (span: Span) =
       // Ignored root spans still need to be set on scope so that `getActiveSpan()` returns them
       // and descendants are also non-recording. Ignored child spans don't need this because
       // the parent span is already on scope.
-      if (!spanIsIgnored(activeSpan) || !parentSpan) {
-        _setSpanForScope(scope, activeSpan);
-      }
+      const makeSpanActive = !spanIsIgnored(activeSpan) || !parentSpan;
 
-      return handleCallbackErrors(
+      return runCallback(
+        activeSpan,
+        makeSpanActive,
         () => callback(activeSpan),
-        () => {
-          // Only update the span status if it hasn't been changed yet, and the span is not yet finished
-          const { status } = spanToJSON(activeSpan);
-          if (activeSpan.isRecording() && status === 'ok') {
-            activeSpan.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
-          }
-        },
-        () => {
-          activeSpan.end();
-        },
+        () => activeSpan.end(),
       );
     });
   });
@@ -150,24 +141,13 @@ export function startSpanManual<T>(options: StartSpanOptions, callback: (span: S
 
       // We don't set ignored child spans onto the scope because there likely is an active,
       // unignored span on the scope already.
-      if (!spanIsIgnored(activeSpan) || !parentSpan) {
-        _setSpanForScope(scope, activeSpan);
-      }
+      const makeSpanActive = !spanIsIgnored(activeSpan) || !parentSpan;
 
-      return handleCallbackErrors(
-        // We pass the `finish` function to the callback, so the user can finish the span manually
-        // this is mainly here for historic purposes because previously, we instructed users to call
-        // `finish` instead of `span.end()` to also clean up the scope. Nowadays, calling `span.end()`
-        // or `finish` has the same effect and we simply leave it here to avoid breaking user code.
-        () => callback(activeSpan, () => activeSpan.end()),
-        () => {
-          // Only update the span status if it hasn't been changed yet, and the span is not yet finished
-          const { status } = spanToJSON(activeSpan);
-          if (activeSpan.isRecording() && status === 'ok') {
-            activeSpan.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
-          }
-        },
-      );
+      // We pass the `finish` function to the callback, so the user can finish the span manually
+      // this is mainly here for historic purposes because previously, we instructed users to call
+      // `finish` instead of `span.end()` to also clean up the scope. Nowadays, calling `span.end()`
+      // or `finish` has the same effect and we simply leave it here to avoid breaking user code.
+      return runCallback(activeSpan, makeSpanActive, () => callback(activeSpan, () => activeSpan.end()));
     });
   });
 }
@@ -677,3 +657,29 @@ function _shouldIgnoreStreamedSpan(client: Client | undefined, spanArguments: Se
 export function spanIsIgnored(span: Span): span is SentryNonRecordingSpan {
   return spanIsNonRecordingSpan(span) && span.dropReason === 'ignored';
 }
+
+function runCallback<T>(span: Span, makeSpanActive: boolean, callback: () => T, finallyCallback?: () => void): T {
+  const wrapper = makeSpanActive
+    ? (callback: () => T) => {
+        return withActiveSpan(span, () => {
+          // Make sure the correct scope is captured on the span, since withActiveSpan forks the scope
+          setCapturedScopesOnSpan(span, getCurrentScope(), getIsolationScope());
+          return callback();
+        });
+      }
+    : (callback: () => T) => callback();
+
+  return wrapper(() =>
+    handleCallbackErrors(
+      () => callback(),
+      () => {
+        // Only update the span status if it hasn't been changed yet, and the span is not yet finished
+        const { status } = spanToJSON(span);
+        if (span.isRecording() && status === 'ok') {
+          span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
+        }
+      },
+      finallyCallback,
+    ),
+  );
+}