Add per-isolate SSE + OTel memory counters to cloud worker (#1255)

* Add per-isolate SSE + OTel memory counters to cloud worker

* Address review: abort-free byte counts, SSE-only wrap guard, no drop/export double-count

Author: RhysSullivan

apps/cloud/src/mcp/agent-handler.ts

@@ -14,6 +14,7 @@ import {
 } from "@executor-js/cloudflare/mcp/do-headers";
 import type { McpSessionProps } from "@executor-js/cloudflare/mcp/agent-durable-object";
 
+import { wrapMcpSseResponse } from "../observability/memory-metrics";
 import { cloudMcpAuth } from "./auth-provider";
 import { McpSessionDOSqlite } from "./session-durable-object";
 
@@ -158,6 +159,7 @@ export const makeCloudMcpAgentHandler = () => {
       },
       resource,
     );
-    return serve.fetch(forwarded, env, ctx);
+    const response = await serve.fetch(forwarded, env, ctx);
+    return wrapMcpSseResponse(request, env, response);
   };
 };

apps/cloud/src/observability/memory-metrics.ts

@@ -0,0 +1,329 @@
+import type { Context } from "@opentelemetry/api";
+import type {
+  ReadableSpan,
+  Span,
+  SpanExporter,
+  SpanProcessor,
+} from "@opentelemetry/sdk-trace-base";
+
+const LOG_PREFIX = "[executor:mem-metrics]";
+const PERIODIC_EMIT_INTERVAL_MS = 60_000;
+const STALLED_WRITE_MS = 30_000;
+const EXPORT_SUCCESS_CODE = 0;
+export const OTEL_MAX_SPAN_QUEUE_SIZE = 2_048;
+
+type CloseReason = "normal" | "cancel" | "error";
+
+type CfRequestMetadata = {
+  readonly colo?: string;
+};
+
+type VersionMetadataEnv = {
+  readonly VERSION_METADATA?: WorkerVersionMetadata;
+  readonly CF_VERSION_METADATA?: WorkerVersionMetadata;
+  readonly SCRIPT_VERSION?: string;
+};
+
+type SseConnection = {
+  readonly id: number;
+  readonly startedAt: number;
+  lastWriteAt: number;
+  chunks: number;
+  bytes: number;
+  closed: boolean;
+  colo: string;
+  scriptVersion: string;
+};
+
+type AgeBuckets = {
+  readonly lt1m: number;
+  readonly m1to5: number;
+  readonly m5to30: number;
+  readonly m30to60: number;
+  readonly gt60m: number;
+};
+
+type OtelMetricsSnapshot = {
+  readonly spansEnded: number;
+  readonly spansExported: number;
+  readonly spansDropped: number;
+  readonly spanExportFailures: number;
+  readonly forceFlushCalls: number;
+  readonly forceFlushFailures: number;
+  readonly forceFlushDurationMsTotal: number;
+  readonly forceFlushDurationMsLast: number;
+  readonly maxQueueSize: number;
+};
+
+const activeSse = new Map<number, SseConnection>();
+
+let nextConnectionId = 1;
+let totalBytesForwarded = 0;
+let lastSnapshotEmittedAt = 0;
+
+const otelMetrics = {
+  spansEnded: 0,
+  spansExported: 0,
+  spansDropped: 0,
+  spanExportFailures: 0,
+  forceFlushCalls: 0,
+  forceFlushFailures: 0,
+  forceFlushDurationMsTotal: 0,
+  forceFlushDurationMsLast: 0,
+};
+
+const requestWithCf = (request: Request): Request & { readonly cf?: CfRequestMetadata } =>
+  request as Request & { readonly cf?: CfRequestMetadata };
+
+const scriptVersionFromEnv = (env: Env): string => {
+  const metadataEnv = env as Env & VersionMetadataEnv;
+  return (
+    metadataEnv.VERSION_METADATA?.id ??
+    metadataEnv.VERSION_METADATA?.tag ??
+    metadataEnv.CF_VERSION_METADATA?.id ??
+    metadataEnv.CF_VERSION_METADATA?.tag ??
+    metadataEnv.SCRIPT_VERSION ??
+    ""
+  );
+};
+
+const utf8ByteCounter = new TextEncoder();
+
+const byteLengthOf = (chunk: unknown): number => {
+  if (typeof chunk === "string") return utf8ByteCounter.encode(chunk).length;
+  if (chunk instanceof ArrayBuffer) return chunk.byteLength;
+  if (ArrayBuffer.isView(chunk)) return chunk.byteLength;
+  return 0;
+};
+
+const emptyAgeBuckets = (): AgeBuckets => ({
+  lt1m: 0,
+  m1to5: 0,
+  m5to30: 0,
+  m30to60: 0,
+  gt60m: 0,
+});
+
+const bucketAge = (buckets: AgeBuckets, ageMs: number): AgeBuckets => {
+  const next = { ...buckets };
+  if (ageMs < 60_000) next.lt1m += 1;
+  else if (ageMs < 5 * 60_000) next.m1to5 += 1;
+  else if (ageMs < 30 * 60_000) next.m5to30 += 1;
+  else if (ageMs < 60 * 60_000) next.m30to60 += 1;
+  else next.gt60m += 1;
+  return next;
+};
+
+const otelSnapshot = (maxQueueSize: number): OtelMetricsSnapshot => ({
+  ...otelMetrics,
+  maxQueueSize,
+});
+
+const snapshot = (now: number, maxQueueSize: number) => {
+  let oldestConnectionAgeMs = 0;
+  let stalledConnections = 0;
+  let ageBuckets = emptyAgeBuckets();
+  const colos = new Set<string>();
+  const scriptVersions = new Set<string>();
+
+  for (const connection of activeSse.values()) {
+    const ageMs = now - connection.startedAt;
+    oldestConnectionAgeMs = Math.max(oldestConnectionAgeMs, ageMs);
+    if (now - connection.lastWriteAt > STALLED_WRITE_MS) stalledConnections += 1;
+    ageBuckets = bucketAge(ageBuckets, ageMs);
+    if (connection.colo) colos.add(connection.colo);
+    if (connection.scriptVersion) scriptVersions.add(connection.scriptVersion);
+  }
+  const sortedColos = Array.from(colos).sort();
+  const sortedScriptVersions = Array.from(scriptVersions).sort();
+
+  return {
+    activeSseConnections: activeSse.size,
+    ageBuckets,
+    oldestConnectionAgeMs,
+    totalBytesForwarded,
+    stalledConnections,
+    stalledWriteThresholdMs: STALLED_WRITE_MS,
+    colo: sortedColos[0] ?? "",
+    colos: sortedColos,
+    scriptVersion: sortedScriptVersions[0] ?? "",
+    scriptVersions: sortedScriptVersions,
+    otel: otelSnapshot(maxQueueSize),
+  };
+};
+
+const emitSnapshot = (event: string, now: number, maxQueueSize: number): void => {
+  lastSnapshotEmittedAt = now;
+  console.log(
+    `${LOG_PREFIX} ${JSON.stringify({
+      event,
+      ...snapshot(now, maxQueueSize),
+    })}`,
+  );
+};
+
+const maybeEmitPeriodicSnapshot = (now: number, maxQueueSize: number): void => {
+  if (now - lastSnapshotEmittedAt < PERIODIC_EMIT_INTERVAL_MS) return;
+  emitSnapshot("snapshot", now, maxQueueSize);
+};
+
+const closeConnection = (
+  connection: SseConnection,
+  reason: CloseReason,
+  maxQueueSize: number,
+): void => {
+  if (connection.closed) return;
+  connection.closed = true;
+  activeSse.delete(connection.id);
+  const now = Date.now();
+  console.log(
+    `${LOG_PREFIX} ${JSON.stringify({
+      event: "sse_close",
+      connectionId: connection.id,
+      reason,
+      ageMs: now - connection.startedAt,
+      bytesForwarded: connection.bytes,
+      chunksForwarded: connection.chunks,
+      lastWriteAgeMs: now - connection.lastWriteAt,
+      colo: connection.colo,
+      scriptVersion: connection.scriptVersion,
+    })}`,
+  );
+  maybeEmitPeriodicSnapshot(now, maxQueueSize);
+};
+
+const sseHeaders = (headers: Headers): Headers => {
+  const next = new Headers(headers);
+  next.delete("content-length");
+  return next;
+};
+
+export const wrapMcpSseResponse = (request: Request, env: Env, response: Response): Response => {
+  if (
+    request.method !== "GET" ||
+    response.body === null ||
+    !(response.headers.get("content-type") ?? "").includes("text/event-stream")
+  )
+    return response;
+
+  const now = Date.now();
+  const connection: SseConnection = {
+    id: nextConnectionId++,
+    startedAt: now,
+    lastWriteAt: now,
+    chunks: 0,
+    bytes: 0,
+    closed: false,
+    colo: requestWithCf(request).cf?.colo ?? "",
+    scriptVersion: scriptVersionFromEnv(env),
+  };
+  activeSse.set(connection.id, connection);
+  emitSnapshot("sse_open", now, OTEL_MAX_SPAN_QUEUE_SIZE);
+
+  const counting = new TransformStream<unknown, unknown>({
+    transform(chunk, controller) {
+      const written = byteLengthOf(chunk);
+      const writeAt = Date.now();
+      connection.chunks += 1;
+      connection.bytes += written;
+      connection.lastWriteAt = writeAt;
+      totalBytesForwarded += written;
+      maybeEmitPeriodicSnapshot(writeAt, OTEL_MAX_SPAN_QUEUE_SIZE);
+      controller.enqueue(chunk);
+    },
+    flush() {
+      closeConnection(connection, "normal", OTEL_MAX_SPAN_QUEUE_SIZE);
+    },
+    cancel(reason) {
+      closeConnection(
+        connection,
+        reason === undefined ? "cancel" : "error",
+        OTEL_MAX_SPAN_QUEUE_SIZE,
+      );
+    },
+  });
+
+  const body = response.body.pipeThrough(counting);
+  return new Response(body, {
+    status: response.status,
+    statusText: response.statusText,
+    headers: sseHeaders(response.headers),
+  });
+};
+
+export const recordForceFlush = (durationMs: number, failed: boolean): void => {
+  otelMetrics.forceFlushCalls += 1;
+  otelMetrics.forceFlushDurationMsLast = durationMs;
+  otelMetrics.forceFlushDurationMsTotal += durationMs;
+  if (failed) otelMetrics.forceFlushFailures += 1;
+};
+
+export class CountingSpanExporter implements SpanExporter {
+  constructor(
+    private readonly inner: SpanExporter,
+    private readonly onExportAttempt: (spans: number) => void,
+  ) {}
+
+  export(spans: ReadableSpan[], resultCallback: Parameters<SpanExporter["export"]>[1]): void {
+    this.onExportAttempt(spans.length);
+    this.inner.export(spans, (result) => {
+      if (result.code === EXPORT_SUCCESS_CODE) {
+        otelMetrics.spansExported += spans.length;
+      } else {
+        otelMetrics.spansDropped += spans.length;
+        otelMetrics.spanExportFailures += 1;
+      }
+      resultCallback(result);
+    });
+  }
+
+  shutdown(): Promise<void> {
+    return this.inner.shutdown();
+  }
+
+  forceFlush(): Promise<void> {
+    return this.inner.forceFlush?.() ?? Promise.resolve();
+  }
+}
+
+export class CountingSpanProcessor implements SpanProcessor {
+  private estimatedQueuedSpans = 0;
+
+  constructor(
+    private readonly inner: SpanProcessor,
+    private readonly maxQueueSize: number,
+  ) {}
+
+  forceFlush(): Promise<void> {
+    return this.inner.forceFlush();
+  }
+
+  onStart(span: Span, parentContext: Context): void {
+    this.inner.onStart(span, parentContext);
+  }
+
+  onEnding(span: Span): void {
+    this.inner.onEnding?.(span);
+  }
+
+  onEnd(span: ReadableSpan): void {
+    otelMetrics.spansEnded += 1;
+    if (this.estimatedQueuedSpans >= this.maxQueueSize) {
+      // At the bound the inner processor would drop this span anyway; drop it
+      // here instead of forwarding, so a dropped span is never also counted as
+      // exported (which would double-count and inflate the drop signal).
+      otelMetrics.spansDropped += 1;
+      return;
+    }
+    this.estimatedQueuedSpans += 1;
+    this.inner.onEnd(span);
+  }
+
+  shutdown(): Promise<void> {
+    return this.inner.shutdown();
+  }
+
+  recordExportAttempt(spans: number): void {
+    this.estimatedQueuedSpans = Math.max(0, this.estimatedQueuedSpans - spans);
+  }
+}