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

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(