Associate child runs with the span ID of the span in the parent run that triggered the child run

Author: ericallam

.changeset/short-tomatoes-beam.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/core": patch
+---
+
+Add otel propagation headers "below" the API fetch span, to attribute the child runs with the proper parent span ID

apps/webapp/app/v3/eventRepository.server.ts

@@ -835,7 +835,8 @@ export class EventRepository {
     options: TraceEventOptions & { incomplete?: boolean },
     callback: (
       e: EventBuilder,
-      traceContext: Record<string, string | undefined>
+      traceContext: Record<string, string | undefined>,
+      traceparent?: { traceId: string; spanId: string }
     ) => Promise<TResult>
   ): Promise<TResult> {
     const propagatedContext = extractContextFromCarrier(options.context ?? {});
@@ -892,7 +893,7 @@ export class EventRepository {
       },
     };
 
-    const result = await callback(eventBuilder, traceContext);
+    const result = await callback(eventBuilder, traceContext, propagatedContext?.traceparent);
 
     const duration = process.hrtime.bigint() - start;
 

apps/webapp/app/v3/services/triggerTask.server.ts

@@ -240,7 +240,7 @@ export class TriggerTaskService extends BaseService {
           incomplete: true,
           immediate: true,
         },
-        async (event, traceContext) => {
+        async (event, traceContext, traceparent) => {
           const run = await autoIncrementCounter.incrementInTransaction(
             `v3-run:${environment.id}:${taskId}`,
             async (num, tx) => {
@@ -307,6 +307,8 @@ export class TriggerTaskService extends BaseService {
                   traceContext: traceContext,
                   traceId: event.traceId,
                   spanId: event.spanId,
+                  parentSpanId:
+                    options.parentAsLinkType === "replay" ? undefined : traceparent?.spanId,
                   lockedToVersionId: lockedToBackgroundWorker?.id,
                   concurrencyKey: body.options?.concurrencyKey,
                   queue: queueName,

packages/core/src/v3/apiClient/core.ts

@@ -4,7 +4,7 @@ import { RetryOptions } from "../schemas/index.js";
 import { calculateNextRetryDelay } from "../utils/retries.js";
 import { ApiConnectionError, ApiError, ApiSchemaValidationError } from "./errors.js";
 
-import { Attributes, Span } from "@opentelemetry/api";
+import { Attributes, Span, context, propagation } from "@opentelemetry/api";
 import { SemanticInternalAttributes } from "../semanticInternalAttributes.js";
 import { TriggerTracer } from "../tracer.js";
 import { accessoryAttributes } from "../utils/styleAttributes.js";
@@ -184,9 +184,11 @@ async function _doZodFetch<TResponseBodySchema extends z.ZodTypeAny>(
   requestInit?: PromiseOrValue<RequestInit>,
   options?: ZodFetchOptions
 ): Promise<ZodFetchResult<z.output<TResponseBodySchema>>> {
-  const $requestInit = await requestInit;
+  let $requestInit = await requestInit;
 
   return traceZodFetch({ url, requestInit: $requestInit, options }, async (span) => {
+    $requestInit = injectPropagationHeadersIfInWorker($requestInit);
+
     const result = await _doZodFetchWithRetries(schema, url, $requestInit, options);
 
     if (options?.onResponseBody && span) {
@@ -577,3 +579,23 @@ export function isEmptyObj(obj: Object | null | undefined): boolean {
 export function hasOwn(obj: Object, key: string): boolean {
   return Object.prototype.hasOwnProperty.call(obj, key);
 }
+
+// If the requestInit has a header x-trigger-worker = true, then we will do
+// propagation.inject(context.active(), headers);
+// and return the new requestInit.
+function injectPropagationHeadersIfInWorker(requestInit?: RequestInit): RequestInit | undefined {
+  const headers = new Headers(requestInit?.headers);
+
+  if (headers.get("x-trigger-worker") !== "true") {
+    return requestInit;
+  }
+
+  const headersObject = Object.fromEntries(headers.entries());
+
+  propagation.inject(context.active(), headersObject);
+
+  return {
+    ...requestInit,
+    headers: new Headers(headersObject),
+  };
+}

packages/core/src/v3/apiClient/index.ts

@@ -1,4 +1,3 @@
-import { context, propagation } from "@opentelemetry/api";
 import { z } from "zod";
 import {
   AddTagsRequestBody,
@@ -509,7 +508,6 @@ export class ApiClient {
     // Only inject the context if we are inside a task
     if (taskContext.isInsideTask) {
       headers["x-trigger-worker"] = "true";
-      propagation.inject(context.active(), headers);
 
       if (spanParentAsLink) {
         headers["x-trigger-span-parent-as-link"] = "1";

packages/database/prisma/migrations/20240924130558_add_parent_span_id_to_task_run/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "TaskRun" ADD COLUMN     "parentSpanId" TEXT;

packages/database/prisma/schema.prisma

@@ -1748,6 +1748,9 @@ model TaskRun {
   /// The depth of this task run in the task run hierarchy
   depth Int @default(0)
 
+  /// The span ID of the "trigger" span in the parent task run
+  parentSpanId String?
+
   @@unique([runtimeEnvironmentId, taskIdentifier, idempotencyKey])
   // Finding child runs
   @@index([parentTaskRunId])