Implement load shedding for TaskEvent records

This change introduces load shedding mechanisms to manage TaskEvent
records, particularly those of kind LOG, when the system experiences
high volumes and is unable to flush to the database in a timely
manner. The addition aims to prevent overwhelming the system and
ensure critical tasks are prioritized.

- Added configuration options for `loadSheddingThreshold` and
  `loadSheddingEnabled` in multiple modules to activate load shedding.
- Introduced `isDroppableEvent` function to allow specific events to
  be dropped when load shedding is enabled.
- Ensured metrics are updated to reflect dropped events and load
  shedding status, providing visibility into system performance
during high load conditions.
- Updated loggers to inform about load shedding state changes,
  ensuring timely awareness of load management activities.

Author: ericallam

apps/webapp/app/env.server.ts

@@ -256,6 +256,8 @@ const EnvironmentSchema = z.object({
   EVENTS_MAX_CONCURRENCY: z.coerce.number().int().default(10),
   EVENTS_MAX_BATCH_SIZE: z.coerce.number().int().default(500),
   EVENTS_MEMORY_PRESSURE_THRESHOLD: z.coerce.number().int().default(2000),
+  EVENTS_LOAD_SHEDDING_THRESHOLD: z.coerce.number().int().default(5000),
+  EVENTS_LOAD_SHEDDING_ENABLED: z.coerce.boolean().default(true),
   SHARED_QUEUE_CONSUMER_POOL_SIZE: z.coerce.number().int().default(10),
   SHARED_QUEUE_CONSUMER_INTERVAL_MS: z.coerce.number().int().default(100),
   SHARED_QUEUE_CONSUMER_NEXT_TICK_INTERVAL_MS: z.coerce.number().int().default(100),

apps/webapp/app/v3/dynamicFlushScheduler.server.ts

@@ -1,6 +1,7 @@
 import { nanoid } from "nanoid";
 import pLimit from "p-limit";
 import { logger } from "~/services/logger.server";
+import { TaskEventKind } from "@trigger.dev/database";
 
 export type DynamicFlushSchedulerConfig<T> = {
   batchSize: number;
@@ -11,6 +12,9 @@ export type DynamicFlushSchedulerConfig<T> = {
   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
 };
 
 export class DynamicFlushScheduler<T> {
@@ -35,8 +39,16 @@ export class DynamicFlushScheduler<T> {
     flushedBatches: 0,
     failedBatches: 0,
     totalItemsFlushed: 0,
+    droppedEvents: 0,
+    droppedEventsByKind: new Map<string, number>(),
   };
 
+  // New properties for load shedding
+  private readonly loadSheddingThreshold: number;
+  private readonly loadSheddingEnabled: boolean;
+  private readonly isDroppableEvent?: (item: T) => boolean;
+  private isLoadShedding: boolean = false;
+
   constructor(config: DynamicFlushSchedulerConfig<T>) {
     this.batchQueue = [];
     this.currentBatch = [];
@@ -52,6 +64,11 @@ export class DynamicFlushScheduler<T> {
     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);
 
@@ -60,8 +77,45 @@ export class DynamicFlushScheduler<T> {
   }
 
   addToBatch(items: T[]): void {
-    this.currentBatch.push(...items);
-    this.totalQueuedItems += items.length;
+    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;
+        
+        // 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;
+          logger.warn("Load shedding activated", {
+            totalQueuedItems: this.totalQueuedItems,
+            threshold: this.loadSheddingThreshold,
+            droppedCount: dropped.length,
+          });
+        }
+      }
+    } else if (this.isLoadShedding && this.totalQueuedItems < this.loadSheddingThreshold * 0.8) {
+      this.isLoadShedding = false;
+      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 (this.currentBatch.length >= this.currentBatchSize) {
@@ -217,6 +271,11 @@ export class DynamicFlushScheduler<T> {
   private startMetricsReporter(): void {
     // Report metrics every 30 seconds
     setInterval(() => {
+      const droppedByKind: Record<string, number> = {};
+      this.metrics.droppedEventsByKind.forEach((count, kind) => {
+        droppedByKind[kind] = count;
+      });
+      
       logger.info("DynamicFlushScheduler metrics", {
         totalQueuedItems: this.totalQueuedItems,
         batchQueueLength: this.batchQueue.length,
@@ -225,21 +284,62 @@ export class DynamicFlushScheduler<T> {
         activeConcurrent: this.limiter.activeCount,
         pendingConcurrent: this.limiter.pendingCount,
         currentBatchSize: this.currentBatchSize,
-        metrics: this.metrics,
+        isLoadShedding: this.isLoadShedding,
+        metrics: {
+          ...this.metrics,
+          droppedEventsByKind,
+        },
       });
     }, 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,
-      metrics: { ...this.metrics },
+      isLoadShedding: this.isLoadShedding,
+      metrics: { 
+        ...this.metrics,
+        droppedEventsByKind: droppedByKind,
+      },
     };
   }
 

apps/webapp/app/v3/eventRepository.server.ts

@@ -111,6 +111,8 @@ export type EventRepoConfig = {
   maxConcurrency?: number;
   maxBatchSize?: number;
   memoryPressureThreshold?: number;
+  loadSheddingThreshold?: number;
+  loadSheddingEnabled?: boolean;
 };
 
 export type QueryOptions = Prisma.TaskEventWhereInput;
@@ -220,6 +222,12 @@ export class EventRepository {
       maxConcurrency: _config.maxConcurrency,
       maxBatchSize: _config.maxBatchSize,
       memoryPressureThreshold: _config.memoryPressureThreshold,
+      loadSheddingThreshold: _config.loadSheddingThreshold,
+      loadSheddingEnabled: _config.loadSheddingEnabled,
+      isDroppableEvent: (event: CreatableEvent) => {
+        // Only drop LOG events during load shedding
+        return event.kind === TaskEventKind.LOG;
+      },
     });
 
     this._redisPublishClient = createRedisClient("trigger:eventRepoPublisher", this._config.redis);
@@ -1340,6 +1348,8 @@ function initializeEventRepo() {
     maxConcurrency: env.EVENTS_MAX_CONCURRENCY,
     maxBatchSize: env.EVENTS_MAX_BATCH_SIZE,
     memoryPressureThreshold: env.EVENTS_MEMORY_PRESSURE_THRESHOLD,
+    loadSheddingThreshold: env.EVENTS_LOAD_SHEDDING_THRESHOLD,
+    loadSheddingEnabled: env.EVENTS_LOAD_SHEDDING_ENABLED,
     redis: {
       port: env.PUBSUB_REDIS_PORT,
       host: env.PUBSUB_REDIS_HOST,
@@ -1400,6 +1410,26 @@ function initializeEventRepo() {
     registers: [metricsRegister],
   });
 
+  new Gauge({
+    name: "event_flush_scheduler_dropped_events",
+    help: "Total number of events dropped due to load shedding",
+    collect() {
+      const status = repo.flushSchedulerStatus;
+      this.set(status.metrics.droppedEvents);
+    },
+    registers: [metricsRegister],
+  });
+
+  new Gauge({
+    name: "event_flush_scheduler_is_load_shedding",
+    help: "Whether load shedding is currently active (1 = active, 0 = inactive)",
+    collect() {
+      const status = repo.flushSchedulerStatus;
+      this.set(status.isLoadShedding ? 1 : 0);
+    },
+    registers: [metricsRegister],
+  });
+
   return repo;
 }