Mr0grog · March 14, 2023 16:24 · Mar 14, 2023 · Mar 13, 2023
diff --git a/sentry-tracing.ts b/sentry-tracing.ts
@@ -1,4 +1,6 @@
+import EventEmitter from "node:events";
 import * as Sentry from "@sentry/node";
+import { Transaction } from "@sentry/tracing";
 import type { Span, SpanStatusType } from "@sentry/tracing";
 import type { SpanContext } from "@sentry/types";
 export type { Span } from "@sentry/tracing";
@@ -8,6 +10,39 @@ interface SpanOptions extends SpanContext {
   timeout?: number;
 }
 
+// Monkey-patch Transaction to add an event to notify listeners (in our case,
+// child spans) when the transaction is finishing. This is super hacky.
+//
+// The most recent release of the Sentry SDK has a `finishTransaction` event
+// on the hub's client, but there are a lot of guard clauses anywhere the client
+// gets used internally, and I'm not sure how reliable it is for this use case.
+// We want to wrap up before the transaction finishes, not when a client that
+// may-or-may-not exist depending on configuring is finishing a transaction.
+// (Also, the transaction's `_hub` is a nullable private property, so it would
+// still be hacky to grab it and add a listener anyway.)
+interface PatchedTransaction extends Transaction {
+  _onFinish(listener: (...args: any[]) => void): void;
+  _emitter?: EventEmitter;
+}
+
+const _transactionFinish = Transaction.prototype.finish;
+Transaction.prototype.finish = function finish(
+  this: PatchedTransaction,
+  endTimestamp?: number
+) {
+  this._emitter?.emit("finish", this);
+  return _transactionFinish.call(this, endTimestamp);
+};
+
+// @ts-expect-error: _onFinish doesn't exist; we're adding it.
+Transaction.prototype._onFinish = function onFinish(
+  this: PatchedTransaction,
+  listener: (...args: any[]) => void
+) {
+  if (!this._emitter) this._emitter = new EventEmitter();
+  return this._emitter.once("finish", listener);
+};
+
 /**
  * Start a tracing span. The returned span should be explicitly ended with
  * `finishSpan`. The created span will be a child of whatever span is currently
@@ -53,16 +88,10 @@ export function startSpan(options: SpanOptions): Span {
     setTimeout(() => {
       cancelSpan(newSpan, "deadline_exceeded");
     }, timeout).unref();
-  } else if (newSpan.transaction && newSpan.transaction !== newSpan) {
-    // FIXME: newer Sentry has an event for this: "finishTransaction" emitted
-    // on the hub's client:
-    // - Code: https://github.com/getsentry/sentry-javascript/blob/ba99e7cdf725725e5a1b99e9d814353dbb3ae2b6/packages/core/src/tracing/transaction.ts#L144-L147
-    // - Feature: https://github.com/getsentry/sentry-javascript/issues/7262
-    const _finishTransaction = newSpan.transaction.finish;
-    newSpan.transaction.finish = function finish(...args) {
-      cancelSpan(newSpan, "cancelled", "did_not_finish");
-      return _finishTransaction.call(this, ...args);
-    };
+  } else if (newSpan.transaction !== newSpan) {
+    (newSpan.transaction as PatchedTransaction)?._onFinish(() =>
+      cancelSpan(newSpan, "cancelled", "did_not_finish")
+    );
   }
 
   return newSpan;
@@ -73,6 +102,8 @@ export function startSpan(options: SpanOptions): Span {
  * this span's parent (if it has a parent).
  */
 export function finishSpan(span: Span, timestamp?: number): void {
+  if (span.endTimestamp) return;
+
   span.finish(timestamp);
 
   let parent;
@@ -96,13 +127,17 @@ export function finishSpan(span: Span, timestamp?: number): void {
 function cancelSpan(
   span: Span,
   status: SpanStatusType = "cancelled",
-  tag: string = status
+  tag: any = status
 ): void {
-  if (!span.endTimestamp) {
-    span.setStatus(status);
-    span.setTag("cancel", tag);
-    finishSpan(span);
+  if (span.endTimestamp) return;
+
+  if (typeof tag !== "string") {
+    tag = `error:${tag?.code || tag?.constructor?.name || "?"}`;
   }
+
+  span.setStatus(status);
+  span.setTag("cancel", tag);
+  finishSpan(span);
 }
 
 /**
@@ -137,12 +172,21 @@ export function withSpan<T extends (span?: Span) => any>(
   try {
     callbackResult = callback(span);
   } catch (error) {
-    finishSpan(span);
+    cancelSpan(span, "unknown_error", error);
     throw error;
   }
 
-  if ("then" in callbackResult && "finally" in callbackResult) {
-    return callbackResult.finally(() => finishSpan(span));
+  if ("then" in callbackResult) {
+    return callbackResult.then(
+      (result: any) => {
+        finishSpan(span);
+        return result;
+      },
+      (error: any) => {
+        cancelSpan(span, "internal_error", error);
+        throw error;
+      }
+    );
   } else {
     finishSpan(span);
     return callbackResult;

diff --git a/sentry-tracing.ts b/sentry-tracing.ts
@@ -0,0 +1,150 @@
+import * as Sentry from "@sentry/node";
+import type { Span, SpanStatusType } from "@sentry/tracing";
+import type { SpanContext } from "@sentry/types";
+export type { Span } from "@sentry/tracing";
+
+interface SpanOptions extends SpanContext {
+  parentSpan?: Span;
+  timeout?: number;
+}
+
+/**
+ * Start a tracing span. The returned span should be explicitly ended with
+ * `finishSpan`. The created span will be a child of whatever span is currently
+ * active (and then become the current span itself), or if there is no current
+ * span or transaction, this will start one for you.
+ *
+ * Alternatively, can explicitly pass an actual span object to be the parent:
+ * `startSpan({ parentSpan: yourSpan })`. In this case, the new span won't
+ * automatically become the global "current" span.
+ *
+ * Spans created this way will be automatically canceled when their transaction
+ * finishes (Sentry will drop any unfinished spans). Alternatively, you can
+ * set `timeout` to a number of milliseconds, and the span will be canceled
+ * after that time (this should prevent you from accidentally leaving a span
+ * open forever). Canceled spans have a non-ok status set and a `cancel` tag
+ * with a reason.
+ *
+ * More on why spans need canceling:
+ * https://github.com/getsentry/sentry-javascript/issues/4165#issuecomment-971424754
+ */
+export function startSpan(options: SpanOptions): Span {
+  let { parentSpan, timeout, ...spanOptions } = options;
+
+  let scope;
+  if (!parentSpan) {
+    scope = Sentry.getCurrentHub().getScope();
+    parentSpan = scope.getSpan();
+  }
+
+  let newSpan: Span;
+  if (parentSpan) {
+    newSpan = parentSpan.startChild(spanOptions);
+  } else {
+    newSpan = Sentry.startTransaction(spanOptions as any);
+  }
+
+  // If we retrieved the span from the scope, update the scope.
+  if (scope) {
+    scope.setSpan(newSpan);
+  }
+
+  if (timeout) {
+    setTimeout(() => {
+      cancelSpan(newSpan, "deadline_exceeded");
+    }, timeout).unref();
+  } else if (newSpan.transaction && newSpan.transaction !== newSpan) {
+    // FIXME: newer Sentry has an event for this: "finishTransaction" emitted
+    // on the hub's client:
+    // - Code: https://github.com/getsentry/sentry-javascript/blob/ba99e7cdf725725e5a1b99e9d814353dbb3ae2b6/packages/core/src/tracing/transaction.ts#L144-L147
+    // - Feature: https://github.com/getsentry/sentry-javascript/issues/7262
+    const _finishTransaction = newSpan.transaction.finish;
+    newSpan.transaction.finish = function finish(...args) {
+      cancelSpan(newSpan, "cancelled", "did_not_finish");
+      return _finishTransaction.call(this, ...args);
+    };
+  }
+
+  return newSpan;
+}
+
+/**
+ * Finish a tracing span. This will also replace the global "current" span with
+ * this span's parent (if it has a parent).
+ */
+export function finishSpan(span: Span, timestamp?: number): void {
+  span.finish(timestamp);
+
+  let parent;
+  if (span.parentSpanId) {
+    // FIXME: abstract this with a nice name, even though it's simple.
+    parent = span.spanRecorder?.spans?.find(
+      (s) => s.spanId === span.parentSpanId
+    );
+  }
+
+  const scope = Sentry.getCurrentHub().getScope();
+  if (scope.getSpan() === span) {
+    scope.setSpan(parent);
+  }
+}
+
+/**
+ * If a span is not finished, finish it and set its status and tags to indicate
+ * that it was canceled. If the span is already finished, do nothing.
+ */
+function cancelSpan(
+  span: Span,
+  status: SpanStatusType = "cancelled",
+  tag: string = status
+): void {
+  if (!span.endTimestamp) {
+    span.setStatus(status);
+    span.setTag("cancel", tag);
+    finishSpan(span);
+  }
+}
+
+/**
+ * Create a new span, run the provided function inside of it, and finish the
+ * span afterward. The function can be async, in which case this will return an
+ * awaitable promise.
+ *
+ * The provided function can take the span as the first argument, in case it
+ * needs to modify the span in some way. If the function returns a value,
+ * `withSpan` will return that value as well.
+ *
+ * @example
+ * withSpan({ op: "validateData" }, (span) => {
+ *   doSomeDataValidation();
+ * });
+ *
+ * let data = { some: "data" };
+ * const id = await withSpan({ op: "saveData" }, async (span) => {
+ *   const id = await saveData(data);
+ *   await updateSomeRelatedRecord(id, otherData);
+ *   return id;
+ * });
+ */
+export function withSpan<T extends (span?: Span) => any>(
+  options: SpanOptions | string,
+  callback: T
+): ReturnType<T> {
+  if (typeof options === "string") options = { op: options };
+  const span = startSpan(options);
+
+  let callbackResult;
+  try {
+    callbackResult = callback(span);
+  } catch (error) {
+    finishSpan(span);
+    throw error;
+  }
+
+  if ("then" in callbackResult && "finally" in callbackResult) {
+    return callbackResult.finally(() => finishSpan(span));
+  } else {
+    finishSpan(span);
+    return callbackResult;
+  }
+}
No results found