chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,837 @@
|
||||
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();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user