Speed up the Run log queries (#1126)

* Added traceId index to TaskEvent

* Get less data and get the trace from the run

* Use the read replica in the EventRepository

* Use the regular client to get the run because the redirect from test can be very fast

Author: matt-aitken

apps/webapp/app/db.server.ts

@@ -70,7 +70,9 @@ export { Prisma };
 
 export const prisma = singleton("prisma", getClient);
 
-export const $replica: Omit<PrismaClient, "$transaction"> = singleton(
+export type PrismaReplicaClient = Omit<PrismaClient, "$transaction">;
+
+export const $replica: PrismaReplicaClient = singleton(
   "replica",
   () => getReplicaClient() ?? prisma
 );

apps/webapp/app/presenters/v3/SpanPresenter.server.ts

@@ -1,29 +1,26 @@
 import { prettyPrintPacket } from "@trigger.dev/core/v3";
 import { PrismaClient, prisma } from "~/db.server";
 import { eventRepository } from "~/v3/eventRepository.server";
+import { BasePresenter } from "./basePresenter.server";
 
 type Result = Awaited<ReturnType<SpanPresenter["call"]>>;
 export type Span = NonNullable<Result>["event"];
 
-export class SpanPresenter {
-  #prismaClient: PrismaClient;
-
-  constructor(prismaClient: PrismaClient = prisma) {
-    this.#prismaClient = prismaClient;
-  }
-
+export class SpanPresenter extends BasePresenter {
   public async call({
     userId,
     projectSlug,
     organizationSlug,
     spanId,
+    runFriendlyId,
   }: {
     userId: string;
     projectSlug: string;
     organizationSlug: string;
     spanId: string;
+    runFriendlyId: string;
   }) {
-    const project = await this.#prismaClient.project.findUnique({
+    const project = await this._replica.project.findUnique({
       where: {
         slug: projectSlug,
       },
@@ -33,7 +30,20 @@ export class SpanPresenter {
       throw new Error("Project not found");
     }
 
-    const span = await eventRepository.getSpan(spanId);
+    const run = await this._prisma.taskRun.findFirst({
+      select: {
+        traceId: true,
+      },
+      where: {
+        friendlyId: runFriendlyId,
+      },
+    });
+
+    if (!run) {
+      return;
+    }
+
+    const span = await eventRepository.getSpan(spanId, run.traceId);
 
     if (!span) {
       return;

apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam.spans.$spanParam/route.tsx

@@ -4,11 +4,11 @@ import {
   QueueListIcon,
   StopCircleIcon,
 } from "@heroicons/react/20/solid";
-import { useFetcher, useParams } from "@remix-run/react";
+import { useParams } from "@remix-run/react";
 import { LoaderFunctionArgs } from "@remix-run/server-runtime";
 import { formatDurationNanoseconds, nanosecondsToMilliseconds } from "@trigger.dev/core/v3";
 import { useEffect } from "react";
-import { typedjson, useTypedFetcher, useTypedLoaderData } from "remix-typedjson";
+import { typedjson, useTypedFetcher } from "remix-typedjson";
 import { ExitIcon } from "~/assets/icons/ExitIcon";
 import { CodeBlock } from "~/components/code/CodeBlock";
 import { EnvironmentLabel } from "~/components/environments/EnvironmentLabel";
@@ -46,6 +46,7 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
     organizationSlug,
     projectSlug: projectParam,
     spanId: spanParam,
+    runFriendlyId: runParam,
   });
 
   if (!span) {

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

@@ -22,7 +22,7 @@ import { Prisma, TaskEvent, TaskEventStatus, type TaskEventKind } from "@trigger
 import Redis, { RedisOptions } from "ioredis";
 import { createHash } from "node:crypto";
 import { EventEmitter } from "node:stream";
-import { PrismaClient, prisma } from "~/db.server";
+import { $replica, PrismaClient, PrismaReplicaClient, prisma } from "~/db.server";
 import { env } from "~/env.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
@@ -101,9 +101,26 @@ export type QueryOptions = Prisma.TaskEventWhereInput;
 
 export type TaskEventRecord = TaskEvent;
 
-export type QueriedEvent = TaskEvent;
+export type QueriedEvent = Prisma.TaskEventGetPayload<{
+  select: {
+    id: true;
+    spanId: true;
+    parentId: true;
+    runId: true;
+    idempotencyKey: true;
+    message: true;
+    style: true;
+    startTime: true;
+    duration: true;
+    isError: true;
+    isPartial: true;
+    isCancelled: true;
+    level: true;
+    events: true;
+  };
+}>;
 
-export type PreparedEvent = Omit<TaskEventRecord, "events" | "style" | "duration"> & {
+export type PreparedEvent = Omit<QueriedEvent, "events" | "style" | "duration"> & {
   duration: number;
   events: SpanEvents;
   style: TaskEventStyle;
@@ -160,7 +177,11 @@ export class EventRepository {
     return this._subscriberCount;
   }
 
-  constructor(private db: PrismaClient = prisma, private readonly _config: EventRepoConfig) {
+  constructor(
+    private db: PrismaClient = prisma,
+    private readReplica: PrismaReplicaClient = $replica,
+    private readonly _config: EventRepoConfig
+  ) {
     this._flushScheduler = new DynamicFlushScheduler({
       batchSize: _config.batchSize,
       flushInterval: _config.batchInterval,
@@ -349,7 +370,23 @@ export class EventRepository {
   }
 
   public async getTraceSummary(traceId: string): Promise<TraceSummary | undefined> {
-    const events = await this.db.taskEvent.findMany({
+    const events = await this.readReplica.taskEvent.findMany({
+      select: {
+        id: true,
+        spanId: true,
+        parentId: true,
+        runId: true,
+        idempotencyKey: true,
+        message: true,
+        style: true,
+        startTime: true,
+        duration: true,
+        isError: true,
+        isPartial: true,
+        isCancelled: true,
+        level: true,
+        events: true,
+      },
       where: {
         traceId,
       },
@@ -407,29 +444,16 @@ export class EventRepository {
 
   // A Span can be cancelled if it is partial and has a parent that is cancelled
   // And a span's duration, if it is partial and has a cancelled parent, is the time between the start of the span and the time of the cancellation event of the parent
-  public async getSpan(spanId: string) {
-    const traceSearch = await this.db.taskEvent.findFirst({
-      where: {
-        spanId,
-      },
-      select: {
-        traceId: true,
-      },
-    });
-
-    if (!traceSearch) {
-      return;
-    }
-
-    const traceSummary = await this.getTraceSummary(traceSearch.traceId);
+  public async getSpan(spanId: string, traceId: string) {
+    const traceSummary = await this.getTraceSummary(traceId);
 
     const span = traceSummary?.spans.find((span) => span.id === spanId);
 
     if (!span) {
       return;
     }
 
-    const fullEvent = await this.db.taskEvent.findUnique({
+    const fullEvent = await this.readReplica.taskEvent.findUnique({
       where: {
         id: span.recordId,
       },
@@ -814,7 +838,7 @@ export class EventRepository {
 export const eventRepository = singleton("eventRepo", initializeEventRepo);
 
 function initializeEventRepo() {
-  const repo = new EventRepository(prisma, {
+  const repo = new EventRepository(prisma, $replica, {
     batchSize: env.EVENTS_BATCH_SIZE,
     batchInterval: env.EVENTS_BATCH_INTERVAL,
     retentionInDays: env.EVENTS_DEFAULT_LOG_RETENTION,

packages/database/prisma/migrations/20240523135511_added_task_event_trace_id_index/migration.sql

@@ -0,0 +1,2 @@
+-- CreateIndex
+CREATE INDEX "TaskEvent_traceId_idx" ON "TaskEvent"("traceId");

packages/database/prisma/schema.prisma

@@ -1898,6 +1898,9 @@ model TaskEvent {
   payloadType String?
 
   createdAt DateTime @default(now())
+
+  /// Used on the run page
+  @@index([traceId])
 }
 
 enum TaskEventLevel {