Files
2026-07-13 13:32:57 +08:00

500 lines
17 KiB
TypeScript

import { Logger } from "@trigger.dev/core/logger";
import { tryCatch } from "@trigger.dev/core/utils";
import { getMeter, type Counter, type Histogram, type Meter } from "@internal/tracing";
import { nanoid } from "nanoid";
import pLimit from "p-limit";
import { signalsEmitter } from "~/services/signals.server";
export type DynamicFlushSchedulerConfig<T> = {
batchSize: number;
flushInterval: number;
callback: (flushId: string, batch: T[]) => Promise<void>;
// New configuration options
minConcurrency?: number;
maxConcurrency?: number;
maxBatchSize?: number;
memoryPressureThreshold?: number; // Number of items that triggers increased concurrency
loadSheddingThreshold?: number; // Number of items that triggers load shedding
loadSheddingEnabled?: boolean;
isDroppableEvent?: (item: T) => boolean; // Function to determine if an event can be dropped
// Self-observability. `name` is the low-cardinality `scheduler` label that separates the
// task_events / llm_metrics / otlp_metrics instances in the same process. `meter` defaults to
// the global provider; inject one in tests. Instruments are no-op unless metrics are enabled.
meter?: Meter;
name?: string;
};
export class DynamicFlushScheduler<T> {
private batchQueue: T[][];
private currentBatch: T[];
private readonly BATCH_SIZE: number;
private readonly FLUSH_INTERVAL: number;
private flushTimer: NodeJS.Timeout | null;
private metricsReporterTimer: NodeJS.Timeout | undefined;
private readonly callback: (flushId: string, batch: T[]) => Promise<void>;
// New properties for dynamic scaling
private readonly minConcurrency: number;
private readonly maxConcurrency: number;
private readonly maxBatchSize: number;
private readonly memoryPressureThreshold: number;
private limiter: ReturnType<typeof pLimit>;
private currentBatchSize: number;
private totalQueuedItems: number = 0;
private consecutiveFlushFailures: number = 0;
private lastFlushTime: number = Date.now();
private metrics = {
flushedBatches: 0,
failedBatches: 0,
totalItemsFlushed: 0,
droppedEvents: 0,
droppedEventsByKind: new Map<string, number>(),
};
private isShuttingDown: boolean = false;
// New properties for load shedding
private readonly loadSheddingThreshold: number;
private readonly loadSheddingEnabled: boolean;
private readonly isDroppableEvent?: (item: T) => boolean;
private isLoadShedding: boolean = false;
private readonly logger: Logger = new Logger("EventRepo.DynamicFlushScheduler", "info");
// Pre-allocated attribute objects (closed label sets) so the hot flush path never allocates.
private readonly _metricAttrs: { scheduler: string };
private readonly _batchOkAttrs: { scheduler: string; outcome: string };
private readonly _batchFailedAttrs: { scheduler: string; outcome: string };
private _batchesCounter?: Counter;
private _itemsCounter?: Counter;
private _flushDurationHistogram?: Histogram;
private _batchSizeHistogram?: Histogram;
private _droppedEventsCounter?: Counter;
constructor(config: DynamicFlushSchedulerConfig<T>) {
const schedulerName = config.name ?? "unknown";
this._metricAttrs = { scheduler: schedulerName };
this._batchOkAttrs = { scheduler: schedulerName, outcome: "ok" };
this._batchFailedAttrs = { scheduler: schedulerName, outcome: "failed" };
this.batchQueue = [];
this.currentBatch = [];
this.BATCH_SIZE = config.batchSize;
this.currentBatchSize = config.batchSize;
this.FLUSH_INTERVAL = config.flushInterval;
this.callback = config.callback;
this.flushTimer = null;
// Initialize dynamic scaling parameters
this.minConcurrency = config.minConcurrency ?? 1;
this.maxConcurrency = config.maxConcurrency ?? 10;
this.maxBatchSize = config.maxBatchSize ?? config.batchSize * 5;
this.memoryPressureThreshold = config.memoryPressureThreshold ?? config.batchSize * 20;
// Initialize load shedding parameters
this.loadSheddingThreshold = config.loadSheddingThreshold ?? config.batchSize * 50;
this.loadSheddingEnabled = config.loadSheddingEnabled ?? true;
this.isDroppableEvent = config.isDroppableEvent;
// Start with minimum concurrency
this.limiter = pLimit(this.minConcurrency);
this.startFlushTimer();
this.startMetricsReporter();
this.setupShutdownHandlers();
this.#setupOtelMetrics(config.meter, schedulerName);
}
#setupOtelMetrics(meterOverride: Meter | undefined, name: string): void {
const meter = meterOverride ?? getMeter("ingest-flush");
this._batchesCounter = meter.createCounter("ingest.flush.batches", {
description: "Batches flushed to the sink, by outcome",
unit: "batches",
});
this._itemsCounter = meter.createCounter("ingest.flush.items", {
description: "Items successfully flushed to the sink",
unit: "items",
});
this._flushDurationHistogram = meter.createHistogram("ingest.flush.duration", {
description: "Wall-clock duration of a single batch flush",
unit: "ms",
});
this._batchSizeHistogram = meter.createHistogram("ingest.flush.batch_size", {
description: "Number of items in a flushed batch",
unit: "items",
});
this._droppedEventsCounter = meter.createCounter("ingest.flush.dropped_events", {
description: "Events dropped by load shedding before they reached the sink",
unit: "events",
});
// Pull-based gauges: read at export time only, so they add zero hot-path cost.
const queueDepthGauge = meter.createObservableGauge("ingest.flush.queue_depth", {
description: "Items queued and awaiting flush",
unit: "items",
});
const concurrencyGauge = meter.createObservableGauge("ingest.flush.concurrency", {
description: "Current concurrent-flush limit",
unit: "flushes",
});
const loadSheddingGauge = meter.createObservableGauge("ingest.flush.load_shedding", {
description: "1 while actively shedding load, otherwise 0",
});
meter.addBatchObservableCallback(
(result) => {
result.observe(queueDepthGauge, this.totalQueuedItems, this._metricAttrs);
result.observe(concurrencyGauge, this.limiter.concurrency, this._metricAttrs);
result.observe(loadSheddingGauge, this.isLoadShedding ? 1 : 0, this._metricAttrs);
},
[queueDepthGauge, concurrencyGauge, loadSheddingGauge]
);
}
addToBatch(items: T[]): void {
let itemsToAdd = items;
// Apply load shedding if enabled and we're over the threshold
if (this.loadSheddingEnabled && this.totalQueuedItems >= this.loadSheddingThreshold) {
const { kept, dropped } = this.applyLoadShedding(items);
itemsToAdd = kept;
if (dropped.length > 0) {
this.metrics.droppedEvents += dropped.length;
this._droppedEventsCounter?.add(dropped.length, this._metricAttrs);
// Track dropped events by kind if possible
dropped.forEach((item) => {
const kind = this.getEventKind(item);
if (kind) {
const currentCount = this.metrics.droppedEventsByKind.get(kind) || 0;
this.metrics.droppedEventsByKind.set(kind, currentCount + 1);
}
});
if (!this.isLoadShedding) {
this.isLoadShedding = true;
}
this.logger.warn("Load shedding", {
totalQueuedItems: this.totalQueuedItems,
threshold: this.loadSheddingThreshold,
droppedCount: dropped.length,
});
}
} else if (this.isLoadShedding && this.totalQueuedItems < this.loadSheddingThreshold * 0.8) {
this.isLoadShedding = false;
this.logger.info("Load shedding deactivated", {
totalQueuedItems: this.totalQueuedItems,
threshold: this.loadSheddingThreshold,
totalDropped: this.metrics.droppedEvents,
});
}
this.currentBatch.push(...itemsToAdd);
this.totalQueuedItems += itemsToAdd.length;
// Check if we need to create a batch (if we are shutting down, create a batch immediately because the flush timer is stopped)
if (this.currentBatch.length >= this.currentBatchSize || this.isShuttingDown) {
this.createBatch();
}
// Adjust concurrency based on queue pressure
this.adjustConcurrency();
}
private createBatch(): void {
if (this.currentBatch.length === 0) return;
this.batchQueue.push(this.currentBatch);
this.currentBatch = [];
this.flushBatches();
this.resetFlushTimer();
}
private setupShutdownHandlers(): void {
signalsEmitter.on("SIGTERM", () =>
this.shutdown().catch((error) => {
this.logger.error("Error shutting down dynamic flush scheduler", {
error,
});
})
);
signalsEmitter.on("SIGINT", () =>
this.shutdown().catch((error) => {
this.logger.error("Error shutting down dynamic flush scheduler", {
error,
});
})
);
}
private startFlushTimer(): void {
this.flushTimer = setInterval(() => this.checkAndFlush(), this.FLUSH_INTERVAL);
}
private resetFlushTimer(): void {
if (this.flushTimer) {
clearInterval(this.flushTimer);
}
if (this.isShuttingDown) return;
this.startFlushTimer();
}
private checkAndFlush(): void {
if (this.currentBatch.length > 0) {
this.createBatch();
}
this.flushBatches();
}
private async flushBatches(): Promise<void> {
const batchesToFlush: T[][] = [];
// Dequeue all available batches up to current concurrency limit
while (this.batchQueue.length > 0 && batchesToFlush.length < this.limiter.concurrency) {
const batch = this.batchQueue.shift();
if (batch) {
batchesToFlush.push(batch);
}
}
if (batchesToFlush.length === 0) return;
// Schedule all batches for concurrent processing
const flushPromises = batchesToFlush.map((batch) =>
this.limiter(async () => {
const itemCount = batch.length;
// eslint-disable-next-line no-this-alias
const self = this;
async function tryFlush(flushId: string, batchToFlush: T[], attempt: number = 1) {
try {
const startTime = Date.now();
await self.callback(flushId, batchToFlush);
const duration = Date.now() - startTime;
self.totalQueuedItems -= itemCount;
self.consecutiveFlushFailures = 0;
self.lastFlushTime = Date.now();
self.metrics.flushedBatches++;
self.metrics.totalItemsFlushed += itemCount;
self._flushDurationHistogram?.record(duration, self._metricAttrs);
self._batchSizeHistogram?.record(itemCount, self._metricAttrs);
self._itemsCounter?.add(itemCount, self._metricAttrs);
self._batchesCounter?.add(1, self._batchOkAttrs);
self.logger.debug("Batch flushed successfully", {
flushId,
itemCount,
duration,
remainingQueueDepth: self.totalQueuedItems,
activeConcurrency: self.limiter.activeCount,
pendingConcurrency: self.limiter.pendingCount,
});
} catch (error) {
self.consecutiveFlushFailures++;
self.metrics.failedBatches++;
self.logger.error("Error attempting to flush batch", {
flushId,
itemCount,
error,
consecutiveFailures: self.consecutiveFlushFailures,
attempt,
});
// Back off on failures
if (self.consecutiveFlushFailures > 5) {
self.adjustConcurrency(true);
}
if (attempt <= 3) {
await new Promise((resolve) => setTimeout(resolve, 500));
return await tryFlush(flushId, batchToFlush, attempt + 1);
} else {
throw error;
}
}
}
const [flushError] = await tryCatch(tryFlush(nanoid(), batch));
if (flushError) {
this.logger.error("Error flushing batch", {
error: flushError,
});
this._batchesCounter?.add(1, this._batchFailedAttrs);
}
})
);
// Don't await here - let them run concurrently
Promise.allSettled(flushPromises).then(() => {
const shouldContinueFlushing =
this.batchQueue.length > 0 && (this.consecutiveFlushFailures < 3 || this.isShuttingDown);
// After flush completes, check if we need to flush more
if (shouldContinueFlushing) {
this.flushBatches();
}
});
}
private lastConcurrencyAdjustment: number = Date.now();
private adjustConcurrency(backOff: boolean = false): void {
const currentConcurrency = this.limiter.concurrency;
let newConcurrency = currentConcurrency;
// Calculate pressure metrics - moved outside the if/else block
const queuePressure = this.totalQueuedItems / this.memoryPressureThreshold;
const timeSinceLastFlush = Date.now() - this.lastFlushTime;
const timeSinceLastAdjustment = Date.now() - this.lastConcurrencyAdjustment;
// Don't adjust too frequently (except for backoff)
if (!backOff && timeSinceLastAdjustment < 1000) {
return;
}
if (backOff) {
// Reduce concurrency on failures
newConcurrency = Math.max(this.minConcurrency, Math.floor(currentConcurrency * 0.75));
} else {
if (queuePressure > 0.8 || timeSinceLastFlush > this.FLUSH_INTERVAL * 2) {
// High pressure - increase concurrency
newConcurrency = Math.min(this.maxConcurrency, currentConcurrency + 2);
} else if (queuePressure < 0.2 && currentConcurrency > this.minConcurrency) {
// Low pressure - decrease concurrency
newConcurrency = Math.max(this.minConcurrency, currentConcurrency - 1);
}
}
// Adjust batch size based on pressure
if (this.totalQueuedItems > this.memoryPressureThreshold) {
this.currentBatchSize = Math.min(
this.maxBatchSize,
Math.floor(this.BATCH_SIZE * (1 + queuePressure))
);
} else {
this.currentBatchSize = this.BATCH_SIZE;
}
// Update concurrency if changed
if (newConcurrency !== currentConcurrency) {
this.limiter = pLimit(newConcurrency);
this.logger.debug("Adjusted flush concurrency", {
previousConcurrency: currentConcurrency,
newConcurrency,
queuePressure,
totalQueuedItems: this.totalQueuedItems,
currentBatchSize: this.currentBatchSize,
memoryPressureThreshold: this.memoryPressureThreshold,
});
}
}
private startMetricsReporter(): void {
// Report metrics every 30 seconds
this.metricsReporterTimer = setInterval(() => {
const droppedByKind: Record<string, number> = {};
this.metrics.droppedEventsByKind.forEach((count, kind) => {
droppedByKind[kind] = count;
});
this.logger.debug("DynamicFlushScheduler metrics", {
totalQueuedItems: this.totalQueuedItems,
batchQueueLength: this.batchQueue.length,
currentBatchLength: this.currentBatch.length,
currentConcurrency: this.limiter.concurrency,
activeConcurrent: this.limiter.activeCount,
pendingConcurrent: this.limiter.pendingCount,
currentBatchSize: this.currentBatchSize,
isLoadShedding: this.isLoadShedding,
metrics: {
...this.metrics,
droppedByKind,
},
});
}, 30000);
}
private applyLoadShedding(items: T[]): { kept: T[]; dropped: T[] } {
if (!this.isDroppableEvent) {
// If no function provided to determine droppable events, keep all
return { kept: items, dropped: [] };
}
const kept: T[] = [];
const dropped: T[] = [];
for (const item of items) {
if (this.isDroppableEvent(item)) {
dropped.push(item);
} else {
kept.push(item);
}
}
return { kept, dropped };
}
private getEventKind(item: T): string | undefined {
// Try to extract the kind from the event if it has one
if (item && typeof item === "object" && "kind" in item) {
return String(item.kind);
}
return undefined;
}
// Method to get current status
getStatus() {
const droppedByKind: Record<string, number> = {};
this.metrics.droppedEventsByKind.forEach((count, kind) => {
droppedByKind[kind] = count;
});
return {
queuedItems: this.totalQueuedItems,
batchQueueLength: this.batchQueue.length,
currentBatchSize: this.currentBatch.length,
concurrency: this.limiter.concurrency,
activeFlushes: this.limiter.activeCount,
pendingFlushes: this.limiter.pendingCount,
isLoadShedding: this.isLoadShedding,
metrics: {
...this.metrics,
droppedEventsByKind: droppedByKind,
},
};
}
// Graceful shutdown
async shutdown(): Promise<void> {
if (this.isShuttingDown) return;
this.isShuttingDown = true;
if (this.flushTimer) {
clearInterval(this.flushTimer);
}
if (this.metricsReporterTimer) {
clearInterval(this.metricsReporterTimer);
}
// Flush any remaining items
if (this.currentBatch.length > 0) {
this.createBatch();
}
// Wait for all pending flushes to complete
while (this.batchQueue.length > 0 || this.limiter.activeCount > 0) {
await new Promise((resolve) => setTimeout(resolve, 100));
}
}
}