Move metrics sending to the event repository

Author: matt-aitken

apps/webapp/app/services/clickhouse/clickhouseFactory.server.ts

@@ -324,6 +324,9 @@ function buildEventRepository(store: string, clickhouse: ClickHouse): Clickhouse
         llmMetricsFlushInterval: env.LLM_METRICS_FLUSH_INTERVAL_MS,
         llmMetricsMaxBatchSize: env.LLM_METRICS_MAX_BATCH_SIZE,
         llmMetricsMaxConcurrency: env.LLM_METRICS_MAX_CONCURRENCY,
+        otlpMetricsBatchSize: env.METRICS_CLICKHOUSE_BATCH_SIZE,
+        otlpMetricsFlushInterval: env.METRICS_CLICKHOUSE_FLUSH_INTERVAL_MS,
+        otlpMetricsMaxConcurrency: env.METRICS_CLICKHOUSE_MAX_CONCURRENCY,
         version: "v1",
       });
     }
@@ -344,6 +347,9 @@ function buildEventRepository(store: string, clickhouse: ClickHouse): Clickhouse
         llmMetricsFlushInterval: env.LLM_METRICS_FLUSH_INTERVAL_MS,
         llmMetricsMaxBatchSize: env.LLM_METRICS_MAX_BATCH_SIZE,
         llmMetricsMaxConcurrency: env.LLM_METRICS_MAX_CONCURRENCY,
+        otlpMetricsBatchSize: env.METRICS_CLICKHOUSE_BATCH_SIZE,
+        otlpMetricsFlushInterval: env.METRICS_CLICKHOUSE_FLUSH_INTERVAL_MS,
+        otlpMetricsMaxConcurrency: env.METRICS_CLICKHOUSE_MAX_CONCURRENCY,
         version: "v2",
       });
     }

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

@@ -1,6 +1,7 @@
 import type {
   ClickHouse,
   LlmMetricsV1Input,
+  MetricsV1Input,
   TaskEventDetailedSummaryV1Result,
   TaskEventDetailsV1Result,
   TaskEventSummaryV1Result,
@@ -91,6 +92,10 @@ export type ClickhouseEventRepositoryConfig = {
   llmMetricsFlushInterval?: number;
   llmMetricsMaxBatchSize?: number;
   llmMetricsMaxConcurrency?: number;
+  /** OTLP / task metrics_v1 flush scheduler config */
+  otlpMetricsBatchSize?: number;
+  otlpMetricsFlushInterval?: number;
+  otlpMetricsMaxConcurrency?: number;
 };
 
 /**
@@ -102,6 +107,7 @@ export class ClickhouseEventRepository implements IEventRepository {
   private _config: ClickhouseEventRepositoryConfig;
   private readonly _flushScheduler: DynamicFlushScheduler<TaskEventV1Input | TaskEventV2Input>;
   private readonly _llmMetricsFlushScheduler: DynamicFlushScheduler<LlmMetricsV1Input>;
+  private readonly _otlpMetricsFlushScheduler: DynamicFlushScheduler<MetricsV1Input>;
   private _tracer: Tracer;
   private _version: "v1" | "v2";
 
@@ -137,6 +143,15 @@ export class ClickhouseEventRepository implements IEventRepository {
       memoryPressureThreshold: config.llmMetricsMaxBatchSize ?? 10000,
       loadSheddingEnabled: false,
     });
+
+    this._otlpMetricsFlushScheduler = new DynamicFlushScheduler({
+      batchSize: config.otlpMetricsBatchSize ?? 10000,
+      flushInterval: config.otlpMetricsFlushInterval ?? 1000,
+      callback: this.#flushOtelMetricsBatch.bind(this),
+      minConcurrency: 1,
+      maxConcurrency: config.otlpMetricsMaxConcurrency ?? 3,
+      loadSheddingEnabled: false,
+    });
   }
 
   get version() {
@@ -251,6 +266,27 @@ export class ClickhouseEventRepository implements IEventRepository {
     });
   }
 
+  async #flushOtelMetricsBatch(flushId: string, rows: MetricsV1Input[]) {
+    await startSpan(this._tracer, "flushOtelMetricsBatch", async (span) => {
+      span.setAttribute("flush_id", flushId);
+      span.setAttribute("row_count", rows.length);
+
+      const [insertError] = await this._clickhouse.metrics.insert(rows, {
+        params: {
+          clickhouse_settings: this.#getClickhouseInsertSettings(),
+        },
+      });
+
+      if (insertError) {
+        throw insertError;
+      }
+
+      logger.info("ClickhouseEventRepository.flushOtelMetricsBatch Inserted OTLP metrics batch", {
+        rows: rows.length,
+      });
+    });
+  }
+
   #createLlmMetricsInput(event: CreateEventInput): LlmMetricsV1Input {
     const llmMetrics = event._llmMetrics!;
 
@@ -326,6 +362,11 @@ export class ClickhouseEventRepository implements IEventRepository {
     this.insertMany(events);
   }
 
+  insertManyMetrics(rows: MetricsV1Input[]): void {
+    if (rows.length === 0) return;
+    this._otlpMetricsFlushScheduler.addToBatch(rows);
+  }
+
   private createEventToTaskEventV1Input(event: CreateEventInput): TaskEventV1Input[] {
     return [
       {

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

@@ -18,6 +18,7 @@ import {
   unflattenAttributes,
 } from "@trigger.dev/core/v3";
 import { serializeTraceparent } from "@trigger.dev/core/v3/isomorphic";
+import type { MetricsV1Input } from "@internal/clickhouse";
 import { Prisma, TaskEvent, TaskEventKind } from "@trigger.dev/database";
 import { nanoid } from "nanoid";
 import { Gauge } from "prom-client";
@@ -159,6 +160,8 @@ export class EventRepository implements IEventRepository {
     await this.#flushBatchWithReturn(nanoid(), events.map(this.#createableEventToPrismaEvent));
   }
 
+  insertManyMetrics(_rows: MetricsV1Input[]): void {}
+
   async completeSuccessfulRunEvent({ run, endTime }: { run: CompleteableTaskRun; endTime?: Date }) {
     const startTime = convertDateToNanoseconds(run.createdAt);
 

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

@@ -14,6 +14,7 @@ import type {
   TaskEventStatus,
   TaskRun,
 } from "@trigger.dev/database";
+import type { MetricsV1Input } from "@internal/clickhouse";
 import type { DetailedTraceEvent, TaskEventStoreTable } from "../taskEventStore.server";
 export type { ExceptionEventProperties };
 
@@ -347,6 +348,7 @@ export interface IEventRepository {
   // Event insertion methods
   insertMany(events: CreateEventInput[]): void;
   insertManyImmediate(events: CreateEventInput[]): Promise<void>;
+  insertManyMetrics(rows: MetricsV1Input[]): void;
 
   // Run event completion methods
   completeSuccessfulRunEvent(params: { run: CompleteableTaskRun; endTime?: Date }): Promise<void>;

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

@@ -19,14 +19,8 @@ import {
   Status_StatusCode,
 } from "@trigger.dev/otlp-importer";
 import type { MetricsV1Input } from "@internal/clickhouse";
-import { ClickHouse } from "@internal/clickhouse";
 import { logger } from "~/services/logger.server";
-import {
-  clickhouseFactory,
-  ClickhouseFactory,
-  getDefaultClickhouseClient,
-} from "~/services/clickhouse/clickhouseFactory.server";
-import { DynamicFlushScheduler } from "./dynamicFlushScheduler.server";
+import { clickhouseFactory, type ClickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server";
 
 import { generateSpanId } from "./eventRepository/common.server";
 import type {
@@ -43,19 +37,13 @@ import { singleton } from "~/utils/singleton";
 
 type OTLPExporterConfig = {
   clickhouseFactory: ClickhouseFactory;
-  metrics: {
-    batchSize: number;
-    flushInterval: number;
-    maxConcurrency: number;
-  };
   verbose: boolean;
   spanAttributeValueLengthLimit: number;
 };
 
 class OTLPExporter {
   private _tracer: Tracer;
   private readonly _clickhouseFactory: ClickhouseFactory;
-  private readonly _defaultMetricsFlushScheduler: DynamicFlushScheduler<MetricsV1Input>;
   private readonly _verbose: boolean;
   private readonly _spanAttributeValueLengthLimit: number;
 
@@ -64,16 +52,6 @@ class OTLPExporter {
     this._clickhouseFactory = config.clickhouseFactory;
     this._verbose = config.verbose;
     this._spanAttributeValueLengthLimit = config.spanAttributeValueLengthLimit;
-    this._defaultMetricsFlushScheduler = new DynamicFlushScheduler<MetricsV1Input>({
-      batchSize: config.metrics.batchSize,
-      flushInterval: config.metrics.flushInterval,
-      callback: async (_flushId, batch) => {
-        await config.metrics.clickhouse.metrics.insert(batch);
-      },
-      minConcurrency: 1,
-      maxConcurrency: config.metrics.maxConcurrency,
-      loadSheddingEnabled: false,
-    });
   }
 
   async exportTraces(request: ExportTraceServiceRequest): Promise<ExportTraceServiceResponse> {
@@ -96,19 +74,19 @@ class OTLPExporter {
 
   async exportMetrics(request: ExportMetricsServiceRequest): Promise<ExportMetricsServiceResponse> {
     return await startSpan(this._tracer, "exportMetrics", async (span) => {
-      const rows = this.#filterResourceMetrics(request.resourceMetrics).flatMap(
-        (resourceMetrics) => {
-          return convertMetricsToClickhouseRows(
+      const metricsWithStores = this.#filterResourceMetrics(request.resourceMetrics).map(
+        (resourceMetrics) =>
+          convertResourceMetricsToRowsWithStore(
             resourceMetrics,
             this._spanAttributeValueLengthLimit
-          );
-        }
+          )
       );
 
-      span.setAttribute("metric_row_count", rows.length);
+      const rowCount = metricsWithStores.reduce((acc, m) => acc + m.rows.length, 0);
+      span.setAttribute("metric_row_count", rowCount);
 
-      if (rows.length > 0) {
-        this._defaultMetricsFlushScheduler.addToBatch(rows);
+      if (rowCount > 0) {
+        await this.#exportMetricRows(metricsWithStores);
       }
 
       return ExportMetricsServiceResponse.create();
@@ -177,6 +155,37 @@ class OTLPExporter {
     return eventCount;
   }
 
+  async #exportMetricRows(
+    metricsWithStores: { rows: MetricsV1Input[]; taskEventStore: string }[]
+  ): Promise<void> {
+    const routeCache = new Map<string, { key: string; repository: IEventRepository }>();
+    const groups = new Map<string, { repository: IEventRepository; rows: MetricsV1Input[] }>();
+    for (const { rows, taskEventStore } of metricsWithStores) {
+      for (const row of rows) {
+        const routeKey = `${row.organization_id}\0${taskEventStore}`;
+        let resolved = routeCache.get(routeKey);
+        if (!resolved) {
+          resolved = this._clickhouseFactory.getEventRepositoryForOrganizationSync(
+            taskEventStore,
+            row.organization_id
+          );
+          routeCache.set(routeKey, resolved);
+        }
+
+        let group = groups.get(resolved.key);
+        if (!group) {
+          group = { repository: resolved.repository, rows: [] };
+          groups.set(resolved.key, group);
+        }
+        group.rows.push(row);
+      }
+    }
+
+    for (const [, { repository, rows }] of groups) {
+      repository.insertManyMetrics(rows);
+    }
+  }
+
   #logEventsVerbose(events: CreateEventInput[], prefix: string) {
     if (!this._verbose) return;
 
@@ -592,6 +601,21 @@ function convertMetricsToClickhouseRows(
   return rows;
 }
 
+function convertResourceMetricsToRowsWithStore(
+  resourceMetrics: ResourceMetrics,
+  spanAttributeValueLengthLimit: number
+): { rows: MetricsV1Input[]; taskEventStore: string } {
+  const resourceAttributes = resourceMetrics.resource?.attributes ?? [];
+  const taskEventStore =
+    extractStringAttribute(resourceAttributes, [SemanticInternalAttributes.TASK_EVENT_STORE]) ??
+    env.EVENT_REPOSITORY_DEFAULT_STORE;
+
+  return {
+    rows: convertMetricsToClickhouseRows(resourceMetrics, spanAttributeValueLengthLimit),
+    taskEventStore,
+  };
+}
+
 // Prefixes injected by TaskContextMetricExporter — these are extracted into
 // the nested `trigger` key and should not appear as top-level user attributes.
 const INTERNAL_METRIC_ATTRIBUTE_PREFIXES = ["ctx.", "worker."];
@@ -1203,11 +1227,6 @@ export const otlpExporter = singleton("otlpExporter", initializeOTLPExporter);
 function initializeOTLPExporter() {
   return new OTLPExporter({
     clickhouseFactory,
-    metrics: {
-      batchSize: env.METRICS_CLICKHOUSE_BATCH_SIZE,
-      flushInterval: env.METRICS_CLICKHOUSE_FLUSH_INTERVAL_MS,
-      maxConcurrency: env.METRICS_CLICKHOUSE_MAX_CONCURRENCY,
-    },
     verbose: process.env.OTLP_EXPORTER_VERBOSE === "1",
     spanAttributeValueLengthLimit: process.env.SERVER_OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT
       ? parseInt(process.env.SERVER_OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT, 10)