fix: Incorrect/missing logs on new schedule engine runs (#2185)

* Fix incorrect logs on new schedule engine triggered taskss

Also added the ability to recover schedules in the schedule engine via an Admin API endpoint in the new schedule engine

* Fixed schedule recovery failing test

Author: ericallam

apps/webapp/app/routes/admin.api.v1.environments.$environmentId.schedules.recover.ts

@@ -0,0 +1,58 @@
+import { ActionFunctionArgs, json, LoaderFunctionArgs } from "@remix-run/server-runtime";
+import { z } from "zod";
+import { prisma } from "~/db.server";
+import { authenticateApiRequestWithPersonalAccessToken } from "~/services/personalAccessToken.server";
+import { scheduleEngine } from "~/v3/scheduleEngine.server";
+
+const ParamsSchema = z.object({
+  environmentId: z.string(),
+});
+
+export async function action({ request, params }: ActionFunctionArgs) {
+  // Next authenticate the request
+  const authenticationResult = await authenticateApiRequestWithPersonalAccessToken(request);
+
+  if (!authenticationResult) {
+    return json({ error: "Invalid or Missing API key" }, { status: 401 });
+  }
+
+  const user = await prisma.user.findUnique({
+    where: {
+      id: authenticationResult.userId,
+    },
+  });
+
+  if (!user) {
+    return json({ error: "Invalid or Missing API key" }, { status: 401 });
+  }
+
+  if (!user.admin) {
+    return json({ error: "You must be an admin to perform this action" }, { status: 403 });
+  }
+
+  const parsedParams = ParamsSchema.parse(params);
+
+  const environment = await prisma.runtimeEnvironment.findFirst({
+    where: {
+      id: parsedParams.environmentId,
+    },
+    include: {
+      organization: true,
+      project: true,
+    },
+  });
+
+  if (!environment) {
+    return json({ error: "Environment not found" }, { status: 404 });
+  }
+
+  const results = await scheduleEngine.recoverSchedulesInEnvironment(
+    environment.projectId,
+    environment.id
+  );
+
+  return json({
+    success: true,
+    results,
+  });
+}

apps/webapp/app/runEngine/concerns/traceEvents.server.ts

@@ -39,6 +39,9 @@ export class DefaultTraceEventsConcern implements TraceEventConcern {
         },
         incomplete: true,
         immediate: true,
+        startTime: request.options?.overrideCreatedAt
+          ? BigInt(request.options.overrideCreatedAt.getTime()) * BigInt(1000000)
+          : undefined,
       },
       async (event, traceContext, traceparent) => {
         return await callback({

apps/webapp/app/runEngine/types.ts

@@ -21,6 +21,7 @@ export type TriggerTaskServiceOptions = {
   runFriendlyId?: string;
   skipChecks?: boolean;
   oneTimeUseToken?: string;
+  overrideCreatedAt?: Date;
 };
 
 // domain/triggerTask.ts

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

@@ -970,7 +970,7 @@ export class EventRepository {
     const propagatedContext = extractContextFromCarrier(options.context ?? {});
 
     const start = process.hrtime.bigint();
-    const startTime = getNowInNanoseconds();
+    const startTime = options.startTime ?? getNowInNanoseconds();
 
     const traceId = options.spanParentAsLink
       ? this.generateTraceId()

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

@@ -312,6 +312,9 @@ export class TriggerTaskServiceV1 extends BaseService {
             },
             incomplete: true,
             immediate: true,
+            startTime: options.overrideCreatedAt
+              ? BigInt(options.overrideCreatedAt.getTime()) * BigInt(1000000)
+              : undefined,
           },
           async (event, traceContext, traceparent) => {
             const run = await autoIncrementCounter.incrementInTransaction(

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

@@ -82,17 +82,20 @@ export class TaskEventStore {
     let finalWhere: Prisma.TaskEventWhereInput = where;
 
     if (table === "taskEventPartitioned") {
-      // Add 1 minute to endCreatedAt to make sure we include all events in the range.
+      // Add buffer to start and end of the range to make sure we include all events in the range.
       const end = endCreatedAt
         ? new Date(endCreatedAt.getTime() + env.TASK_EVENT_PARTITIONED_WINDOW_IN_SECONDS * 1000)
         : new Date();
+      const startCreatedAtWithBuffer = new Date(
+        startCreatedAt.getTime() - env.TASK_EVENT_PARTITIONED_WINDOW_IN_SECONDS * 1000
+      );
 
       finalWhere = {
         AND: [
           where,
           {
             createdAt: {
-              gte: startCreatedAt,
+              gte: startCreatedAtWithBuffer,
               lt: end,
             },
           },
@@ -138,6 +141,11 @@ export class TaskEventStore {
       options?.includeDebugLogs === false || options?.includeDebugLogs === undefined;
 
     if (table === "taskEventPartitioned") {
+      const createdAtBufferInMillis = env.TASK_EVENT_PARTITIONED_WINDOW_IN_SECONDS * 1000;
+      const startCreatedAtWithBuffer = new Date(startCreatedAt.getTime() - createdAtBufferInMillis);
+      const $endCreatedAt = endCreatedAt ?? new Date();
+      const endCreatedAtWithBuffer = new Date($endCreatedAt.getTime() + createdAtBufferInMillis);
+
       return await this.readReplica.$queryRaw<TraceEvent[]>`
         SELECT
           "spanId",
@@ -158,11 +166,8 @@ export class TaskEventStore {
         FROM "TaskEventPartitioned"
         WHERE
           "traceId" = ${traceId}
-          AND "createdAt" >= ${startCreatedAt.toISOString()}::timestamp
-          AND "createdAt" < ${(endCreatedAt
-            ? new Date(endCreatedAt.getTime() + env.TASK_EVENT_PARTITIONED_WINDOW_IN_SECONDS * 1000)
-            : new Date()
-          ).toISOString()}::timestamp
+          AND "createdAt" >= ${startCreatedAtWithBuffer.toISOString()}::timestamp
+          AND "createdAt" < ${endCreatedAtWithBuffer.toISOString()}::timestamp
           ${
             filterDebug
               ? Prisma.sql`AND \"kind\" <> CAST('LOG'::text AS "public"."TaskEventKind")`

internal-packages/schedule-engine/src/engine/index.ts

@@ -8,7 +8,7 @@ import {
   Tracer,
 } from "@internal/tracing";
 import { Logger } from "@trigger.dev/core/logger";
-import { PrismaClient } from "@trigger.dev/database";
+import { PrismaClient, TaskSchedule, TaskScheduleInstance } from "@trigger.dev/database";
 import { Worker, type JobHandlerParams } from "@trigger.dev/redis-worker";
 import { calculateDistributedExecutionTime } from "./distributedScheduling.js";
 import { calculateNextScheduledTimestamp, nextScheduledTimestamps } from "./scheduleCalculation.js";
@@ -645,6 +645,140 @@ export class ScheduleEngine {
     });
   }
 
+  public recoverSchedulesInEnvironment(projectId: string, environmentId: string) {
+    return startSpan(this.tracer, "recoverSchedulesInEnvironment", async (span) => {
+      this.logger.info("Recovering schedules in environment", {
+        environmentId,
+        projectId,
+      });
+
+      span.setAttribute("environmentId", environmentId);
+
+      const schedules = await this.prisma.taskSchedule.findMany({
+        where: {
+          projectId,
+          instances: {
+            some: {
+              environmentId,
+            },
+          },
+        },
+        select: {
+          id: true,
+          generatorExpression: true,
+          instances: {
+            select: {
+              id: true,
+              environmentId: true,
+              lastScheduledTimestamp: true,
+              nextScheduledTimestamp: true,
+            },
+          },
+        },
+      });
+
+      const instancesWithSchedule = schedules
+        .map((schedule) => ({
+          schedule,
+          instance: schedule.instances.find((instance) => instance.environmentId === environmentId),
+        }))
+        .filter((instance) => instance.instance) as Array<{
+        schedule: Omit<(typeof schedules)[number], "instances">;
+        instance: NonNullable<(typeof schedules)[number]["instances"][number]>;
+      }>;
+
+      if (instancesWithSchedule.length === 0) {
+        this.logger.info("No instances found for environment", {
+          environmentId,
+          projectId,
+        });
+
+        return {
+          recovered: [],
+          skipped: [],
+        };
+      }
+
+      const results = {
+        recovered: [],
+        skipped: [],
+      } as { recovered: string[]; skipped: string[] };
+
+      for (const { instance, schedule } of instancesWithSchedule) {
+        this.logger.info("Recovering schedule", {
+          schedule,
+          instance,
+        });
+
+        const [recoverError, result] = await tryCatch(
+          this.#recoverTaskScheduleInstance({ instance, schedule })
+        );
+
+        if (recoverError) {
+          this.logger.error("Error recovering schedule", {
+            error: recoverError instanceof Error ? recoverError.message : String(recoverError),
+          });
+
+          span.setAttribute("recover_error", true);
+          span.setAttribute(
+            "recover_error_message",
+            recoverError instanceof Error ? recoverError.message : String(recoverError)
+          );
+        } else {
+          span.setAttribute("recover_success", true);
+
+          if (result === "recovered") {
+            results.recovered.push(instance.id);
+          } else {
+            results.skipped.push(instance.id);
+          }
+        }
+      }
+
+      return results;
+    });
+  }
+
+  async #recoverTaskScheduleInstance({
+    instance,
+    schedule,
+  }: {
+    instance: {
+      id: string;
+      environmentId: string;
+      lastScheduledTimestamp: Date | null;
+      nextScheduledTimestamp: Date | null;
+    };
+    schedule: { id: string; generatorExpression: string };
+  }) {
+    // inspect the schedule worker to see if there is a job for this instance
+    const job = await this.worker.getJob(`scheduled-task-instance:${instance.id}`);
+
+    if (job) {
+      this.logger.info("Job already exists for instance", {
+        instanceId: instance.id,
+        job,
+        schedule,
+      });
+
+      return "skipped";
+    }
+
+    this.logger.info("No job found for instance, registering next run", {
+      instanceId: instance.id,
+      schedule,
+    });
+
+    // If the job does not exist, register the next run
+    await this.registerNextTaskScheduleInstance({ instanceId: instance.id });
+
+    return "recovered";
+  }
+
+  async getJob(id: string) {
+    return this.worker.getJob(id);
+  }
+
   async quit() {
     this.logger.info("Shutting down schedule engine");
 

internal-packages/schedule-engine/test/scheduleEngine.test.ts

@@ -1,9 +1,9 @@
 import { containerTest } from "@internal/testcontainers";
 import { trace } from "@internal/tracing";
-import { describe, expect, vi } from "vitest";
-import { ScheduleEngine } from "../src/index.js";
 import { setTimeout } from "timers/promises";
+import { describe, expect, vi } from "vitest";
 import { TriggerScheduledTaskParams } from "../src/engine/types.js";
+import { ScheduleEngine } from "../src/index.js";
 
 describe("ScheduleEngine Integration", () => {
   containerTest(