wehub-resource-sync
838 lines
29 KiB
TypeScript
838 lines
29 KiB
TypeScript
import {
|
||
type StreamBatchItemsResponse,
|
||
BatchItemNDJSON as BatchItemNDJSONSchema,
|
||
} from "@trigger.dev/core/v3";
|
||
import { BatchId } from "@trigger.dev/core/v3/isomorphic";
|
||
import type { BatchItem, RunEngine } from "@internal/run-engine";
|
||
import pMap from "p-map";
|
||
import type { BatchTaskRunStatus } from "@trigger.dev/database";
|
||
import { prisma, type PrismaClientOrTransaction } from "~/db.server";
|
||
import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
|
||
import { logger } from "~/services/logger.server";
|
||
import { ServiceValidationError, WithRunEngine } from "../../v3/services/baseService.server";
|
||
import { BatchPayloadProcessor } from "../concerns/batchPayloads.server";
|
||
|
||
/**
|
||
* Phase 2 retry idempotency check.
|
||
*
|
||
* Returns true when the batch is in a state that means the Phase 2 stream's
|
||
* job has already been done — every item has a TaskRun record (real or
|
||
* pre-failed) for the customer to monitor. A retry, or the original call
|
||
* racing against a fast-completing BatchQueue, should return sealed:true
|
||
* in these states so the SDK stops retrying.
|
||
*
|
||
* Three "work is done" shapes:
|
||
* - status moved out of PENDING into PROCESSING/COMPLETED/PARTIAL_FAILED
|
||
* (PROCESSING via our seal, COMPLETED via tryCompleteBatch, PARTIAL_FAILED
|
||
* via the V2 batchCompletionCallback).
|
||
* - status stuck at PENDING but `sealed=true`: another concurrent
|
||
* streamBatchItems call sealed the batch and then the callback's
|
||
* happy-path branch reset status to PENDING ("all runs created").
|
||
* - status stuck at PENDING with `sealed=false` but `processingCompletedAt`
|
||
* set: the cleanup-race. BatchQueue rushed through all items, callback
|
||
* fired (setting processingCompletedAt), cleanup deleted the Redis
|
||
* metadata — all before our service got the chance to seal. The work
|
||
* is done; the discriminator is processingCompletedAt which is set
|
||
* exclusively by the V2 completion callback.
|
||
*
|
||
* ABORTED is excluded — it means ZERO TaskRun records were created (every
|
||
* per-item attempt failed AND the pre-failed-TaskRun fallback also failed,
|
||
* or queue-overload on every item). The customer has nothing to monitor
|
||
* at the run level, so the trigger call must throw to give their retry/
|
||
* error handling a chance to create a fresh batch.
|
||
*/
|
||
export function isIdempotentRetrySuccess(
|
||
status: BatchTaskRunStatus | null | undefined,
|
||
sealed: boolean | null | undefined,
|
||
processingCompletedAt: Date | null | undefined
|
||
): boolean {
|
||
return (
|
||
status === "PROCESSING" ||
|
||
status === "COMPLETED" ||
|
||
status === "PARTIAL_FAILED" ||
|
||
(status === "PENDING" && (sealed === true || processingCompletedAt != null))
|
||
);
|
||
}
|
||
|
||
export type StreamBatchItemsServiceOptions = {
|
||
maxItemBytes: number;
|
||
/** Max items processed concurrently. The route wires this to STREAMING_BATCH_INGEST_CONCURRENCY. */
|
||
concurrency: number;
|
||
};
|
||
|
||
export type OversizedItemMarker = {
|
||
__batchItemError: "OVERSIZED";
|
||
index: number;
|
||
task: string;
|
||
actualSize: number;
|
||
maxSize: number;
|
||
};
|
||
|
||
export type StreamBatchItemsServiceConstructorOptions = {
|
||
prisma?: PrismaClientOrTransaction;
|
||
engine?: RunEngine;
|
||
/** Override the payload processor (used in tests to observe ingest concurrency). */
|
||
payloadProcessor?: BatchPayloadProcessor;
|
||
};
|
||
|
||
/**
|
||
* Stream Batch Items Service (Phase 2 of 2-phase batch API).
|
||
*
|
||
* This service handles Phase 2 of the streaming batch API:
|
||
* 1. Validates batch exists and is in PENDING status
|
||
* 2. Processes NDJSON stream item by item
|
||
* 3. Calls engine.enqueueBatchItem() for each item
|
||
* 4. Tracks accepted/deduplicated counts
|
||
* 5. On completion: validates count, seals the batch
|
||
*
|
||
* The service is designed for streaming and processes items as they arrive,
|
||
* providing backpressure through the async iterator pattern.
|
||
*/
|
||
export class StreamBatchItemsService extends WithRunEngine {
|
||
private readonly payloadProcessor: BatchPayloadProcessor;
|
||
|
||
constructor(opts: StreamBatchItemsServiceConstructorOptions = {}) {
|
||
super({ prisma: opts.prisma ?? prisma, engine: opts.engine });
|
||
this.payloadProcessor = opts.payloadProcessor ?? new BatchPayloadProcessor();
|
||
}
|
||
|
||
/**
|
||
* Parse a batch friendly ID to its internal ID format.
|
||
* Throws a ServiceValidationError with 400 status if the ID is malformed.
|
||
*/
|
||
private parseBatchFriendlyId(friendlyId: string): string {
|
||
try {
|
||
return BatchId.fromFriendlyId(friendlyId);
|
||
} catch {
|
||
throw new ServiceValidationError(`Invalid batchFriendlyId: ${friendlyId}`, 400);
|
||
}
|
||
}
|
||
|
||
/**
|
||
* Process a stream of batch items from an async iterator.
|
||
* Each item is validated and enqueued to the BatchQueue.
|
||
* The batch is sealed when the stream completes.
|
||
*/
|
||
public async call(
|
||
environment: AuthenticatedEnvironment,
|
||
batchFriendlyId: string,
|
||
itemsIterator: AsyncIterable<unknown>,
|
||
options: StreamBatchItemsServiceOptions
|
||
): Promise<StreamBatchItemsResponse> {
|
||
return this.traceWithEnv<StreamBatchItemsResponse>(
|
||
"streamBatchItems()",
|
||
environment,
|
||
async (span) => {
|
||
span.setAttribute("batchId", batchFriendlyId);
|
||
|
||
// Convert friendly ID to internal ID
|
||
const batchId = this.parseBatchFriendlyId(batchFriendlyId);
|
||
|
||
// Validate batch exists and belongs to this environment. Routed by batch id so a
|
||
// run-ops id (NEW-resident) batch is found on the owning DB; the env-ownership check that
|
||
// was in the where clause is enforced app-side below.
|
||
const batch = await this._engine.runStore.findBatchTaskRunById(batchId);
|
||
|
||
if (!batch || batch.runtimeEnvironmentId !== environment.id) {
|
||
throw new ServiceValidationError(`Batch ${batchFriendlyId} not found`);
|
||
}
|
||
|
||
if (isIdempotentRetrySuccess(batch.status, batch.sealed, batch.processingCompletedAt)) {
|
||
logger.info("Batch already sealed/completed - treating Phase 2 retry as success", {
|
||
batchId: batchFriendlyId,
|
||
batchSealed: batch.sealed,
|
||
batchStatus: batch.status,
|
||
processingCompletedAt: batch.processingCompletedAt,
|
||
});
|
||
|
||
return {
|
||
id: batchFriendlyId,
|
||
itemsAccepted: 0,
|
||
itemsDeduplicated: 0,
|
||
sealed: true,
|
||
runCount: batch.runCount,
|
||
};
|
||
}
|
||
|
||
if (batch.status !== "PENDING") {
|
||
// ABORTED or any other unexpected non-PENDING state — surface as an error.
|
||
// For ABORTED specifically, throwing is required so the customer's
|
||
// batchTrigger() retries (a new batch) can recreate the runs.
|
||
throw new ServiceValidationError(
|
||
`Batch ${batchFriendlyId} is not in PENDING status (current: ${batch.status})`
|
||
);
|
||
}
|
||
|
||
// Process items from the stream with bounded concurrency.
|
||
//
|
||
// Ordering and idempotency do NOT depend on processing order:
|
||
// - The BatchQueue derives run order from each item's index
|
||
// (enqueue timestamp = batch.createdAt + itemIndex), not enqueue order.
|
||
// - enqueueBatchItem() dedups atomically per index.
|
||
// We cap concurrency to bound peak in-flight memory (≈ concurrency ×
|
||
// maxItemBytes) and to keep backpressure on the request body stream.
|
||
// p-map pulls lazily from the async iterator — at most `concurrency`
|
||
// items are read and in flight at once. stopOnError aborts ingestion on
|
||
// the first failure (the batch is left unsealed; the SDK's retry
|
||
// re-streams and dedups already-enqueued items).
|
||
const outcomes = await pMap(
|
||
itemsIterator,
|
||
(rawItem) => this.#processItem(rawItem, batchId, environment, batch.runCount),
|
||
{ concurrency: options.concurrency, stopOnError: true }
|
||
);
|
||
|
||
let itemsAccepted = 0;
|
||
let itemsDeduplicated = 0;
|
||
for (const outcome of outcomes) {
|
||
if (outcome === "accepted") {
|
||
itemsAccepted++;
|
||
} else {
|
||
itemsDeduplicated++;
|
||
}
|
||
}
|
||
|
||
// Get the actual enqueued count from Redis
|
||
const enqueuedCount = await this._engine.getBatchEnqueuedCount(batchId);
|
||
|
||
// Validate we received the expected number of items
|
||
if (enqueuedCount !== batch.runCount) {
|
||
// The batch queue consumers may have already processed all items and
|
||
// cleaned up the Redis keys before we got here. This happens when all
|
||
// runs complete fast enough that cleanup() deletes the enqueuedItemsKey
|
||
// before we read it — typically when the last item executes in the
|
||
// milliseconds between the loop ending and getBatchEnqueuedCount() being called.
|
||
// Check both sealed (sealed by this endpoint on a concurrent request) and
|
||
// COMPLETED (sealed by the BatchQueue completion path before we got here).
|
||
const currentBatch = await this._engine.runStore.findBatchTaskRunById(batchId);
|
||
|
||
if (
|
||
isIdempotentRetrySuccess(
|
||
currentBatch?.status,
|
||
currentBatch?.sealed,
|
||
currentBatch?.processingCompletedAt
|
||
)
|
||
) {
|
||
logger.info("Batch already sealed before count check (fast completion)", {
|
||
batchId: batchFriendlyId,
|
||
itemsAccepted,
|
||
itemsDeduplicated,
|
||
enqueuedCount,
|
||
expectedCount: batch.runCount,
|
||
batchStatus: currentBatch?.status,
|
||
processingCompletedAt: currentBatch?.processingCompletedAt,
|
||
});
|
||
|
||
return {
|
||
id: batchFriendlyId,
|
||
itemsAccepted,
|
||
itemsDeduplicated,
|
||
sealed: true,
|
||
runCount: batch.runCount,
|
||
};
|
||
}
|
||
|
||
if (currentBatch?.status === "ABORTED") {
|
||
// Zero TaskRuns exist — the count-mismatch sealed:false semantics
|
||
// ("retry with missing items") would mislead the SDK. Throw so the
|
||
// customer's batchTrigger() retry creates a fresh batch.
|
||
throw new ServiceValidationError(
|
||
`Batch ${batchFriendlyId} is not in PENDING status (current: ABORTED)`
|
||
);
|
||
}
|
||
|
||
logger.warn("Batch item count mismatch", {
|
||
batchId: batchFriendlyId,
|
||
expected: batch.runCount,
|
||
received: enqueuedCount,
|
||
itemsAccepted,
|
||
itemsDeduplicated,
|
||
});
|
||
|
||
// Don't seal the batch if count doesn't match
|
||
// Return sealed: false so client knows to retry with missing items
|
||
return {
|
||
id: batchFriendlyId,
|
||
itemsAccepted,
|
||
itemsDeduplicated,
|
||
sealed: false,
|
||
enqueuedCount,
|
||
expectedCount: batch.runCount,
|
||
runCount: batch.runCount,
|
||
};
|
||
}
|
||
|
||
// Seal the batch - use conditional update to prevent TOCTOU race
|
||
// Another concurrent request may have already sealed this batch
|
||
const now = new Date();
|
||
const sealResult = await this._engine.runStore.updateManyBatchTaskRun({
|
||
where: {
|
||
id: batchId,
|
||
sealed: false,
|
||
status: "PENDING",
|
||
},
|
||
data: {
|
||
sealed: true,
|
||
sealedAt: now,
|
||
status: "PROCESSING",
|
||
processingStartedAt: now,
|
||
},
|
||
});
|
||
|
||
// Check if we won the race to seal the batch
|
||
if (sealResult.count === 0) {
|
||
// The conditional update failed because the batch was no longer in
|
||
// PENDING status. Re-query to determine which path got there first:
|
||
// - A concurrent streaming request already sealed and moved it to
|
||
// PROCESSING.
|
||
// - The BatchQueue completion path finished all runs and set it to
|
||
// COMPLETED (without setting sealed=true — that's this endpoint's
|
||
// job). This window exists between completionCallback (which calls
|
||
// tryCompleteBatch) and cleanup() in BatchQueue — see
|
||
// batch-queue/index.ts.
|
||
// Either way the goal — a durable batch that the SDK stops retrying —
|
||
// has been achieved, so we return sealed: true.
|
||
const currentBatch = await this._engine.runStore.findBatchTaskRunById(batchId);
|
||
|
||
if (
|
||
isIdempotentRetrySuccess(
|
||
currentBatch?.status,
|
||
currentBatch?.sealed,
|
||
currentBatch?.processingCompletedAt
|
||
)
|
||
) {
|
||
logger.info("Batch already sealed/completed by concurrent path", {
|
||
batchId: batchFriendlyId,
|
||
itemsAccepted,
|
||
itemsDeduplicated,
|
||
envId: environment.id,
|
||
batchStatus: currentBatch?.status,
|
||
batchSealed: currentBatch?.sealed,
|
||
processingCompletedAt: currentBatch?.processingCompletedAt,
|
||
});
|
||
|
||
span.setAttribute("itemsAccepted", itemsAccepted);
|
||
span.setAttribute("itemsDeduplicated", itemsDeduplicated);
|
||
span.setAttribute("sealedByConcurrentRequest", true);
|
||
|
||
return {
|
||
id: batchFriendlyId,
|
||
itemsAccepted,
|
||
itemsDeduplicated,
|
||
sealed: true,
|
||
runCount: batch.runCount,
|
||
};
|
||
}
|
||
|
||
// Batch is in an unexpected state - fail with error
|
||
const actualStatus = currentBatch?.status ?? "unknown";
|
||
const actualSealed = currentBatch?.sealed ?? "unknown";
|
||
logger.error("Batch seal race condition: unexpected state", {
|
||
batchId: batchFriendlyId,
|
||
expectedStatus: "PENDING",
|
||
actualStatus,
|
||
expectedSealed: false,
|
||
actualSealed,
|
||
envId: environment.id,
|
||
});
|
||
|
||
throw new ServiceValidationError(
|
||
`Batch ${batchFriendlyId} is in unexpected state (status: ${actualStatus}, sealed: ${actualSealed}). Cannot seal batch.`
|
||
);
|
||
}
|
||
|
||
logger.info("Batch sealed and ready for processing", {
|
||
batchId: batchFriendlyId,
|
||
itemsAccepted,
|
||
itemsDeduplicated,
|
||
totalEnqueued: enqueuedCount,
|
||
envId: environment.id,
|
||
});
|
||
|
||
span.setAttribute("itemsAccepted", itemsAccepted);
|
||
span.setAttribute("itemsDeduplicated", itemsDeduplicated);
|
||
|
||
return {
|
||
id: batchFriendlyId,
|
||
itemsAccepted,
|
||
itemsDeduplicated,
|
||
sealed: true,
|
||
runCount: batch.runCount,
|
||
};
|
||
}
|
||
);
|
||
}
|
||
|
||
/**
|
||
* Process a single streamed batch item: validate it, offload its payload to
|
||
* object storage if oversized, and enqueue it. Returns whether the item was
|
||
* newly enqueued ("accepted") or was a duplicate ("deduplicated"). Throws
|
||
* ServiceValidationError for invalid items, which aborts the stream.
|
||
*
|
||
* Safe to run concurrently: enqueueBatchItem() is atomic and order-independent
|
||
* per item index, and each item carries its own index (real items from the
|
||
* SDK; oversized markers are stamped by the NDJSON parser).
|
||
*/
|
||
async #processItem(
|
||
rawItem: unknown,
|
||
batchId: string,
|
||
environment: AuthenticatedEnvironment,
|
||
runCount: number
|
||
): Promise<"accepted" | "deduplicated"> {
|
||
// Oversized item marker emitted by the NDJSON parser
|
||
if (rawItem && typeof rawItem === "object" && "__batchItemError" in rawItem) {
|
||
const marker = rawItem as OversizedItemMarker;
|
||
|
||
// Same out-of-range guard as normal items: an oversized item with an
|
||
// out-of-range index must 4xx rather than create a stray pre-failed run.
|
||
if (marker.index >= runCount) {
|
||
throw new ServiceValidationError(
|
||
`Item index ${marker.index} exceeds batch runCount ${runCount}`
|
||
);
|
||
}
|
||
|
||
const errorMessage = `Batch item payload is too large (${(marker.actualSize / 1024).toFixed(
|
||
1
|
||
)} KB). Maximum allowed size is ${(marker.maxSize / 1024).toFixed(
|
||
1
|
||
)} KB. Reduce the payload size or offload large data to external storage.`;
|
||
|
||
// Enqueue with __error metadata - processItemCallback will detect this
|
||
// and use TriggerFailedTaskService to create a pre-failed run
|
||
const batchItem: BatchItem = {
|
||
task: marker.task,
|
||
payload: "{}",
|
||
payloadType: "application/json",
|
||
options: {
|
||
__error: errorMessage,
|
||
__errorCode: "PAYLOAD_TOO_LARGE",
|
||
},
|
||
};
|
||
|
||
const result = await this._engine.enqueueBatchItem(
|
||
batchId,
|
||
environment.id,
|
||
marker.index,
|
||
batchItem
|
||
);
|
||
|
||
return result.enqueued ? "accepted" : "deduplicated";
|
||
}
|
||
|
||
// Parse and validate the item
|
||
const parseResult = BatchItemNDJSONSchema.safeParse(rawItem);
|
||
if (!parseResult.success) {
|
||
const rawIndex = (rawItem as { index?: unknown } | null)?.index;
|
||
const where = typeof rawIndex === "number" ? `index ${rawIndex}` : "unknown index";
|
||
throw new ServiceValidationError(`Invalid item at ${where}: ${parseResult.error.message}`);
|
||
}
|
||
|
||
const item = parseResult.data;
|
||
|
||
// Validate index is within expected range
|
||
if (item.index >= runCount) {
|
||
throw new ServiceValidationError(
|
||
`Item index ${item.index} exceeds batch runCount ${runCount}`
|
||
);
|
||
}
|
||
|
||
// Get the original payload type
|
||
const originalPayloadType = (item.options?.payloadType as string) ?? "application/json";
|
||
|
||
// Process payload - offload to object storage if it exceeds threshold
|
||
const processedPayload = await this.payloadProcessor.process(
|
||
item.payload,
|
||
originalPayloadType,
|
||
batchId,
|
||
item.index,
|
||
environment
|
||
);
|
||
|
||
// Convert to BatchItem format with potentially offloaded payload
|
||
const batchItem: BatchItem = {
|
||
task: item.task,
|
||
payload: processedPayload.payload,
|
||
payloadType: processedPayload.payloadType,
|
||
options: item.options,
|
||
};
|
||
|
||
// Enqueue the item
|
||
const result = await this._engine.enqueueBatchItem(
|
||
batchId,
|
||
environment.id,
|
||
item.index,
|
||
batchItem
|
||
);
|
||
|
||
return result.enqueued ? "accepted" : "deduplicated";
|
||
}
|
||
}
|
||
|
||
/**
|
||
* Extract `index` and `task` from raw JSON bytes without decoding the full line.
|
||
* Scans at most 512 bytes, tracking JSON nesting depth to only match top-level keys.
|
||
*/
|
||
export function extractIndexAndTask(bytes: Uint8Array): { index: number; task: string } {
|
||
let index = -1;
|
||
let task = "unknown";
|
||
let depth = 0;
|
||
let foundIndex = false;
|
||
let foundTask = false;
|
||
const limit = Math.min(bytes.byteLength, 512);
|
||
|
||
const QUOTE = 0x22; // "
|
||
const COLON = 0x3a; // :
|
||
const LBRACE = 0x7b; // {
|
||
const RBRACE = 0x7d; // }
|
||
const LBRACKET = 0x5b; // [
|
||
const RBRACKET = 0x5d; // ]
|
||
const BACKSLASH = 0x5c; // \
|
||
|
||
// Byte patterns for "index" and "task" (without quotes)
|
||
const INDEX_BYTES = [0x69, 0x6e, 0x64, 0x65, 0x78]; // index
|
||
const TASK_BYTES = [0x74, 0x61, 0x73, 0x6b]; // task
|
||
|
||
let i = 0;
|
||
while (i < limit && !(foundIndex && foundTask)) {
|
||
const b = bytes[i];
|
||
|
||
if (b === LBRACE || b === LBRACKET) {
|
||
depth++;
|
||
i++;
|
||
continue;
|
||
}
|
||
if (b === RBRACE || b === RBRACKET) {
|
||
depth--;
|
||
i++;
|
||
continue;
|
||
}
|
||
|
||
// Only match keys at depth 1 (top-level object)
|
||
if (b === QUOTE && depth === 1) {
|
||
// Read the key inside quotes
|
||
const keyStart = i + 1;
|
||
let keyEnd = keyStart;
|
||
while (keyEnd < limit && bytes[keyEnd] !== QUOTE) {
|
||
if (bytes[keyEnd] === BACKSLASH) keyEnd++; // skip escaped char
|
||
keyEnd++;
|
||
}
|
||
|
||
const keyLen = keyEnd - keyStart;
|
||
|
||
// Check if this key matches "index" or "task"
|
||
const isIndex =
|
||
!foundIndex &&
|
||
keyLen === INDEX_BYTES.length &&
|
||
INDEX_BYTES.every((b, j) => bytes[keyStart + j] === b);
|
||
const isTask =
|
||
!foundTask &&
|
||
keyLen === TASK_BYTES.length &&
|
||
TASK_BYTES.every((b, j) => bytes[keyStart + j] === b);
|
||
|
||
if (isIndex || isTask) {
|
||
// Skip past closing quote and find colon
|
||
let pos = keyEnd + 1;
|
||
while (pos < limit && bytes[pos] !== COLON) pos++;
|
||
pos++; // skip colon
|
||
// Skip whitespace
|
||
while (pos < limit && (bytes[pos] === 0x20 || bytes[pos] === 0x09)) pos++;
|
||
|
||
if (isIndex) {
|
||
// Parse digits
|
||
let num = 0;
|
||
let hasDigit = false;
|
||
while (pos < limit && bytes[pos] >= 0x30 && bytes[pos] <= 0x39) {
|
||
num = num * 10 + (bytes[pos] - 0x30);
|
||
hasDigit = true;
|
||
pos++;
|
||
}
|
||
if (hasDigit) {
|
||
index = num;
|
||
foundIndex = true;
|
||
}
|
||
} else {
|
||
// Parse quoted string value
|
||
if (pos < limit && bytes[pos] === QUOTE) {
|
||
const valStart = pos + 1;
|
||
let valEnd = valStart;
|
||
while (valEnd < limit && bytes[valEnd] !== QUOTE) {
|
||
if (bytes[valEnd] === BACKSLASH) valEnd++;
|
||
valEnd++;
|
||
}
|
||
// Decode just this slice
|
||
try {
|
||
task = new TextDecoder("utf-8", { fatal: true }).decode(
|
||
bytes.slice(valStart, valEnd)
|
||
);
|
||
foundTask = true;
|
||
} catch {
|
||
// Leave as "unknown"
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Skip past the key's closing quote
|
||
i = keyEnd + 1;
|
||
continue;
|
||
}
|
||
|
||
i++;
|
||
}
|
||
|
||
return { index, task };
|
||
}
|
||
|
||
/**
|
||
* Create an NDJSON parser transform stream.
|
||
*
|
||
* Converts a stream of Uint8Array chunks into parsed JSON objects.
|
||
* Each line in the NDJSON is parsed independently.
|
||
*
|
||
* Uses byte-buffer accumulation to:
|
||
* - Prevent OOM from unbounded string buffers
|
||
* - Properly handle multibyte UTF-8 characters across chunk boundaries
|
||
* - Check size limits on raw bytes before decoding
|
||
*
|
||
* @param maxItemBytes - Maximum allowed bytes per line (item)
|
||
* @returns TransformStream that outputs parsed JSON objects
|
||
*/
|
||
export function createNdjsonParserStream(
|
||
maxItemBytes: number
|
||
): TransformStream<Uint8Array, unknown> {
|
||
// Single decoder instance, reused for all lines
|
||
const decoder = new TextDecoder("utf-8", { fatal: true });
|
||
|
||
// Byte buffer: array of chunks with tracked total length
|
||
let chunks: Uint8Array[] = [];
|
||
let totalBytes = 0;
|
||
let lineNumber = 0;
|
||
// 0-based position of the next object we emit (parsed item or oversized
|
||
// marker). The parser is the single sequential point in the pipeline, so this
|
||
// is the authoritative source of item ordering — downstream consumers can
|
||
// process items concurrently and must not rely on processing order to derive
|
||
// an item's index. Used to back-fill an oversized marker's index when it
|
||
// couldn't be extracted from the (truncated) raw bytes.
|
||
let emittedCount = 0;
|
||
// When an oversized incomplete line is detected (Case 2), we must discard
|
||
// all remaining bytes of that line until the next newline delimiter.
|
||
let skipUntilNewline = false;
|
||
|
||
const NEWLINE_BYTE = 0x0a; // '\n'
|
||
|
||
/**
|
||
* Emit a parsed object or marker downstream and advance the emit position.
|
||
* Every emitted object MUST go through here so `emittedCount` stays aligned
|
||
* with item position (empty/skipped lines never emit, so they don't count).
|
||
*/
|
||
function emit(controller: TransformStreamDefaultController<unknown>, obj: unknown): void {
|
||
controller.enqueue(obj);
|
||
emittedCount++;
|
||
}
|
||
|
||
/**
|
||
* Concatenate all chunks into a single Uint8Array
|
||
*/
|
||
function concatenateChunks(): Uint8Array {
|
||
if (chunks.length === 0) {
|
||
return new Uint8Array(0);
|
||
}
|
||
if (chunks.length === 1) {
|
||
return chunks[0];
|
||
}
|
||
const result = new Uint8Array(totalBytes);
|
||
let offset = 0;
|
||
for (const chunk of chunks) {
|
||
result.set(chunk, offset);
|
||
offset += chunk.byteLength;
|
||
}
|
||
return result;
|
||
}
|
||
|
||
/**
|
||
* Find the index of the first newline byte in the buffer.
|
||
* Returns -1 if not found.
|
||
*/
|
||
function findNewlineIndex(): number {
|
||
let globalIndex = 0;
|
||
for (const chunk of chunks) {
|
||
for (let i = 0; i < chunk.byteLength; i++) {
|
||
if (chunk[i] === NEWLINE_BYTE) {
|
||
return globalIndex + i;
|
||
}
|
||
}
|
||
globalIndex += chunk.byteLength;
|
||
}
|
||
return -1;
|
||
}
|
||
|
||
/**
|
||
* Extract bytes from the buffer up to (but not including) the given index,
|
||
* and remove those bytes plus the delimiter from the buffer.
|
||
*/
|
||
function extractLine(newlineIndex: number): Uint8Array {
|
||
const fullBuffer = concatenateChunks();
|
||
const lineBytes = fullBuffer.slice(0, newlineIndex);
|
||
const remaining = fullBuffer.slice(newlineIndex + 1); // Skip the newline
|
||
|
||
// Reset buffer with remaining bytes
|
||
if (remaining.byteLength > 0) {
|
||
chunks = [remaining];
|
||
totalBytes = remaining.byteLength;
|
||
} else {
|
||
chunks = [];
|
||
totalBytes = 0;
|
||
}
|
||
|
||
return lineBytes;
|
||
}
|
||
|
||
/**
|
||
* Parse a line from bytes, handling whitespace trimming.
|
||
* Returns the parsed object or null for empty lines.
|
||
*/
|
||
function parseLine(
|
||
lineBytes: Uint8Array,
|
||
controller: TransformStreamDefaultController<unknown>
|
||
): void {
|
||
lineNumber++;
|
||
|
||
// Decode the line bytes (stream: false since this is a complete line)
|
||
let lineText: string;
|
||
try {
|
||
lineText = decoder.decode(lineBytes, { stream: false });
|
||
} catch (err) {
|
||
throw new Error(`Invalid UTF-8 at line ${lineNumber}: ${(err as Error).message}`);
|
||
}
|
||
|
||
const trimmed = lineText.trim();
|
||
if (!trimmed) {
|
||
return; // Skip empty lines
|
||
}
|
||
|
||
try {
|
||
const obj = JSON.parse(trimmed);
|
||
emit(controller, obj);
|
||
} catch (err) {
|
||
throw new Error(`Invalid JSON at line ${lineNumber}: ${(err as Error).message}`);
|
||
}
|
||
}
|
||
|
||
return new TransformStream<Uint8Array, unknown>({
|
||
transform(chunk, controller) {
|
||
// If we're skipping the remainder of an oversized line, scan for the
|
||
// next newline in this chunk and discard everything before it.
|
||
if (skipUntilNewline) {
|
||
const nlPos = chunk.indexOf(NEWLINE_BYTE);
|
||
if (nlPos === -1) {
|
||
// Entire chunk is still part of the oversized line — discard it
|
||
return;
|
||
}
|
||
// Found the newline — keep everything after it
|
||
skipUntilNewline = false;
|
||
const remaining = chunk.slice(nlPos + 1);
|
||
if (remaining.byteLength === 0) {
|
||
return;
|
||
}
|
||
// Replace chunk with the remainder and fall through to normal processing
|
||
chunk = remaining;
|
||
}
|
||
|
||
// Append chunk to buffer
|
||
chunks.push(chunk);
|
||
totalBytes += chunk.byteLength;
|
||
|
||
// Process all complete lines in the buffer
|
||
let newlineIndex: number;
|
||
while ((newlineIndex = findNewlineIndex()) !== -1) {
|
||
// Check size limit BEFORE extracting/decoding (bytes up to newline)
|
||
if (newlineIndex > maxItemBytes) {
|
||
// Case 1: Complete line exceeds limit - emit marker instead of throwing
|
||
const lineBytes = extractLine(newlineIndex);
|
||
const extracted = extractIndexAndTask(lineBytes);
|
||
const marker: OversizedItemMarker = {
|
||
__batchItemError: "OVERSIZED",
|
||
index: extracted.index >= 0 ? extracted.index : emittedCount,
|
||
task: extracted.task,
|
||
actualSize: newlineIndex,
|
||
maxSize: maxItemBytes,
|
||
};
|
||
emit(controller, marker);
|
||
lineNumber++;
|
||
continue;
|
||
}
|
||
|
||
const lineBytes = extractLine(newlineIndex);
|
||
parseLine(lineBytes, controller);
|
||
}
|
||
|
||
// Check if the remaining buffer (incomplete line) exceeds the limit
|
||
// This prevents OOM from a single huge line without newlines
|
||
if (totalBytes > maxItemBytes) {
|
||
// Case 2: Incomplete line exceeds limit - emit marker instead of throwing
|
||
const extracted = extractIndexAndTask(concatenateChunks());
|
||
const marker: OversizedItemMarker = {
|
||
__batchItemError: "OVERSIZED",
|
||
index: extracted.index >= 0 ? extracted.index : emittedCount,
|
||
task: extracted.task,
|
||
actualSize: totalBytes,
|
||
maxSize: maxItemBytes,
|
||
};
|
||
emit(controller, marker);
|
||
lineNumber++;
|
||
// Clear buffer and skip remaining bytes of this oversized line
|
||
// until the next newline delimiter is found in a subsequent chunk
|
||
chunks = [];
|
||
totalBytes = 0;
|
||
skipUntilNewline = true;
|
||
return;
|
||
}
|
||
},
|
||
|
||
flush(controller) {
|
||
// Flush any remaining bytes from the decoder's internal state
|
||
// This handles multibyte characters that may have been split across chunks
|
||
decoder.decode(new Uint8Array(0), { stream: false });
|
||
|
||
// Process any remaining buffered data (no trailing newline case)
|
||
if (totalBytes === 0) {
|
||
return;
|
||
}
|
||
|
||
// Check size limit before processing final line
|
||
if (totalBytes > maxItemBytes) {
|
||
// Case 3: Flush with oversized remaining - emit marker instead of throwing
|
||
const extracted = extractIndexAndTask(concatenateChunks());
|
||
const marker: OversizedItemMarker = {
|
||
__batchItemError: "OVERSIZED",
|
||
index: extracted.index >= 0 ? extracted.index : emittedCount,
|
||
task: extracted.task,
|
||
actualSize: totalBytes,
|
||
maxSize: maxItemBytes,
|
||
};
|
||
emit(controller, marker);
|
||
return;
|
||
}
|
||
|
||
const finalBytes = concatenateChunks();
|
||
parseLine(finalBytes, controller);
|
||
},
|
||
});
|
||
}
|
||
|
||
/**
|
||
* Convert a ReadableStream into an AsyncIterable.
|
||
* Useful for processing streams with for-await-of loops.
|
||
*/
|
||
export async function* streamToAsyncIterable<T>(stream: ReadableStream<T>): AsyncIterable<T> {
|
||
const reader = stream.getReader();
|
||
try {
|
||
while (true) {
|
||
const { done, value } = await reader.read();
|
||
if (done) break;
|
||
yield value;
|
||
}
|
||
} finally {
|
||
reader.releaseLock();
|
||
}
|
||
}
|