TS/JS Sentry Tracing Helpers

sentry-tracing.ts

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";

interface SpanOptions extends SpanContext {
  parentSpan?: Span;
  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
 * 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) {
    (newSpan.transaction as PatchedTransaction)?._onFinish(() =>
      cancelSpan(newSpan, "cancelled", "did_not_finish")
    );
  }

  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 {
  if (span.endTimestamp) return;

  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: any = status
): void {
  if (span.endTimestamp) return;

  if (typeof tag !== "string") {
    tag = `error:${tag?.code || tag?.constructor?.name || "?"}`;
  }

  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) {
    cancelSpan(span, "unknown_error", error);
    throw error;
  }

  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;
  }
}

Update: this note about using events may not actually be feasible:

// 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);
};

The event is on the client, which belongs to the transaction’s hub. You could have lots of hubs, and at this point in the code we don’t reliably know whether the transaction’s hub is the current one or another one. The transaction holds a reference to its hub, but it’s a private property named _hub, and so may not be stable or reliable. So… this monkey-patching of the finish() method may actually still be the best approach (although it could be done in a slightly more clean way).