Fixing the dev & shared queue consumer telemetry so we can see what’s going on

Author: ericallam

apps/webapp/app/env.server.ts

@@ -137,7 +137,7 @@ const EnvironmentSchema = z.object({
   INTERNAL_OTEL_TRACE_EXPORTER_AUTH_HEADER_VALUE: z.string().optional(),
   INTERNAL_OTEL_TRACE_LOGGING_ENABLED: z.string().default("1"),
   // this means 1/20 traces or 5% of traces will be sampled (sampled = recorded)
-  INTERNAL_OTEL_TRACE_SAMPING_RATE: z.string().default("20"),
+  INTERNAL_OTEL_TRACE_SAMPLING_RATE: z.string().default("20"),
   INTERNAL_OTEL_TRACE_INSTRUMENT_PRISMA_ENABLED: z.string().default("0"),
 });
 

apps/webapp/app/v3/marqs/devQueueConsumer.server.ts

@@ -16,12 +16,14 @@ import { generateFriendlyId } from "../friendlyIdentifiers";
 import { marqs } from "~/v3/marqs/index.server";
 import { CancelAttemptService } from "../services/cancelAttempt.server";
 import { CompleteAttemptService } from "../services/completeAttempt.server";
-import { attributesFromAuthenticatedEnv } from "../tracer.server";
+import {
+  SEMINTATTRS_FORCE_RECORDING,
+  attributesFromAuthenticatedEnv,
+  tracer,
+} from "../tracer.server";
 import { DevSubscriber, devPubSub } from "./devPubSub.server";
 import { CancelTaskRunService } from "../services/cancelTaskRun.server";
 
-const tracer = trace.getTracer("devQueueConsumer");
-
 const MessageBody = z.discriminatedUnion("type", [
   z.object({
     type: z.literal("EXECUTE"),
@@ -165,6 +167,11 @@ export class DevQueueConsumer {
 
       logger.debug("Unsubscribed from background worker channel", { id });
     }
+
+    // We need to end the current span
+    if (this._currentSpan) {
+      this._currentSpan.end();
+    }
   }
 
   async #cancelInProgressRunsAndAttempts(reason: string) {
@@ -284,6 +291,10 @@ export class DevQueueConsumer {
         this._currentSpan.setAttribute("tasks.period.failures", this._taskFailures);
         this._currentSpan.setAttribute("tasks.period.successes", this._taskSuccesses);
 
+        logger.debug("Ending DevQueueConsumer.doWork() trace", {
+          isRecording: this._currentSpan.isRecording(),
+        });
+
         this._currentSpan.end();
       }
 
@@ -294,6 +305,7 @@ export class DevQueueConsumer {
           kind: SpanKind.CONSUMER,
           attributes: {
             ...attributesFromAuthenticatedEnv(this.env),
+            [SEMINTATTRS_FORCE_RECORDING]: true,
           },
         },
         ROOT_CONTEXT

apps/webapp/app/v3/marqs/index.server.ts

@@ -139,7 +139,7 @@ export class MarQS {
   public async dequeueMessageInEnv(env: AuthenticatedEnvironment) {
     return this.#trace(
       "dequeueMessageInEnv",
-      async (span, abort) => {
+      async (span) => {
         const parentQueue = this.keys.envSharedQueueKey(env);
 
         // Read the parent queue for matching queues
@@ -150,7 +150,6 @@ export class MarQS {
         );
 
         if (!messageQueue) {
-          abort();
           return;
         }
 
@@ -167,7 +166,6 @@ export class MarQS {
         });
 
         if (!messageData) {
-          abort();
           return;
         }
 
@@ -181,8 +179,6 @@ export class MarQS {
             [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey,
             [SemanticAttributes.PARENT_QUEUE]: message.parentQueue,
           });
-        } else {
-          abort();
         }
 
         return message;
@@ -204,7 +200,7 @@ export class MarQS {
   public async dequeueMessageInSharedQueue() {
     return this.#trace(
       "dequeueMessageInSharedQueue",
-      async (span, abort) => {
+      async (span) => {
         const parentQueue = constants.SHARED_QUEUE;
 
         // Read the parent queue for matching queues
@@ -215,7 +211,6 @@ export class MarQS {
         );
 
         if (!messageQueue) {
-          abort();
           return;
         }
 
@@ -233,7 +228,6 @@ export class MarQS {
         });
 
         if (!messageData) {
-          abort();
           return;
         }
 
@@ -247,8 +241,6 @@ export class MarQS {
             [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey,
             [SemanticAttributes.PARENT_QUEUE]: message.parentQueue,
           });
-        } else {
-          abort();
         }
 
         return message;
@@ -355,17 +347,12 @@ export class MarQS {
 
   async #trace<T>(
     name: string,
-    fn: (span: Span, abort: () => void) => Promise<T>,
-    options?: SpanOptions
+    fn: (span: Span) => Promise<T>,
+    options?: SpanOptions & { sampleRate?: number }
   ): Promise<T> {
     return tracer.startActiveSpan(name, options ?? {}, async (span) => {
-      let _abort = false;
-      let aborter = () => {
-        _abort = true;
-      };
-
       try {
-        return await fn(span, aborter);
+        return await fn(span);
       } catch (e) {
         if (e instanceof Error) {
           span.recordException(e);
@@ -375,9 +362,7 @@ export class MarQS {
 
         throw e;
       } finally {
-        if (!_abort) {
-          span.end();
-        }
+        span.end();
       }
     });
   }
@@ -480,7 +465,7 @@ export class MarQS {
   ) {
     return this.#trace(
       "getRandomQueueFromParentQueue",
-      async (span, abort) => {
+      async (span) => {
         const { range, selectionId } = await queuePriorityStrategy.nextCandidateSelection(
           parentQueue
         );
@@ -497,7 +482,6 @@ export class MarQS {
         );
 
         if (typeof choice !== "string") {
-          abort();
           return;
         }
 

apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts

@@ -1,12 +1,4 @@
-import {
-  Context,
-  ROOT_CONTEXT,
-  Span,
-  SpanKind,
-  context,
-  propagation,
-  trace,
-} from "@opentelemetry/api";
+import { Context, ROOT_CONTEXT, Span, SpanKind, context, trace } from "@opentelemetry/api";
 import {
   Machine,
   ProdTaskRunExecution,
@@ -28,16 +20,15 @@ import {
 import { z } from "zod";
 import { prisma } from "~/db.server";
 import { logger } from "~/services/logger.server";
-import { generateFriendlyId } from "../friendlyIdentifiers";
+import { singleton } from "~/utils/singleton";
 import { marqs } from "~/v3/marqs/index.server";
 import { EnvironmentVariablesRepository } from "../environmentVariables/environmentVariablesRepository.server";
-import { CancelAttemptService } from "../services/cancelAttempt.server";
+import { generateFriendlyId } from "../friendlyIdentifiers";
 import { socketIo } from "../handleSocketIo.server";
-import { singleton } from "~/utils/singleton";
-import { RestoreCheckpointService } from "../services/restoreCheckpoint.server";
 import { findCurrentWorkerDeployment } from "../models/workerDeployment.server";
-
-const tracer = trace.getTracer("sharedQueueConsumer");
+import { CancelAttemptService } from "../services/cancelAttempt.server";
+import { RestoreCheckpointService } from "../services/restoreCheckpoint.server";
+import { tracer } from "../tracer.server";
 
 const WithTraceContext = z.object({
   traceparent: z.string().optional(),
@@ -154,6 +145,10 @@ export class SharedQueueConsumer {
 
     logger.debug("Stopping shared queue consumer");
     this._enabled = false;
+
+    if (this._currentSpan) {
+      this._currentSpan.end();
+    }
   }
 
   async #cancelInProgressAttempts(reason: string) {

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

@@ -30,6 +30,8 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { singleton } from "~/utils/singleton";
 import { LoggerSpanExporter } from "./telemetry/loggerExporter.server";
 
+export const SEMINTATTRS_FORCE_RECORDING = "forceRecording";
+
 class CustomWebappSampler implements Sampler {
   constructor(private readonly _baseSampler: Sampler) {}
 
@@ -49,8 +51,22 @@ class CustomWebappSampler implements Sampler {
       return { decision: SamplingDecision.NOT_RECORD };
     }
 
+    // If the span has the forceRecording attribute, always record it
+    if (attributes[SEMINTATTRS_FORCE_RECORDING]) {
+      return { decision: SamplingDecision.RECORD_AND_SAMPLED };
+    }
+
     // For all other spans, defer to the base sampler
-    return this._baseSampler.shouldSample(context, traceId, name, spanKind, attributes, links);
+    const result = this._baseSampler.shouldSample(
+      context,
+      traceId,
+      name,
+      spanKind,
+      attributes,
+      links
+    );
+
+    return result;
   }
 
   toString(): string {
@@ -63,16 +79,16 @@ export const tracer = singleton("tracer", getTracer);
 function getTracer() {
   diag.setLogger(new DiagConsoleLogger(), DiagLogLevel.ERROR);
 
-  const samplingRate = 1.0 / Math.max(parseInt(env.INTERNAL_OTEL_TRACE_SAMPING_RATE, 10), 1);
+  const samplingRate = 1.0 / Math.max(parseInt(env.INTERNAL_OTEL_TRACE_SAMPLING_RATE, 10), 1);
 
   const provider = new NodeTracerProvider({
     forceFlushTimeoutMillis: 500,
     resource: new Resource({
       [SEMRESATTRS_SERVICE_NAME]: env.SERVICE_NAME,
     }),
     sampler: new ParentBasedSampler({
-      root: new CustomWebappSampler(new TraceIdRatioBasedSampler(samplingRate)), // 5% sampling
-    }), // 5% sampling
+      root: new CustomWebappSampler(new TraceIdRatioBasedSampler(samplingRate)),
+    }),
   });
 
   if (env.INTERNAL_OTEL_TRACE_EXPORTER_URL) {
@@ -92,13 +108,15 @@ function getTracer() {
     provider.addSpanProcessor(
       new BatchSpanProcessor(exporter, {
         maxExportBatchSize: 512,
-        scheduledDelayMillis: 200,
+        scheduledDelayMillis: 1000,
         exportTimeoutMillis: 30000,
         maxQueueSize: 2048,
       })
     );
 
-    console.log(`🔦 Tracer: OTLP exporter enabled to ${env.INTERNAL_OTEL_TRACE_EXPORTER_URL}`);
+    console.log(
+      `🔦 Tracer: OTLP exporter enabled to ${env.INTERNAL_OTEL_TRACE_EXPORTER_URL} (sampling = ${samplingRate})`
+    );
   } else {
     if (env.INTERNAL_OTEL_TRACE_LOGGING_ENABLED === "1") {
       console.log(`🔦 Tracer: Logger exporter enabled`);