Batch support for Job Runs

apps/webapp/app/services/events/deliverEvent.server.ts

@@ -74,7 +74,7 @@ export class DeliverEventService {
         await Promise.all(
           matchingEventDispatchers.map((eventDispatcher) => {
             if (eventDispatcher.batcher) {
-              return workerQueue.batchEnqueue("simulateBatch", [{ seconds: 20 }], {
+              return workerQueue.batchEnqueue("events.invokeDispatchChunker", [eventRecord.id], {
                 tx,
                 jobKey: eventDispatcher.id,
                 maxPayloads: eventDispatcher.batcher.maxPayloads ?? DEFAULT_MAX_PAYLOADS,

apps/webapp/app/services/events/dispatchChunker.server.ts

@@ -0,0 +1,93 @@
+import type { PrismaClientOrTransaction } from "~/db.server";
+import { Prisma, prisma } from "~/db.server";
+import { logger } from "~/services/logger.server";
+import { workerQueue } from "../worker.server";
+
+const DEFAULT_MAX_PAYLOAD_SIZE = 2 * 1024 * 1024; // 2MB
+const DEFAULT_MAX_INTERVAL_IN_SECONDS = 20;
+
+export class DispatchChunkerService {
+  #prismaClient: PrismaClientOrTransaction;
+
+  constructor(
+    private maxPayloadSize = DEFAULT_MAX_PAYLOAD_SIZE,
+    prismaClient: PrismaClientOrTransaction = prisma
+  ) {
+    this.#prismaClient = prismaClient;
+  }
+
+  public async call(id: string, eventRecordIds: string[]) {
+    const eventDispatcher = await this.#prismaClient.eventDispatcher.findUniqueOrThrow({
+      where: {
+        id,
+      },
+      include: {
+        batcher: true,
+      },
+    });
+
+    if (!eventDispatcher.enabled) {
+      logger.debug("Event dispatcher is disabled", {
+        eventDispatcher,
+      });
+
+      return;
+    }
+
+    if (!eventDispatcher.batcher) {
+      logger.debug("Dispatcher has no batcher", {
+        eventDispatcher,
+      });
+
+      return;
+    }
+
+    const eventRecords = await prisma.$queryRaw<{ id: string; payloadSize: number }[]>`
+      SELECT id, LENGTH(payload::text) AS payloadSize FROM "EventRecord"
+      WHERE id IN (${Prisma.join(eventRecordIds)});
+    `;
+
+    let chunkSize = 0;
+    let chunkIndex = 0;
+
+    const chunks: Record<number, string[]> = { 0: [] };
+
+    const maxInterval = eventDispatcher.batcher.maxInterval ?? DEFAULT_MAX_INTERVAL_IN_SECONDS;
+
+    for (const event of eventRecords) {
+      if (chunkSize + event.payloadSize > this.maxPayloadSize) {
+        // enqueue full chunk
+        await this.#enqueueChunk(eventDispatcher.id, chunks[chunkIndex], maxInterval);
+
+        // start new chunk
+        chunkIndex++;
+        chunkSize = 0;
+        chunks[chunkIndex] = [];
+      }
+
+      chunkSize += event.payloadSize;
+      chunks[chunkIndex].push(event.id);
+    }
+
+    if (chunks[chunkIndex].length) {
+      await this.#enqueueChunk(eventDispatcher.id, chunks[chunkIndex], maxInterval);
+    }
+  }
+
+  async #enqueueChunk(dispatcherId: string, eventRecordIds: string[], deliverAfter?: number) {
+    logger.debug("Invoking batch event dispatcher", {
+      dispatcherId,
+      totalEvents: eventRecordIds.length,
+    });
+
+    await workerQueue.enqueue(
+      "events.invokeBatchDispatcher",
+      { id: dispatcherId, eventRecordIds },
+      // {
+      //   runAt: deliverAfter ? deliverAfterToDate(deliverAfter) : undefined,
+      // }
+    );
+  }
+}
+
+const deliverAfterToDate = (seconds: number) => new Date(Date.now() + seconds * 1000);

apps/webapp/app/services/events/invokeDispatcher.server.ts

@@ -1,10 +1,9 @@
-import { z } from "zod";
 import type { PrismaClientOrTransaction } from "~/db.server";
 import { prisma } from "~/db.server";
 import { logger } from "~/services/logger.server";
 import { CreateRunService } from "~/services/runs/createRun.server";
-import { InvokeEphemeralDispatcherService } from "../dispatchers/invokeEphemeralEventDispatcher.server";
 import { DispatchableSchema } from "~/models/eventDispatcher.server";
+import { InvokeEphemeralDispatcherService } from "../dispatchers/invokeEphemeralEventDispatcher.server";
 
 export class InvokeDispatcherService {
   #prismaClient: PrismaClientOrTransaction;
@@ -13,7 +12,7 @@ export class InvokeDispatcherService {
     this.#prismaClient = prismaClient;
   }
 
-  public async call(id: string, eventRecordId: string) {
+  public async call(id: string, eventRecordIds: string[]) {
     const eventDispatcher = await this.#prismaClient.eventDispatcher.findUniqueOrThrow({
       where: {
         id,
@@ -36,15 +35,27 @@ export class InvokeDispatcherService {
       return;
     }
 
-    const eventRecord = await this.#prismaClient.eventRecord.findUniqueOrThrow({
+    const eventRecords = await this.#prismaClient.eventRecord.findMany({
       where: {
-        id: eventRecordId,
+        id: { in: eventRecordIds },
+      },
+      select: {
+        id: true,
       },
     });
 
-    logger.debug("Invoking event dispatcher", {
+    if (!eventRecords.length) {
+      logger.debug("No event records found", {
+        eventDispatcher,
+        eventRecordIds,
+      });
+
+      return;
+    }
+
+    logger.debug("Invoking batch event dispatcher", {
       eventDispatcher,
-      eventRecord: eventRecord.id,
+      eventRecordIds,
     });
 
     const dispatchable = DispatchableSchema.safeParse(eventDispatcher.dispatchable);
@@ -72,7 +83,7 @@ export class InvokeDispatcherService {
         const createRunService = new CreateRunService(this.#prismaClient);
 
         await createRunService.call({
-          eventId: eventRecord.id,
+          eventIds: eventRecords.map((e) => e.id),
           job: jobVersion.job,
           version: jobVersion,
           environment: eventDispatcher.environment,
@@ -122,7 +133,7 @@ export class InvokeDispatcherService {
           const createRunService = new CreateRunService(this.#prismaClient);
 
           await createRunService.call({
-            eventId: eventRecord.id,
+            eventIds: eventRecords.map((e) => e.id),
             job: job,
             version: latestJobVersion,
             environment: eventDispatcher.environment,
@@ -132,7 +143,11 @@ export class InvokeDispatcherService {
         break;
       }
       case "EPHEMERAL": {
-        await InvokeEphemeralDispatcherService.enqueue(eventDispatcher.id, eventRecord.id);
+        if (eventRecords.length > 1) {
+          throw new Error("Ephemeral dispatcher unsupported when batching is enabled.");
+        }
+
+        await InvokeEphemeralDispatcherService.enqueue(eventDispatcher.id, eventRecords[0].id);
 
         break;
       }

apps/webapp/app/services/jobs/invokeJob.server.ts

@@ -128,7 +128,7 @@ export class InvokeJobService {
       const run = await createRunService.call(
         {
           environment,
-          eventId: eventLog.id,
+          eventIds: [eventLog.id],
           job: job,
           version,
         },

apps/webapp/app/services/jobs/testJob.server.ts

@@ -102,7 +102,7 @@ export class TestJobService {
 
         return await createRunService.call({
           environment,
-          eventId: eventLog.id,
+          eventIds: [eventLog.id],
           job: version.job,
           version,
         });

apps/webapp/app/services/runs/createRun.server.ts

@@ -14,17 +14,21 @@ export class CreateRunService {
   public async call(
     {
       environment,
-      eventId,
+      eventIds,
       job,
       version,
     }: {
       environment: AuthenticatedEnvironment;
-      eventId: string;
+      eventIds: string[];
       job: Job;
       version: JobVersion;
     },
     options: { callbackUrl?: string } = {}
   ) {
+    if (!eventIds.length) {
+      throw new Error("No event IDs provided.");
+    }
+
     const endpoint = await this.#prismaClient.endpoint.findUniqueOrThrow({
       where: {
         id: version.endpointId,
@@ -37,12 +41,18 @@ export class CreateRunService {
       },
     });
 
-    const eventRecord = await this.#prismaClient.eventRecord.findUniqueOrThrow({
+    const eventRecords = await this.#prismaClient.eventRecord.findMany({
       where: {
-        id: eventId,
+        id: { in: eventIds },
       },
     });
 
+    if (!eventRecords.length) {
+      throw new Error("No event records found.");
+    }
+
+    const firstEvent = eventRecords[0];
+
     return await $transaction(this.#prismaClient, async (tx) => {
       // Get the current max number for the given jobId
       const latestJob = await tx.jobRun.findFirst({
@@ -63,16 +73,21 @@ export class CreateRunService {
           preprocess: version.preprocessRuns,
           jobId: job.id,
           versionId: version.id,
-          eventId: eventId,
+          eventId: firstEvent.id,
           environmentId: environment.id,
           organizationId: environment.organizationId,
           projectId: environment.projectId,
           endpointId: endpoint.id,
           queueId: jobQueue.id,
-          externalAccountId: eventRecord.externalAccountId
-            ? eventRecord.externalAccountId
+          payload: JSON.stringify(
+            eventRecords.length > 1
+              ? eventRecords.map((event) => event.payload) ?? [{}]
+              : firstEvent.payload ?? {}
+          ),
+          externalAccountId: firstEvent.externalAccountId
+            ? firstEvent.externalAccountId
             : undefined,
-          isTest: eventRecord.isTest,
+          isTest: firstEvent.isTest,
           internal: job.internal,
         },
       });

apps/webapp/app/services/runs/performRunExecutionV3.server.ts

@@ -545,6 +545,7 @@ export class PerformRunExecutionV3Service {
 
       return {
         event,
+        payload: run.payload,
         job: {
           id: run.version.job.slug,
           version: run.version.version,
@@ -596,6 +597,7 @@ export class PerformRunExecutionV3Service {
 
     return {
       event,
+      payload: run.payload,
       job: {
         id: run.version.job.slug,
         version: run.version.version,

apps/webapp/app/services/runs/reRun.server.ts

@@ -69,7 +69,7 @@ export class ReRunService {
             organization: existingRun.organization,
             project: existingRun.project,
           },
-          eventId: eventLog.id,
+          eventIds: [eventLog.id],
           job: existingRun.job,
           version: existingRun.version,
         });

apps/webapp/app/services/worker.server.ts

@@ -28,6 +28,7 @@ import { ExpireDispatcherService } from "./dispatchers/expireDispatcher.server";
 import { InvokeEphemeralDispatcherService } from "./dispatchers/invokeEphemeralEventDispatcher.server";
 import { DeliverWebhookRequestService } from "./sources/deliverWebhookRequest.server";
 import { GraphileMigrationHelperService } from "./db/graphileMigrationHelper.server";
+import { DispatchChunkerService } from "./events/dispatchChunker.server";
 
 const workerCatalog = {
   indexEndpoint: z.object({
@@ -66,6 +67,11 @@ const workerCatalog = {
     ])
   ),
   deliverEvent: z.object({ id: z.string() }),
+  "events.invokeDispatchChunker": z.array(z.string()),
+  "events.invokeBatchDispatcher": z.object({
+    id: z.string(),
+    eventRecordIds: z.string().array(),
+  }),
   "events.invokeDispatcher": z.object({
     id: z.string(),
     eventRecordId: z.string(),
@@ -230,6 +236,30 @@ function getWorkerQueue() {
       },
     },
     tasks: {
+      "events.invokeDispatchChunker": {
+        priority: 0, // smaller number = higher priority
+        maxAttempts: 6,
+        queueName: (payload, jobKey) => `dispatcher-chunk:${jobKey}`,
+        handler: async (payload, job) => {
+          if (!job.key) {
+            throw new Error("Job key is required for batch jobs.");
+          }
+
+          const service = new DispatchChunkerService();
+
+          await service.call(job.key, payload);
+        },
+      },
+      "events.invokeBatchDispatcher": {
+        priority: 0, // smaller number = higher priority
+        maxAttempts: 6,
+        queueName: (payload) => `dispatcher-batch:${payload.id}`,
+        handler: async (payload, job) => {
+          const service = new InvokeDispatcherService();
+
+          await service.call(payload.id, payload.eventRecordIds);
+        },
+      },
       "events.invokeDispatcher": {
         priority: 0, // smaller number = higher priority
         maxAttempts: 6,
@@ -237,7 +267,7 @@ function getWorkerQueue() {
         handler: async (payload, job) => {
           const service = new InvokeDispatcherService();
 
-          await service.call(payload.id, payload.eventRecordId);
+          await service.call(payload.id, [payload.eventRecordId]);
         },
       },
       "events.deliverScheduled": {

packages/core/src/schemas/api.ts

@@ -576,6 +576,7 @@ export type AutoYieldConfig = z.infer<typeof AutoYieldConfigSchema>;
 
 export const RunJobBodySchema = z.object({
   event: ApiEventLogSchema,
+  payload: z.string().nullable(),
   job: z.object({
     id: z.string(),
     version: z.string(),

packages/database/prisma/migrations/20231128171050_add_payload_to_job_run/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "JobRun" ADD COLUMN     "payload" TEXT;

packages/database/prisma/schema.prisma

@@ -739,6 +739,7 @@ model JobRun {
   id       String  @id @default(cuid())
   number   Int
   internal Boolean @default(false)
+  payload  String?
 
   job   Job    @relation(fields: [jobId], references: [id], onDelete: Cascade, onUpdate: Cascade)
   jobId String

packages/trigger-sdk/src/triggerClient.ts

@@ -1217,7 +1217,9 @@ export class TriggerClient {
     );
 
     try {
-      const parsedPayload = job.trigger.event.parsePayload(body.event.payload ?? {});
+      // For compatibility with old Job Runs where payload is only available on the related Event Record
+      const payload = body.payload ? JSON.parse(body.payload) : body.event.payload;
+      const parsedPayload = job.trigger.event.parsePayload(payload ?? {});
 
       if (!context.run.isTest) {
         const verified = await job.trigger.verifyPayload(parsedPayload);