wehub-resource-sync
2359 lines
78 KiB
TypeScript
2359 lines
78 KiB
TypeScript
import { describe, expect, it, vi } from "vitest";
|
|
|
|
// Mock the db prisma client - needs to be before other imports
|
|
vi.mock("~/db.server", () => ({
|
|
prisma: {},
|
|
$replica: {},
|
|
}));
|
|
|
|
vi.mock("~/services/platform.v3.server", async (importOriginal) => {
|
|
const actual = (await importOriginal()) as Record<string, unknown>;
|
|
return {
|
|
...actual,
|
|
getEntitlement: vi.fn(),
|
|
};
|
|
});
|
|
|
|
import { RunEngine } from "@internal/run-engine";
|
|
import { PostgresRunStore } from "@internal/run-store";
|
|
import { setupAuthenticatedEnvironment } from "@internal/run-engine/tests";
|
|
// Per-test redis isolation: each test runs its own RunEngine whose background work outlives the test
|
|
// body. NoClickhouse because this suite never touches ClickHouse - skips the worker-scoped boot+migrate.
|
|
import { containerTestWithIsolatedRedisNoClickhouse as containerTest } from "@internal/testcontainers";
|
|
import { trace } from "@opentelemetry/api";
|
|
import type { PrismaClient } from "@trigger.dev/database";
|
|
import { BatchId } from "@trigger.dev/core/v3/isomorphic";
|
|
import {
|
|
StreamBatchItemsService,
|
|
createNdjsonParserStream,
|
|
streamToAsyncIterable,
|
|
extractIndexAndTask,
|
|
type OversizedItemMarker,
|
|
} from "../../app/runEngine/services/streamBatchItems.server";
|
|
import { ServiceValidationError } from "../../app/v3/services/baseService.server";
|
|
import {
|
|
BatchPayloadProcessor,
|
|
type BatchPayloadProcessResult,
|
|
} from "../../app/runEngine/concerns/batchPayloads.server";
|
|
import { setTimeout as sleep } from "node:timers/promises";
|
|
|
|
// 120s: a cold per-test container boot counts against the test's own timeout, and under CI Docker
|
|
// contention 30s was too tight. Matches the run-engine convention for this footprint.
|
|
vi.setConfig({ testTimeout: 120_000 });
|
|
|
|
describe("StreamBatchItemsService", () => {
|
|
async function createBatch(
|
|
prisma: PrismaClient,
|
|
environmentId: string,
|
|
options: {
|
|
runCount: number;
|
|
status?: "PENDING" | "PROCESSING" | "COMPLETED" | "PARTIAL_FAILED" | "ABORTED";
|
|
sealed?: boolean;
|
|
processingCompletedAt?: Date | null;
|
|
}
|
|
) {
|
|
const { id, friendlyId } = BatchId.generate();
|
|
|
|
const batch = await prisma.batchTaskRun.create({
|
|
data: {
|
|
id,
|
|
friendlyId,
|
|
runtimeEnvironmentId: environmentId,
|
|
status: options.status ?? "PENDING",
|
|
runCount: options.runCount,
|
|
expectedCount: options.runCount,
|
|
runIds: [],
|
|
batchVersion: "runengine:v2",
|
|
sealed: options.sealed ?? false,
|
|
processingCompletedAt: options.processingCompletedAt ?? null,
|
|
},
|
|
});
|
|
|
|
return batch;
|
|
}
|
|
|
|
async function* itemsToAsyncIterable(
|
|
items: Array<{ task: string; payload: string; index: number }>
|
|
) {
|
|
for (const item of items) {
|
|
yield item;
|
|
}
|
|
}
|
|
|
|
function makeItems(count: number, taskId = "test-task") {
|
|
return Array.from({ length: count }, (_, index) => ({
|
|
task: taskId,
|
|
payload: JSON.stringify({ i: index }),
|
|
index,
|
|
}));
|
|
}
|
|
|
|
/**
|
|
* Build a RunEngine for tests. The worker is disabled and no item-processing
|
|
* callback is wired, so streamed items stay enqueued (the BatchQueue consumer
|
|
* doesn't drain them) — keeping the enqueued-count assertions deterministic.
|
|
*/
|
|
function buildEngine(prisma: PrismaClient, redisOptions: any) {
|
|
return new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: { redis: redisOptions },
|
|
runLock: { redis: redisOptions },
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: { redis: redisOptions },
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
}
|
|
|
|
containerTest(
|
|
"should seal batch successfully when no race condition",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// Create a batch
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
// Initialize the batch in Redis
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: authenticatedEnvironment.id,
|
|
environmentType: authenticatedEnvironment.type,
|
|
organizationId: authenticatedEnvironment.organizationId,
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runCount: 2,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
// Enqueue items directly to Redis (bypassing the service's item processing)
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item1" }),
|
|
payloadType: "application/json",
|
|
});
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item2" }),
|
|
payloadType: "application/json",
|
|
});
|
|
|
|
// Create service with our test engine and prisma
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
});
|
|
|
|
// Create an empty items iterator since items are already enqueued
|
|
const items = itemsToAsyncIterable([]);
|
|
|
|
const result = await service.call(authenticatedEnvironment, batch.friendlyId, items, {
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
});
|
|
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
|
|
// Verify the batch is sealed in the database
|
|
const updatedBatch = await prisma.batchTaskRun.findUnique({
|
|
where: { id: batch.id },
|
|
});
|
|
|
|
expect(updatedBatch?.sealed).toBe(true);
|
|
expect(updatedBatch?.status).toBe("PROCESSING");
|
|
expect(updatedBatch?.sealedAt).toBeDefined();
|
|
expect(updatedBatch?.processingStartedAt).toBeDefined();
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=true when batch is already sealed and PROCESSING (Phase 2 retry idempotency)",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// Simulate the SDK retrying Phase 2 after the original request succeeded:
|
|
// the original request already sealed the batch and moved it to PROCESSING.
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "PROCESSING",
|
|
sealed: true,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
});
|
|
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
// The retry should be treated as success — the original request already
|
|
// enqueued every item, so the SDK should stop retrying.
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
expect(result.itemsAccepted).toBe(0);
|
|
expect(result.itemsDeduplicated).toBe(0);
|
|
expect(result.runCount).toBe(2);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=true when batch is COMPLETED before Phase 2 retry arrives (TRI-9944)",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// The customer-reported scenario: single-item batch where the original
|
|
// Phase 2 request succeeded server-side, the run executed fast, the batch
|
|
// flipped to COMPLETED, then the lost-response SDK retry hits us.
|
|
// Note: tryCompleteBatch sets status=COMPLETED but does NOT set sealed=true.
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 1,
|
|
status: "COMPLETED",
|
|
sealed: false,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
});
|
|
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
expect(result.itemsAccepted).toBe(0);
|
|
expect(result.itemsDeduplicated).toBe(0);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should throw when batch is in ABORTED status",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "ABORTED",
|
|
sealed: false,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
});
|
|
|
|
await expect(
|
|
service.call(authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), {
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
})
|
|
).rejects.toThrow(ServiceValidationError);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should throw when batch is sealed but ABORTED (callback aborted post-seal must surface as error)",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// V2 batch completion callback sets status=ABORTED (failedRunCount > 0 &&
|
|
// successfulRunCount === 0) without touching sealed=true that the seal
|
|
// step previously set. The Phase 2 retry must NOT mask this terminal
|
|
// failure as success — every run failed.
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "ABORTED",
|
|
sealed: true,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
});
|
|
|
|
await expect(
|
|
service.call(authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), {
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
})
|
|
).rejects.toThrow(ServiceValidationError);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=true when batch is PARTIAL_FAILED (Phase 2 retry idempotency)",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// V2 completion callback sets PARTIAL_FAILED when some run-creation
|
|
// attempts failed but at least one succeeded. The Phase 2 stream itself
|
|
// did its job (items were enqueued and processed), so a retry should
|
|
// see this as terminal success — the per-item failures are visible on
|
|
// the individual run records.
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "PARTIAL_FAILED",
|
|
sealed: false,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
});
|
|
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
expect(result.itemsAccepted).toBe(0);
|
|
expect(result.itemsDeduplicated).toBe(0);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=true when batch is PARTIAL_FAILED by callback before seal attempt",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: authenticatedEnvironment.id,
|
|
environmentType: authenticatedEnvironment.type,
|
|
organizationId: authenticatedEnvironment.organizationId,
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runCount: 2,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item1" }),
|
|
payloadType: "application/json",
|
|
});
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item2" }),
|
|
payloadType: "application/json",
|
|
});
|
|
|
|
// Simulate the race where V2's batchCompletionCallback runs between
|
|
// getEnqueuedCount and the seal updateMany — some runs failed to create
|
|
// but at least one succeeded, so callback sets status=PARTIAL_FAILED
|
|
// without setting sealed=true.
|
|
const racingPrisma = {
|
|
...prisma,
|
|
batchTaskRun: {
|
|
...prisma.batchTaskRun,
|
|
findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun),
|
|
updateMany: async () => {
|
|
await prisma.batchTaskRun.update({
|
|
where: { id: batch.id },
|
|
data: {
|
|
status: "PARTIAL_FAILED",
|
|
},
|
|
});
|
|
return { count: 0 };
|
|
},
|
|
findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun),
|
|
},
|
|
} as unknown as PrismaClient;
|
|
|
|
// The batch find + seal updateMany now route through the engine's run-store
|
|
// (route-by-batch-id under the run-ops split), so the racing client must back the
|
|
// store the service reads through.
|
|
engine.runStore = new PostgresRunStore({
|
|
prisma: racingPrisma,
|
|
readOnlyPrisma: racingPrisma,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma: racingPrisma,
|
|
engine,
|
|
});
|
|
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
|
|
const updatedBatch = await prisma.batchTaskRun.findUnique({
|
|
where: { id: batch.id },
|
|
});
|
|
|
|
expect(updatedBatch?.status).toBe("PARTIAL_FAILED");
|
|
expect(updatedBatch?.sealed).toBe(false);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=true when concurrent request already sealed the batch during seal attempt",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// Create a batch in PENDING state
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
// Initialize the batch in Redis
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: authenticatedEnvironment.id,
|
|
environmentType: authenticatedEnvironment.type,
|
|
organizationId: authenticatedEnvironment.organizationId,
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runCount: 2,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
// Enqueue items
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item1" }),
|
|
payloadType: "application/json",
|
|
});
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item2" }),
|
|
payloadType: "application/json",
|
|
});
|
|
|
|
// Create a custom prisma client that simulates a race condition:
|
|
// When updateMany is called on batchTaskRun, it returns count: 0 (as if another request beat us)
|
|
// but the subsequent findUnique shows the batch is sealed and PROCESSING
|
|
const racingPrisma = {
|
|
...prisma,
|
|
batchTaskRun: {
|
|
...prisma.batchTaskRun,
|
|
findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun),
|
|
updateMany: async () => {
|
|
// Simulate another request winning the race - seal the batch first
|
|
await prisma.batchTaskRun.update({
|
|
where: { id: batch.id },
|
|
data: {
|
|
sealed: true,
|
|
sealedAt: new Date(),
|
|
status: "PROCESSING",
|
|
processingStartedAt: new Date(),
|
|
},
|
|
});
|
|
// Return 0 as if the conditional update failed
|
|
return { count: 0 };
|
|
},
|
|
findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun),
|
|
},
|
|
} as unknown as PrismaClient;
|
|
|
|
// The batch find + seal updateMany now route through the engine's run-store
|
|
// (route-by-batch-id under the run-ops split), so the racing client must back the
|
|
// store the service reads through.
|
|
engine.runStore = new PostgresRunStore({
|
|
prisma: racingPrisma,
|
|
readOnlyPrisma: racingPrisma,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma: racingPrisma,
|
|
engine,
|
|
});
|
|
|
|
// Call the service - it should detect the race and return success since batch is sealed
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
// Should return sealed=true because the batch was sealed (by the "other" request)
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
|
|
// Verify the batch is sealed in the database
|
|
const updatedBatch = await prisma.batchTaskRun.findUnique({
|
|
where: { id: batch.id },
|
|
});
|
|
|
|
expect(updatedBatch?.sealed).toBe(true);
|
|
expect(updatedBatch?.status).toBe("PROCESSING");
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=true when batch is COMPLETED by BatchQueue before seal attempt",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// Create a batch in PENDING state
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
// Initialize the batch in Redis
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: authenticatedEnvironment.id,
|
|
environmentType: authenticatedEnvironment.type,
|
|
organizationId: authenticatedEnvironment.organizationId,
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runCount: 2,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
// Enqueue items - the enqueued count check passes but the seal updateMany
|
|
// will race with tryCompleteBatch moving status to COMPLETED.
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item1" }),
|
|
payloadType: "application/json",
|
|
});
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item2" }),
|
|
payloadType: "application/json",
|
|
});
|
|
|
|
// Simulate the race where BatchQueue's completionCallback runs
|
|
// tryCompleteBatch between getEnqueuedCount and the seal updateMany.
|
|
// tryCompleteBatch sets status=COMPLETED but NOT sealed=true.
|
|
const racingPrisma = {
|
|
...prisma,
|
|
batchTaskRun: {
|
|
...prisma.batchTaskRun,
|
|
findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun),
|
|
updateMany: async () => {
|
|
await prisma.batchTaskRun.update({
|
|
where: { id: batch.id },
|
|
data: {
|
|
status: "COMPLETED",
|
|
},
|
|
});
|
|
// The conditional updateMany(where: status="PENDING") would now fail
|
|
return { count: 0 };
|
|
},
|
|
findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun),
|
|
},
|
|
} as unknown as PrismaClient;
|
|
|
|
// The batch find + seal updateMany now route through the engine's run-store
|
|
// (route-by-batch-id under the run-ops split), so the racing client must back the
|
|
// store the service reads through.
|
|
engine.runStore = new PostgresRunStore({
|
|
prisma: racingPrisma,
|
|
readOnlyPrisma: racingPrisma,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma: racingPrisma,
|
|
engine,
|
|
});
|
|
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
// The endpoint should accept the COMPLETED state as a success case so the
|
|
// SDK does not retry a batch whose child runs have already finished.
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
|
|
const updatedBatch = await prisma.batchTaskRun.findUnique({
|
|
where: { id: batch.id },
|
|
});
|
|
|
|
expect(updatedBatch?.status).toBe("COMPLETED");
|
|
// sealed stays false because the BatchQueue completion path does not set
|
|
// it - that's fine, the batch is terminal.
|
|
expect(updatedBatch?.sealed).toBe(false);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should throw error when race condition leaves batch in unexpected state",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// Create a batch in PENDING state
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
// Initialize the batch in Redis
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: authenticatedEnvironment.id,
|
|
environmentType: authenticatedEnvironment.type,
|
|
organizationId: authenticatedEnvironment.organizationId,
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runCount: 2,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
// Enqueue items
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item1" }),
|
|
payloadType: "application/json",
|
|
});
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item2" }),
|
|
payloadType: "application/json",
|
|
});
|
|
|
|
// Create a custom prisma client that simulates a race condition where
|
|
// the batch ends up in an unexpected state (ABORTED instead of PROCESSING)
|
|
const racingPrisma = {
|
|
...prisma,
|
|
batchTaskRun: {
|
|
...prisma.batchTaskRun,
|
|
findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun),
|
|
updateMany: async () => {
|
|
// Simulate the batch being aborted by another process
|
|
await prisma.batchTaskRun.update({
|
|
where: { id: batch.id },
|
|
data: {
|
|
sealed: true,
|
|
status: "ABORTED",
|
|
},
|
|
});
|
|
// Return 0 as if the conditional update failed
|
|
return { count: 0 };
|
|
},
|
|
findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun),
|
|
},
|
|
} as unknown as PrismaClient;
|
|
|
|
// The batch find + seal updateMany now route through the engine's run-store
|
|
// (route-by-batch-id under the run-ops split), so the racing client must back the
|
|
// store the service reads through.
|
|
engine.runStore = new PostgresRunStore({
|
|
prisma: racingPrisma,
|
|
readOnlyPrisma: racingPrisma,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma: racingPrisma,
|
|
engine,
|
|
});
|
|
|
|
// Call the service - it should throw because the batch is in an unexpected state
|
|
await expect(
|
|
service.call(authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), {
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
})
|
|
).rejects.toThrow(/unexpected state/);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=false when item count does not match",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// Create a batch expecting 3 items
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 3,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
// Initialize the batch in Redis
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: authenticatedEnvironment.id,
|
|
environmentType: authenticatedEnvironment.type,
|
|
organizationId: authenticatedEnvironment.organizationId,
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runCount: 3,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
// Only enqueue 2 items (1 short of expected)
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item1" }),
|
|
payloadType: "application/json",
|
|
});
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item2" }),
|
|
payloadType: "application/json",
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
});
|
|
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
// Should return sealed=false because item count doesn't match
|
|
expect(result.sealed).toBe(false);
|
|
expect(result.enqueuedCount).toBe(2);
|
|
expect(result.expectedCount).toBe(3);
|
|
|
|
// Verify the batch is NOT sealed in the database
|
|
const updatedBatch = await prisma.batchTaskRun.findUnique({
|
|
where: { id: batch.id },
|
|
});
|
|
|
|
expect(updatedBatch?.sealed).toBe(false);
|
|
expect(updatedBatch?.status).toBe("PENDING");
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=true when seal-failed race produces sealed=true + PENDING (post-callback all-created)",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 2,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: authenticatedEnvironment.id,
|
|
environmentType: authenticatedEnvironment.type,
|
|
organizationId: authenticatedEnvironment.organizationId,
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runCount: 2,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item1" }),
|
|
payloadType: "application/json",
|
|
});
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item2" }),
|
|
payloadType: "application/json",
|
|
});
|
|
|
|
// Simulate the race where a concurrent path seals the batch (sealed=true,
|
|
// PROCESSING), then the V2 batchCompletionCallback fires with all runs
|
|
// created successfully and resets status to PENDING (sealed stays true).
|
|
// Our seal updateMany then fails the conditional (sealed=false no longer
|
|
// matches), and the re-query sees sealed=true + PENDING — a perfectly
|
|
// valid post-callback state that the SDK retry should treat as success.
|
|
const racingPrisma = {
|
|
...prisma,
|
|
batchTaskRun: {
|
|
...prisma.batchTaskRun,
|
|
findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun),
|
|
updateMany: async () => {
|
|
await prisma.batchTaskRun.update({
|
|
where: { id: batch.id },
|
|
data: {
|
|
sealed: true,
|
|
sealedAt: new Date(),
|
|
// Intentionally leave status as PENDING — that's exactly what
|
|
// the V2 batchCompletionCallback does after all runs are
|
|
// created (status PROCESSING → PENDING).
|
|
},
|
|
});
|
|
return { count: 0 };
|
|
},
|
|
findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun),
|
|
},
|
|
} as unknown as PrismaClient;
|
|
|
|
// The batch find + seal updateMany now route through the engine's run-store
|
|
// (route-by-batch-id under the run-ops split), so the racing client must back the
|
|
// store the service reads through.
|
|
engine.runStore = new PostgresRunStore({
|
|
prisma: racingPrisma,
|
|
readOnlyPrisma: racingPrisma,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma: racingPrisma,
|
|
engine,
|
|
});
|
|
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
|
|
const updatedBatch = await prisma.batchTaskRun.findUnique({
|
|
where: { id: batch.id },
|
|
});
|
|
|
|
expect(updatedBatch?.sealed).toBe(true);
|
|
expect(updatedBatch?.status).toBe("PENDING");
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should throw when count-mismatch race produces sealed=true + ABORTED (no TaskRuns created)",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 3,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: authenticatedEnvironment.id,
|
|
environmentType: authenticatedEnvironment.type,
|
|
organizationId: authenticatedEnvironment.organizationId,
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runCount: 3,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
// Only enqueue 2 items so the post-loop count check trips into the
|
|
// mismatch handler. The race we're simulating: between our pre-loop
|
|
// findFirst and the count-mismatch re-query, a concurrent path sealed
|
|
// the batch, runs were attempted, every run-creation failed AND the
|
|
// pre-failed-TaskRun fallback also failed → callback sets ABORTED.
|
|
// The customer has zero TaskRun records to monitor, so the retry must
|
|
// throw rather than silently succeed.
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item1" }),
|
|
payloadType: "application/json",
|
|
});
|
|
await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ data: "item2" }),
|
|
payloadType: "application/json",
|
|
});
|
|
|
|
// Override findFirst to flip the batch to sealed=true + ABORTED on the
|
|
// re-query that happens INSIDE the count-mismatch branch. The first
|
|
// findFirst (pre-loop) must still see PENDING + sealed=false so we
|
|
// pass through and reach the count-mismatch branch.
|
|
let findFirstCallCount = 0;
|
|
const racingPrisma = {
|
|
...prisma,
|
|
batchTaskRun: {
|
|
...prisma.batchTaskRun,
|
|
findFirst: async (args: Parameters<typeof prisma.batchTaskRun.findFirst>[0]) => {
|
|
findFirstCallCount++;
|
|
if (findFirstCallCount >= 2) {
|
|
await prisma.batchTaskRun.update({
|
|
where: { id: batch.id },
|
|
data: {
|
|
sealed: true,
|
|
sealedAt: new Date(),
|
|
status: "ABORTED",
|
|
completedAt: new Date(),
|
|
},
|
|
});
|
|
}
|
|
return prisma.batchTaskRun.findFirst.call(prisma.batchTaskRun, args);
|
|
},
|
|
updateMany: prisma.batchTaskRun.updateMany.bind(prisma.batchTaskRun),
|
|
findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun),
|
|
},
|
|
} as unknown as PrismaClient;
|
|
|
|
// The batch find + seal updateMany now route through the engine's run-store
|
|
// (route-by-batch-id under the run-ops split), so the racing client must back the
|
|
// store the service reads through.
|
|
engine.runStore = new PostgresRunStore({
|
|
prisma: racingPrisma,
|
|
readOnlyPrisma: racingPrisma,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma: racingPrisma,
|
|
engine,
|
|
});
|
|
|
|
await expect(
|
|
service.call(authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), {
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
})
|
|
).rejects.toThrow(ServiceValidationError);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=true when batch is sealed=false + PENDING + processingCompletedAt set (pre-loop post-callback)",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// The V2 batchCompletionCallback set processingCompletedAt without
|
|
// touching sealed (sealed gets set by streamBatchItems, not the callback).
|
|
// Status stays PENDING because every run was created successfully (the
|
|
// callback's "all created, waiting for completion" branch). A Phase 2
|
|
// retry arriving in this state must treat it as success — every item
|
|
// has a TaskRun record for the customer to monitor.
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 4,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
processingCompletedAt: new Date(),
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
});
|
|
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
expect(result.itemsAccepted).toBe(0);
|
|
expect(result.itemsDeduplicated).toBe(0);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should return sealed=true on count-mismatch when callback fired before our getEnqueuedCount (cleanup race — customer scenario)",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = new RunEngine({
|
|
prisma,
|
|
worker: {
|
|
redis: redisOptions,
|
|
workers: 1,
|
|
tasksPerWorker: 10,
|
|
pollIntervalMs: 100,
|
|
disabled: true,
|
|
},
|
|
queue: {
|
|
redis: redisOptions,
|
|
},
|
|
runLock: {
|
|
redis: redisOptions,
|
|
},
|
|
machines: {
|
|
defaultMachine: "small-1x",
|
|
machines: {
|
|
"small-1x": {
|
|
name: "small-1x" as const,
|
|
cpu: 0.5,
|
|
memory: 0.5,
|
|
centsPerMs: 0.0001,
|
|
},
|
|
},
|
|
baseCostInCents: 0.0005,
|
|
},
|
|
batchQueue: {
|
|
redis: redisOptions,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
// The customer's exact case: 4-item batch, BatchQueue rushes through
|
|
// all items before our service finishes its loop, callback fires
|
|
// (setting processingCompletedAt; status stays PENDING because all
|
|
// runs created cleanly), cleanup deletes Redis state. Our service
|
|
// hits the count-mismatch branch because getBatchEnqueuedCount returns
|
|
// 0 (cleaned). The re-query sees sealed=false + PENDING but
|
|
// processingCompletedAt is set — the work is done.
|
|
const batch = await createBatch(prisma, authenticatedEnvironment.id, {
|
|
runCount: 4,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: authenticatedEnvironment.id,
|
|
environmentType: authenticatedEnvironment.type,
|
|
organizationId: authenticatedEnvironment.organizationId,
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runCount: 4,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
// The pre-loop validate-find must see "initial" state (so it passes through to the
|
|
// loop), and the count-mismatch re-query must see "post-callback" state. Use a findFirst
|
|
// counter to flip the DB between those two reads, matching production timing where the
|
|
// callback fires while our loop runs. Both reads route through the engine's run-store
|
|
// (route-by-batch-id under the split), so the racing client backs the store.
|
|
let findFirstCallCount = 0;
|
|
const racingPrisma = {
|
|
...prisma,
|
|
batchTaskRun: {
|
|
...prisma.batchTaskRun,
|
|
findFirst: async (args: Parameters<typeof prisma.batchTaskRun.findFirst>[0]) => {
|
|
findFirstCallCount++;
|
|
if (findFirstCallCount === 2) {
|
|
await prisma.batchTaskRun.update({
|
|
where: { id: batch.id },
|
|
data: {
|
|
processingCompletedAt: new Date(),
|
|
successfulRunCount: 4,
|
|
},
|
|
});
|
|
}
|
|
return prisma.batchTaskRun.findFirst.call(prisma.batchTaskRun, args);
|
|
},
|
|
updateMany: prisma.batchTaskRun.updateMany.bind(prisma.batchTaskRun),
|
|
findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun),
|
|
},
|
|
} as unknown as PrismaClient;
|
|
|
|
engine.runStore = new PostgresRunStore({
|
|
prisma: racingPrisma,
|
|
readOnlyPrisma: racingPrisma,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma: racingPrisma,
|
|
engine,
|
|
});
|
|
|
|
const result = await service.call(
|
|
authenticatedEnvironment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable([]),
|
|
{
|
|
maxItemBytes: 1024 * 1024,
|
|
concurrency: 1,
|
|
}
|
|
);
|
|
|
|
// The retry must be treated as success — every item's TaskRun was
|
|
// created by the original Phase 2 call. Returning sealed:false here
|
|
// (the previous behavior) made the SDK retry the stream against a
|
|
// cleaned-up batch, which then 5xx'd, exhausted SDK retries, and
|
|
// surfaced as BatchTriggerError despite all runs succeeding.
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.id).toBe(batch.friendlyId);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"ingests a large batch with bounded concurrency and seals it",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = buildEngine(prisma, redisOptions);
|
|
const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const runCount = 100;
|
|
const batch = await createBatch(prisma, environment.id, {
|
|
runCount,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: environment.id,
|
|
environmentType: environment.type,
|
|
organizationId: environment.organizationId,
|
|
projectId: environment.projectId,
|
|
runCount,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({ prisma, engine });
|
|
|
|
const result = await service.call(
|
|
environment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable(makeItems(runCount)),
|
|
{ maxItemBytes: 1024 * 1024, concurrency: 10 }
|
|
);
|
|
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.itemsAccepted).toBe(runCount);
|
|
expect(result.itemsDeduplicated).toBe(0);
|
|
|
|
// Every item was enqueued exactly once despite concurrent ingestion.
|
|
const enqueuedCount = await engine.getBatchEnqueuedCount(batch.id);
|
|
expect(enqueuedCount).toBe(runCount);
|
|
|
|
const updatedBatch = await prisma.batchTaskRun.findUnique({ where: { id: batch.id } });
|
|
expect(updatedBatch?.sealed).toBe(true);
|
|
expect(updatedBatch?.status).toBe("PROCESSING");
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"bounds in-flight item processing to the configured concurrency",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = buildEngine(prisma, redisOptions);
|
|
const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const runCount = 40;
|
|
const concurrency = 5;
|
|
const batch = await createBatch(prisma, environment.id, {
|
|
runCount,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: environment.id,
|
|
environmentType: environment.type,
|
|
organizationId: environment.organizationId,
|
|
projectId: environment.projectId,
|
|
runCount,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
// A real payload processor that holds each slot briefly so we can observe
|
|
// how many items are processed simultaneously. This guards against an
|
|
// accidental regression that buffers the whole stream / ignores the cap.
|
|
let inFlight = 0;
|
|
let maxInFlight = 0;
|
|
class TrackingPayloadProcessor extends BatchPayloadProcessor {
|
|
async process(payload: unknown, payloadType: string): Promise<BatchPayloadProcessResult> {
|
|
inFlight++;
|
|
maxInFlight = Math.max(maxInFlight, inFlight);
|
|
try {
|
|
await sleep(20);
|
|
return { payload, payloadType, wasOffloaded: false, size: 0 };
|
|
} finally {
|
|
inFlight--;
|
|
}
|
|
}
|
|
}
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
payloadProcessor: new TrackingPayloadProcessor(),
|
|
});
|
|
|
|
const result = await service.call(
|
|
environment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable(makeItems(runCount)),
|
|
{ maxItemBytes: 1024 * 1024, concurrency }
|
|
);
|
|
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.itemsAccepted).toBe(runCount);
|
|
// Actually ran concurrently, but never exceeded the configured cap.
|
|
expect(maxInFlight).toBeGreaterThan(1);
|
|
expect(maxInFlight).toBeLessThanOrEqual(concurrency);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"perf: higher ingest concurrency processes a batch proportionally faster",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = buildEngine(prisma, redisOptions);
|
|
const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const runCount = 150;
|
|
|
|
// A payload processor that holds each slot for a fixed duration, modelling
|
|
// the per-item object-store offload. perItemLatencyMs=0 models a small
|
|
// payload that never offloads, leaving the real per-item enqueueBatchItem
|
|
// (a Redis round-trip) as the only serialized cost. Local MinIO/Redis
|
|
// latency is sub-millisecond and too noisy to compare directly, so the
|
|
// offload case uses a fixed hold instead.
|
|
class FixedLatencyPayloadProcessor extends BatchPayloadProcessor {
|
|
constructor(private readonly perItemLatencyMs: number) {
|
|
super();
|
|
}
|
|
async process(payload: unknown, payloadType: string): Promise<BatchPayloadProcessResult> {
|
|
if (this.perItemLatencyMs > 0) {
|
|
await sleep(this.perItemLatencyMs);
|
|
}
|
|
return { payload, payloadType, wasOffloaded: this.perItemLatencyMs > 0, size: 0 };
|
|
}
|
|
}
|
|
|
|
async function timeIngest(concurrency: number, perItemLatencyMs: number): Promise<number> {
|
|
// Each run needs its own batch: sealing mutates state and a re-stream
|
|
// of the same batch would dedup every item.
|
|
const batch = await createBatch(prisma, environment.id, {
|
|
runCount,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: environment.id,
|
|
environmentType: environment.type,
|
|
organizationId: environment.organizationId,
|
|
projectId: environment.projectId,
|
|
runCount,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({
|
|
prisma,
|
|
engine,
|
|
payloadProcessor: new FixedLatencyPayloadProcessor(perItemLatencyMs),
|
|
});
|
|
|
|
const start = performance.now();
|
|
const result = await service.call(
|
|
environment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable(makeItems(runCount)),
|
|
{ maxItemBytes: 1024 * 1024, concurrency }
|
|
);
|
|
const elapsedMs = performance.now() - start;
|
|
|
|
// Correctness holds at every concurrency: all items accepted and sealed.
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.itemsAccepted).toBe(runCount);
|
|
expect(result.itemsDeduplicated).toBe(0);
|
|
expect(await engine.getBatchEnqueuedCount(batch.id)).toBe(runCount);
|
|
|
|
return elapsedMs;
|
|
}
|
|
|
|
// Scenario A: large payloads, where each item offloads to object storage.
|
|
const offloadLatencyMs = 10;
|
|
const offloadSeqMs = await timeIngest(1, offloadLatencyMs);
|
|
const offload10Ms = await timeIngest(10, offloadLatencyMs);
|
|
const offload50Ms = await timeIngest(50, offloadLatencyMs);
|
|
|
|
// Scenario B: small payloads (no offload). The only per-item cost is the
|
|
// real Redis enqueue, so this is the floor case that proves the speedup
|
|
// applies to all batch ingest, not just large-payload batches.
|
|
const enqueueSeqMs = await timeIngest(1, 0);
|
|
const enqueue10Ms = await timeIngest(10, 0);
|
|
|
|
// eslint-disable-next-line no-console
|
|
console.log(
|
|
`\n[streamBatchItems perf] runCount=${runCount}\n` +
|
|
` large payloads (${offloadLatencyMs}ms/item offload):\n` +
|
|
` concurrency=1 ${offloadSeqMs.toFixed(0)}ms\n` +
|
|
` concurrency=10 ${offload10Ms.toFixed(0)}ms (${(offloadSeqMs / offload10Ms).toFixed(
|
|
1
|
|
)}x faster)\n` +
|
|
` concurrency=50 ${offload50Ms.toFixed(0)}ms (${(offloadSeqMs / offload50Ms).toFixed(
|
|
1
|
|
)}x faster)\n` +
|
|
` small payloads (Redis enqueue only, no offload):\n` +
|
|
` concurrency=1 ${enqueueSeqMs.toFixed(0)}ms\n` +
|
|
` concurrency=10 ${enqueue10Ms.toFixed(0)}ms (${(enqueueSeqMs / enqueue10Ms).toFixed(
|
|
1
|
|
)}x faster)\n`
|
|
);
|
|
|
|
// Sequential floor: N items each held for offloadLatencyMs cannot finish
|
|
// faster than N x latency. Parallel ingest must beat that floor decisively.
|
|
const sequentialFloorMs = runCount * offloadLatencyMs;
|
|
expect(offloadSeqMs).toBeGreaterThan(sequentialFloorMs * 0.8);
|
|
|
|
// 10x concurrency on a latency-bound workload should be well over 2x faster.
|
|
expect(offload10Ms).toBeLessThan(offloadSeqMs / 2);
|
|
// More concurrency keeps helping (or at least never regresses).
|
|
expect(offload50Ms).toBeLessThanOrEqual(offload10Ms * 1.2);
|
|
|
|
// Even with no offload, overlapping the per-item Redis enqueue is strictly
|
|
// faster than doing them one at a time.
|
|
expect(enqueue10Ms).toBeLessThan(enqueueSeqMs);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"deduplicates already-enqueued items during concurrent ingest (Phase 2 retry)",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = buildEngine(prisma, redisOptions);
|
|
const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const runCount = 20;
|
|
const batch = await createBatch(prisma, environment.id, {
|
|
runCount,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: environment.id,
|
|
environmentType: environment.type,
|
|
organizationId: environment.organizationId,
|
|
projectId: environment.projectId,
|
|
runCount,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
// Pre-enqueue the first half (simulating a partially-completed prior stream
|
|
// that failed before sealing). The retry re-streams everything.
|
|
for (let index = 0; index < 10; index++) {
|
|
await engine.enqueueBatchItem(batch.id, environment.id, index, {
|
|
task: "test-task",
|
|
payload: JSON.stringify({ i: index }),
|
|
payloadType: "application/json",
|
|
});
|
|
}
|
|
|
|
const service = new StreamBatchItemsService({ prisma, engine });
|
|
|
|
const result = await service.call(
|
|
environment,
|
|
batch.friendlyId,
|
|
itemsToAsyncIterable(makeItems(runCount)),
|
|
{ maxItemBytes: 1024 * 1024, concurrency: 8 }
|
|
);
|
|
|
|
expect(result.sealed).toBe(true);
|
|
expect(result.itemsDeduplicated).toBe(10);
|
|
expect(result.itemsAccepted).toBe(10);
|
|
|
|
const enqueuedCount = await engine.getBatchEnqueuedCount(batch.id);
|
|
expect(enqueuedCount).toBe(runCount);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"rejects an oversized item whose index exceeds runCount",
|
|
async ({ prisma, redisOptions }) => {
|
|
const engine = buildEngine(prisma, redisOptions);
|
|
const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const batch = await createBatch(prisma, environment.id, {
|
|
runCount: 2,
|
|
status: "PENDING",
|
|
sealed: false,
|
|
});
|
|
|
|
await engine.initializeBatch({
|
|
batchId: batch.id,
|
|
friendlyId: batch.friendlyId,
|
|
environmentId: environment.id,
|
|
environmentType: environment.type,
|
|
organizationId: environment.organizationId,
|
|
projectId: environment.projectId,
|
|
runCount: 2,
|
|
processingConcurrency: 10,
|
|
});
|
|
|
|
const service = new StreamBatchItemsService({ prisma, engine });
|
|
|
|
// An oversized marker must hit the same out-of-range guard as a normal
|
|
// item rather than slipping through to create a stray pre-failed run.
|
|
async function* oversizedOutOfRange() {
|
|
yield {
|
|
__batchItemError: "OVERSIZED",
|
|
index: 5,
|
|
task: "test-task",
|
|
actualSize: 9999,
|
|
maxSize: 1000,
|
|
};
|
|
}
|
|
|
|
await expect(
|
|
service.call(environment, batch.friendlyId, oversizedOutOfRange(), {
|
|
maxItemBytes: 1024,
|
|
concurrency: 4,
|
|
})
|
|
).rejects.toThrow(ServiceValidationError);
|
|
|
|
await engine.quit();
|
|
}
|
|
);
|
|
});
|
|
|
|
describe("createNdjsonParserStream", () => {
|
|
async function collectStream<T>(stream: ReadableStream<T>): Promise<T[]> {
|
|
const results: T[] = [];
|
|
for await (const item of streamToAsyncIterable(stream)) {
|
|
results.push(item);
|
|
}
|
|
return results;
|
|
}
|
|
|
|
function chunksToStream(chunks: Uint8Array[]): ReadableStream<Uint8Array> {
|
|
let index = 0;
|
|
return new ReadableStream({
|
|
pull(controller) {
|
|
if (index < chunks.length) {
|
|
controller.enqueue(chunks[index++]);
|
|
} else {
|
|
controller.close();
|
|
}
|
|
},
|
|
});
|
|
}
|
|
|
|
it("should parse basic NDJSON correctly", async () => {
|
|
const ndjson = '{"name":"alice"}\n{"name":"bob"}\n{"name":"charlie"}\n';
|
|
const encoder = new TextEncoder();
|
|
const stream = chunksToStream([encoder.encode(ndjson)]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ name: "alice" }, { name: "bob" }, { name: "charlie" }]);
|
|
});
|
|
|
|
it("should handle NDJSON without trailing newline", async () => {
|
|
const ndjson = '{"id":1}\n{"id":2}';
|
|
const encoder = new TextEncoder();
|
|
const stream = chunksToStream([encoder.encode(ndjson)]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ id: 1 }, { id: 2 }]);
|
|
});
|
|
|
|
it("should handle escaped newlines in JSON string values", async () => {
|
|
// JSON.stringify escapes newlines as \n (two chars: backslash + n),
|
|
// so they don't break NDJSON line boundaries. This is the normal case
|
|
// when the SDK serializes payloads containing newlines.
|
|
const item1 = JSON.stringify({ payload: "line1\nline2\nline3" });
|
|
const item2 = JSON.stringify({ payload: "no newlines" });
|
|
const ndjson = item1 + "\n" + item2 + "\n";
|
|
const encoder = new TextEncoder();
|
|
const stream = chunksToStream([encoder.encode(ndjson)]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ payload: "line1\nline2\nline3" }, { payload: "no newlines" }]);
|
|
});
|
|
|
|
it("should skip empty lines", async () => {
|
|
const ndjson = '{"a":1}\n\n{"b":2}\n \n{"c":3}\n';
|
|
const encoder = new TextEncoder();
|
|
const stream = chunksToStream([encoder.encode(ndjson)]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ a: 1 }, { b: 2 }, { c: 3 }]);
|
|
});
|
|
|
|
it("should handle chunks split mid-line", async () => {
|
|
// Split '{"split":"value"}\n' across multiple chunks
|
|
const encoder = new TextEncoder();
|
|
const chunks = [encoder.encode('{"spl'), encoder.encode('it":"va'), encoder.encode('lue"}\n')];
|
|
const stream = chunksToStream(chunks);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ split: "value" }]);
|
|
});
|
|
|
|
it("should handle multibyte UTF-8 characters split across chunks", async () => {
|
|
// Test with emoji and other multibyte characters
|
|
// The emoji "😀" is 4 bytes: 0xF0 0x9F 0x98 0x80
|
|
const json = '{"emoji":"😀"}\n';
|
|
const fullBytes = new TextEncoder().encode(json);
|
|
|
|
// Split in the middle of the emoji (between byte 2 and 3 of the 4-byte sequence)
|
|
// Find where the emoji starts
|
|
const emojiStart = fullBytes.indexOf(0xf0);
|
|
expect(emojiStart).toBeGreaterThan(0);
|
|
|
|
// Split after first 2 bytes of the emoji
|
|
const chunk1 = fullBytes.slice(0, emojiStart + 2);
|
|
const chunk2 = fullBytes.slice(emojiStart + 2);
|
|
|
|
const stream = chunksToStream([chunk1, chunk2]);
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ emoji: "😀" }]);
|
|
});
|
|
|
|
it("should handle multiple multibyte characters across chunks", async () => {
|
|
// Japanese text: "こんにちは" (each hiragana is 3 bytes in UTF-8)
|
|
const json = '{"greeting":"こんにちは"}\n';
|
|
const fullBytes = new TextEncoder().encode(json);
|
|
|
|
// Split into many small chunks to stress test UTF-8 handling
|
|
const chunkSize = 3; // Deliberately misaligned with UTF-8 boundaries
|
|
const chunks: Uint8Array[] = [];
|
|
for (let i = 0; i < fullBytes.length; i += chunkSize) {
|
|
chunks.push(fullBytes.slice(i, i + chunkSize));
|
|
}
|
|
|
|
const stream = chunksToStream(chunks);
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ greeting: "こんにちは" }]);
|
|
});
|
|
|
|
it("should emit OversizedItemMarker for lines exceeding maxItemBytes", async () => {
|
|
const maxBytes = 50;
|
|
// Create a line that exceeds the limit with index and task fields
|
|
const largeJson = JSON.stringify({ index: 3, task: "my-task", data: "x".repeat(100) }) + "\n";
|
|
const encoder = new TextEncoder();
|
|
const stream = chunksToStream([encoder.encode(largeJson)]);
|
|
|
|
const parser = createNdjsonParserStream(maxBytes);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toHaveLength(1);
|
|
const marker = results[0] as OversizedItemMarker;
|
|
expect(marker.__batchItemError).toBe("OVERSIZED");
|
|
expect(marker.index).toBe(3);
|
|
expect(marker.task).toBe("my-task");
|
|
expect(marker.maxSize).toBe(maxBytes);
|
|
expect(marker.actualSize).toBeGreaterThan(maxBytes);
|
|
});
|
|
|
|
it("should emit OversizedItemMarker for unbounded accumulation without newlines", async () => {
|
|
const maxBytes = 50;
|
|
// Send data without any newlines that exceeds the buffer limit
|
|
const encoder = new TextEncoder();
|
|
const chunks = [
|
|
encoder.encode('{"index":7,"task":"big-task","start":"'),
|
|
encoder.encode("x".repeat(60)), // This will push buffer over 50 bytes
|
|
];
|
|
const stream = chunksToStream(chunks);
|
|
|
|
const parser = createNdjsonParserStream(maxBytes);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toHaveLength(1);
|
|
const marker = results[0] as OversizedItemMarker;
|
|
expect(marker.__batchItemError).toBe("OVERSIZED");
|
|
expect(marker.index).toBe(7);
|
|
expect(marker.task).toBe("big-task");
|
|
expect(marker.maxSize).toBe(maxBytes);
|
|
});
|
|
|
|
it("should skip remaining bytes of oversized line arriving in subsequent chunks", async () => {
|
|
const maxBytes = 50;
|
|
const encoder = new TextEncoder();
|
|
// Simulate a normal item, then an oversized item split across many chunks,
|
|
// then another normal item after the newline.
|
|
// The oversized line is: {"index":1,"task":"t","data":"xxxx...120 x's...xxxx"}\n
|
|
const normalItem1 = '{"index":0,"task":"t","x":1}\n';
|
|
const oversizedStart = '{"index":1,"task":"t","data":"';
|
|
const oversizedMiddle = "x".repeat(120); // way over 50 bytes
|
|
const oversizedEnd = '"}\n';
|
|
const normalItem2 = '{"index":2,"task":"t","x":2}\n';
|
|
|
|
// Send as separate chunks to trigger Case 2 (no newline, buffer > limit)
|
|
const chunks = [
|
|
encoder.encode(normalItem1 + oversizedStart),
|
|
encoder.encode(oversizedMiddle.slice(0, 60)),
|
|
encoder.encode(oversizedMiddle.slice(60)),
|
|
encoder.encode(oversizedEnd + normalItem2),
|
|
];
|
|
const stream = chunksToStream(chunks);
|
|
|
|
const parser = createNdjsonParserStream(maxBytes);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
// Should get: normal item 1, oversized marker, normal item 2
|
|
expect(results).toHaveLength(3);
|
|
expect(results[0]).toEqual({ index: 0, task: "t", x: 1 });
|
|
expect((results[1] as OversizedItemMarker).__batchItemError).toBe("OVERSIZED");
|
|
expect((results[1] as OversizedItemMarker).index).toBe(1);
|
|
expect(results[2]).toEqual({ index: 2, task: "t", x: 2 });
|
|
});
|
|
|
|
it("stamps an oversized marker's index by emit position when it can't be extracted", async () => {
|
|
// Concurrency-safety: when the `index` field can't be recovered from the
|
|
// (truncated) raw bytes, the parser must back-fill it from the emit
|
|
// position rather than leaving it for the consumer to infer sequentially.
|
|
const maxBytes = 50;
|
|
const encoder = new TextEncoder();
|
|
// item 0: small + valid
|
|
const item0 = '{"index":0,"task":"t","x":1}\n';
|
|
// item 1: oversized AND missing an `index` field → extractIndexAndTask returns -1
|
|
const item1 = `{"task":"big","data":"${"x".repeat(100)}"}\n`;
|
|
// item 2: small + valid, must keep its own index
|
|
const item2 = '{"index":2,"task":"t","x":2}\n';
|
|
const stream = chunksToStream([encoder.encode(item0 + item1 + item2)]);
|
|
|
|
const parser = createNdjsonParserStream(maxBytes);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toHaveLength(3);
|
|
expect(results[0]).toEqual({ index: 0, task: "t", x: 1 });
|
|
const marker = results[1] as OversizedItemMarker;
|
|
expect(marker.__batchItemError).toBe("OVERSIZED");
|
|
expect(marker.index).toBe(1); // emit-position fallback, not -1
|
|
expect(marker.task).toBe("big");
|
|
expect(results[2]).toEqual({ index: 2, task: "t", x: 2 });
|
|
});
|
|
|
|
it("should check byte size before decoding to prevent OOM", async () => {
|
|
// This test verifies that size is checked on raw bytes, not decoded string length
|
|
// Unicode characters like "🎉" are 4 bytes but 2 UTF-16 code units (string length 2)
|
|
const maxBytes = 30;
|
|
|
|
// Create a line with emojis - each emoji is 4 bytes
|
|
// {"e":"🎉🎉🎉🎉🎉"} = 5 + 20 (5 emojis * 4 bytes) + 2 = 27 bytes - should pass
|
|
const smallJson = '{"e":"🎉🎉🎉🎉🎉"}\n';
|
|
const smallBytes = new TextEncoder().encode(smallJson);
|
|
expect(smallBytes.length).toBeLessThan(maxBytes);
|
|
|
|
// {"e":"🎉🎉🎉🎉🎉🎉🎉"} = 7 emojis * 4 bytes + overhead = 35 bytes - should fail
|
|
const largeJson = '{"e":"🎉🎉🎉🎉🎉🎉🎉"}\n';
|
|
const largeBytes = new TextEncoder().encode(largeJson);
|
|
expect(largeBytes.length).toBeGreaterThan(maxBytes);
|
|
|
|
// Small one should succeed
|
|
const stream1 = chunksToStream([smallBytes]);
|
|
const parser1 = createNdjsonParserStream(maxBytes);
|
|
const results1 = await collectStream(stream1.pipeThrough(parser1));
|
|
expect(results1).toHaveLength(1);
|
|
|
|
// Large one should emit an OversizedItemMarker
|
|
const stream2 = chunksToStream([largeBytes]);
|
|
const parser2 = createNdjsonParserStream(maxBytes);
|
|
const results2 = await collectStream(stream2.pipeThrough(parser2));
|
|
expect(results2).toHaveLength(1);
|
|
expect((results2[0] as OversizedItemMarker).__batchItemError).toBe("OVERSIZED");
|
|
});
|
|
|
|
it("should handle final line in flush without trailing newline", async () => {
|
|
const ndjson = '{"first":1}\n{"second":2}'; // No trailing newline
|
|
const encoder = new TextEncoder();
|
|
const stream = chunksToStream([encoder.encode(ndjson)]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ first: 1 }, { second: 2 }]);
|
|
});
|
|
|
|
it("should reject invalid JSON", async () => {
|
|
const ndjson = '{"valid":true}\n{invalid json}\n';
|
|
const encoder = new TextEncoder();
|
|
const stream = chunksToStream([encoder.encode(ndjson)]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
|
|
await expect(collectStream(stream.pipeThrough(parser))).rejects.toThrow(
|
|
/Invalid JSON at line 2/
|
|
);
|
|
});
|
|
|
|
it("should reject invalid UTF-8 sequences", async () => {
|
|
// Invalid UTF-8: 0xFF is never valid in UTF-8
|
|
const invalidBytes = new Uint8Array([
|
|
0x7b,
|
|
0x22,
|
|
0x78,
|
|
0x22,
|
|
0x3a,
|
|
0xff,
|
|
0x7d,
|
|
0x0a, // {"x":�}\n with invalid byte
|
|
]);
|
|
const stream = chunksToStream([invalidBytes]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
|
|
await expect(collectStream(stream.pipeThrough(parser))).rejects.toThrow(/Invalid UTF-8/);
|
|
});
|
|
|
|
it("should handle many small chunks efficiently", async () => {
|
|
// Simulate streaming byte-by-byte
|
|
const json = '{"test":"value"}\n';
|
|
const bytes = new TextEncoder().encode(json);
|
|
const chunks = Array.from(bytes).map((b) => new Uint8Array([b]));
|
|
|
|
const stream = chunksToStream(chunks);
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ test: "value" }]);
|
|
});
|
|
|
|
it("should handle multiple lines per chunk", async () => {
|
|
const ndjson = '{"a":1}\n{"b":2}\n{"c":3}\n{"d":4}\n{"e":5}\n';
|
|
const encoder = new TextEncoder();
|
|
const stream = chunksToStream([encoder.encode(ndjson)]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([{ a: 1 }, { b: 2 }, { c: 3 }, { d: 4 }, { e: 5 }]);
|
|
});
|
|
|
|
it("should handle empty stream", async () => {
|
|
const stream = chunksToStream([]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([]);
|
|
});
|
|
|
|
it("should pass normal items and emit markers for oversized items in the same stream", async () => {
|
|
const maxBytes = 50;
|
|
const encoder = new TextEncoder();
|
|
// Normal item, then oversized item, then another normal item
|
|
const normalItem1 = '{"index":0,"task":"t","x":1}\n';
|
|
const oversizedItem = JSON.stringify({ index: 1, task: "t", data: "x".repeat(100) }) + "\n";
|
|
const normalItem2 = '{"index":2,"task":"t","x":2}\n';
|
|
const stream = chunksToStream([encoder.encode(normalItem1 + oversizedItem + normalItem2)]);
|
|
|
|
const parser = createNdjsonParserStream(maxBytes);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toHaveLength(3);
|
|
// First: normal parsed object
|
|
expect(results[0]).toEqual({ index: 0, task: "t", x: 1 });
|
|
// Second: oversized marker
|
|
expect((results[1] as OversizedItemMarker).__batchItemError).toBe("OVERSIZED");
|
|
expect((results[1] as OversizedItemMarker).index).toBe(1);
|
|
// Third: normal parsed object
|
|
expect(results[2]).toEqual({ index: 2, task: "t", x: 2 });
|
|
});
|
|
|
|
it("should handle stream with only whitespace", async () => {
|
|
const encoder = new TextEncoder();
|
|
const stream = chunksToStream([encoder.encode(" \n\n \n")]);
|
|
|
|
const parser = createNdjsonParserStream(1024);
|
|
const results = await collectStream(stream.pipeThrough(parser));
|
|
|
|
expect(results).toEqual([]);
|
|
});
|
|
});
|
|
|
|
describe("extractIndexAndTask", () => {
|
|
const encoder = new TextEncoder();
|
|
|
|
it("should extract index and task from JSON bytes", () => {
|
|
const bytes = encoder.encode('{"index":42,"task":"my-task","data":"x"}');
|
|
const result = extractIndexAndTask(bytes);
|
|
expect(result.index).toBe(42);
|
|
expect(result.task).toBe("my-task");
|
|
});
|
|
|
|
it("should return defaults for empty or malformed bytes", () => {
|
|
const result = extractIndexAndTask(new Uint8Array(0));
|
|
expect(result.index).toBe(-1);
|
|
expect(result.task).toBe("unknown");
|
|
});
|
|
|
|
it("should handle keys in any order", () => {
|
|
const bytes = encoder.encode('{"task":"other-task","data":"y","index":99}');
|
|
const result = extractIndexAndTask(bytes);
|
|
expect(result.index).toBe(99);
|
|
expect(result.task).toBe("other-task");
|
|
});
|
|
|
|
it("should not match nested keys", () => {
|
|
const bytes = encoder.encode(
|
|
'{"nested":{"index":999,"task":"inner"},"index":5,"task":"outer"}'
|
|
);
|
|
const result = extractIndexAndTask(bytes);
|
|
expect(result.index).toBe(5);
|
|
expect(result.task).toBe("outer");
|
|
});
|
|
});
|