Fix race condition between response end and the trace end

.changeset/early-coats-pull.md

@@ -0,0 +1,5 @@
+---
+"@vercel/otel": patch
+---
+
+Eliminate race condition between response end and trace end

.github/workflows/release-snapshot.yaml

@@ -45,7 +45,7 @@ on:
 
 jobs:
   release-snapshot:
-    name: Release
+    name: Release Snapshot
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
@@ -68,7 +68,7 @@ jobs:
         run: pnpm install
 
       - name: Build
-        run: pnpm build
+        run: pnpm --filter @vercel/otel build
 
       - name: Add SHORT_SHA env property with commit short sha
         run: echo "SHORT_SHA=`echo ${{ github.sha }} | cut -c1-8`" >> $GITHUB_ENV

packages/otel/src/exporters/otlp-exporter-base.ts

@@ -67,6 +67,8 @@ export abstract class OTLPExporterEdgeBase<
       .then((res) => {
         diag.debug("@vercel/otel/otlp: onSuccess", res.status, res.statusText);
         onSuccess();
+        // Drain the response body.
+        void res.arrayBuffer().catch(() => undefined);
       })
       .catch((err) => {
         diag.error("@vercel/otel/otlp: onError", err);

packages/otel/src/processor/composite-span-processor.ts

@@ -11,6 +11,7 @@ import { getVercelRequestContextAttributes } from "../vercel-request-context/att
 /** @internal */
 export class CompositeSpanProcessor implements SpanProcessor {
   private readonly rootSpanIds = new Map<string, string>();
+  private readonly waitSpanEnd = new Map<string, () => void>();
 
   constructor(private processors: SpanProcessor[]) {}
 
@@ -45,7 +46,29 @@ export class CompositeSpanProcessor implements SpanProcessor {
       // Flush the streams to avoid data loss.
       const vrc = getVercelRequestContext();
       if (vrc) {
-        vrc.waitUntil(() => this.forceFlush());
+        vrc.waitUntil(async () => {
+          if (this.rootSpanIds.has(traceId)) {
+            // Not root has not completed yet, so no point in flushing.
+            // Need to wait for onEnd.
+            const promise = new Promise<void>((resolve) => {
+              this.waitSpanEnd.set(traceId, resolve);
+            });
+            let timer: NodeJS.Timeout | undefined;
+            await Promise.race([
+              promise,
+              new Promise((resolve) => {
+                timer = setTimeout(() => {
+                  this.waitSpanEnd.delete(traceId);
+                  resolve(undefined);
+                }, 50);
+              }),
+            ]);
+            if (timer) {
+              clearTimeout(timer);
+            }
+          }
+          return this.forceFlush();
+        });
       }
     }
 
@@ -57,13 +80,19 @@ export class CompositeSpanProcessor implements SpanProcessor {
   onEnd(span: ReadableSpan): void {
     const { traceId, spanId } = span.spanContext();
     const isRoot = this.rootSpanIds.get(traceId) === spanId;
-    if (isRoot) {
-      this.rootSpanIds.delete(traceId);
-    }
 
     for (const spanProcessor of this.processors) {
       spanProcessor.onEnd(span);
     }
+
+    if (isRoot) {
+      this.rootSpanIds.delete(traceId);
+      const pending = this.waitSpanEnd.get(traceId);
+      if (pending) {
+        this.waitSpanEnd.delete(traceId);
+        pending();
+      }
+    }
   }
 }