Improve dynamic flush scheduler for otel data

ericallam · ericallam · commit a34f03f2a803 · 2025-07-26T09:39:36.000+01:00
The changes introduce a more flexible and adaptive dynamic flush scheduler to address production issues where the system wasn't flushing data fast enough, causing memory growth and crashes. This issue arises from the existing scheduler handling only a single flush at a time, limiting concurrency and failing to cope with the influx of logs.

- Added configuration options for setting minimum and maximum concurrency levels, maximum batch size, and memory pressure threshold. These parameters ensure that flush operations adjust dynamically based on workload and pressure.
- Implemented `pLimit` to facilitate concurrent flush operations, with adjustments made according to batch queue length and memory pressure.
- Metrics reporting improvements were added to monitor the dynamic behavior of the flush scheduler, aiding in identifying performance issues and optimizing the operation accordingly.
diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts
@@ -252,6 +252,10 @@ const EnvironmentSchema = z.object({
   EVENTS_BATCH_SIZE: z.coerce.number().int().default(100),
   EVENTS_BATCH_INTERVAL: z.coerce.number().int().default(1000),
   EVENTS_DEFAULT_LOG_RETENTION: z.coerce.number().int().default(7),
+  EVENTS_MIN_CONCURRENCY: z.coerce.number().int().default(1),
+  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),
   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),
diff --git a/apps/webapp/app/v3/dynamicFlushScheduler.server.ts b/apps/webapp/app/v3/dynamicFlushScheduler.server.ts
@@ -1,38 +1,84 @@
 import { nanoid } from "nanoid";
+import pLimit from "p-limit";
+import { logger } from "~/services/logger.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
 };
 
 export class DynamicFlushScheduler<T> {
-  private batchQueue: T[][]; // Adjust the type according to your data structure
-  private currentBatch: T[]; // Adjust the type according to your data structure
+  private batchQueue: T[][]; 
+  private currentBatch: T[]; 
   private readonly BATCH_SIZE: number;
   private readonly FLUSH_INTERVAL: number;
   private flushTimer: NodeJS.Timeout | null;
   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,
+  };
 
   constructor(config: DynamicFlushSchedulerConfig<T>) {
     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;
+    
+    // Start with minimum concurrency
+    this.limiter = pLimit(this.minConcurrency);
+    
     this.startFlushTimer();
+    this.startMetricsReporter();
   }
 
   addToBatch(items: T[]): void {
     this.currentBatch.push(...items);
+    this.totalQueuedItems += items.length;
 
-    if (this.currentBatch.length >= this.BATCH_SIZE) {
-      this.batchQueue.push(this.currentBatch);
-      this.currentBatch = [];
-      this.flushNextBatch();
-      this.resetFlushTimer();
+    // Check if we need to create a batch
+    if (this.currentBatch.length >= this.currentBatchSize) {
+      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 startFlushTimer(): void {
@@ -48,23 +94,169 @@ export class DynamicFlushScheduler<T> {
 
   private checkAndFlush(): void {
     if (this.currentBatch.length > 0) {
-      this.batchQueue.push(this.currentBatch);
-      this.currentBatch = [];
+      this.createBatch();
     }
-    this.flushNextBatch();
+    this.flushBatches();
   }
 
-  private async flushNextBatch(): Promise<void> {
-    if (this.batchQueue.length === 0) return;
+  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;
 
-    const batchToFlush = this.batchQueue.shift();
-    try {
-      await this.callback(nanoid(), batchToFlush!);
+    // Schedule all batches for concurrent processing
+    const flushPromises = batchesToFlush.map((batch) =>
+      this.limiter(async () => {
+        const flushId = nanoid();
+        const itemCount = batch.length;
+        
+        try {
+          const startTime = Date.now();
+          await this.callback(flushId, batch);
+          
+          const duration = Date.now() - startTime;
+          this.totalQueuedItems -= itemCount;
+          this.consecutiveFlushFailures = 0;
+          this.lastFlushTime = Date.now();
+          this.metrics.flushedBatches++;
+          this.metrics.totalItemsFlushed += itemCount;
+          
+          logger.debug("Batch flushed successfully", {
+            flushId,
+            itemCount,
+            duration,
+            remainingQueueDepth: this.totalQueuedItems,
+            activeConcurrency: this.limiter.activeCount,
+            pendingConcurrency: this.limiter.pendingCount,
+          });
+        } catch (error) {
+          this.consecutiveFlushFailures++;
+          this.metrics.failedBatches++;
+          
+          logger.error("Error flushing batch", {
+            flushId,
+            itemCount,
+            error,
+            consecutiveFailures: this.consecutiveFlushFailures,
+          });
+          
+          // Re-queue the batch at the front if it fails
+          this.batchQueue.unshift(batch);
+          this.totalQueuedItems += itemCount;
+          
+          // Back off on failures
+          if (this.consecutiveFlushFailures > 3) {
+            this.adjustConcurrency(true);
+          }
+        }
+      })
+    );
+    
+    // Don't await here - let them run concurrently
+    Promise.allSettled(flushPromises).then(() => {
+      // After flush completes, check if we need to flush more
       if (this.batchQueue.length > 0) {
-        this.flushNextBatch();
+        this.flushBatches();
       }
-    } catch (error) {
-      console.error("Error inserting batch:", error);
+    });
+  }
+
+  private adjustConcurrency(backOff: boolean = false): void {
+    const currentConcurrency = this.limiter.concurrency;
+    let newConcurrency = currentConcurrency;
+    
+    if (backOff) {
+      // Reduce concurrency on failures
+      newConcurrency = Math.max(this.minConcurrency, Math.floor(currentConcurrency * 0.75));
+    } else {
+      // Calculate pressure metrics
+      const queuePressure = this.totalQueuedItems / this.memoryPressureThreshold;
+      const timeSinceLastFlush = Date.now() - this.lastFlushTime;
+      
+      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);
+      
+      logger.info("Adjusted flush concurrency", {
+        previousConcurrency: currentConcurrency,
+        newConcurrency,
+        queuePressure,
+        totalQueuedItems: this.totalQueuedItems,
+        currentBatchSize: this.currentBatchSize,
+      });
+    }
+  }
+  
+  private startMetricsReporter(): void {
+    // Report metrics every 30 seconds
+    setInterval(() => {
+      logger.info("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,
+        metrics: this.metrics,
+      });
+    }, 30000);
+  }
+
+  // Method to get current status
+  getStatus() {
+    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 },
+    };
+  }
+
+  // Graceful shutdown
+  async shutdown(): Promise<void> {
+    if (this.flushTimer) {
+      clearInterval(this.flushTimer);
+    }
+    
+    // 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));
     }
   }
-}
+}
diff --git a/apps/webapp/app/v3/eventRepository.server.ts b/apps/webapp/app/v3/eventRepository.server.ts
@@ -107,6 +107,10 @@ export type EventRepoConfig = {
   retentionInDays: number;
   partitioningEnabled: boolean;
   tracer?: Tracer;
+  minConcurrency?: number;
+  maxConcurrency?: number;
+  maxBatchSize?: number;
+  memoryPressureThreshold?: number;
 };
 
 export type QueryOptions = Prisma.TaskEventWhereInput;
@@ -199,6 +203,10 @@ export class EventRepository {
     return this._subscriberCount;
   }
 
+  get flushSchedulerStatus() {
+    return this._flushScheduler.getStatus();
+  }
+
   constructor(
     db: PrismaClient = prisma,
     readReplica: PrismaReplicaClient = $replica,
@@ -208,6 +216,10 @@ export class EventRepository {
       batchSize: _config.batchSize,
       flushInterval: _config.batchInterval,
       callback: this.#flushBatch.bind(this),
+      minConcurrency: _config.minConcurrency,
+      maxConcurrency: _config.maxConcurrency,
+      maxBatchSize: _config.maxBatchSize,
+      memoryPressureThreshold: _config.memoryPressureThreshold,
     });
 
     this._redisPublishClient = createRedisClient("trigger:eventRepoPublisher", this._config.redis);
@@ -1324,6 +1336,10 @@ function initializeEventRepo() {
     batchInterval: env.EVENTS_BATCH_INTERVAL,
     retentionInDays: env.EVENTS_DEFAULT_LOG_RETENTION,
     partitioningEnabled: env.TASK_EVENT_PARTITIONING_ENABLED === "1",
+    minConcurrency: env.EVENTS_MIN_CONCURRENCY,
+    maxConcurrency: env.EVENTS_MAX_CONCURRENCY,
+    maxBatchSize: env.EVENTS_MAX_BATCH_SIZE,
+    memoryPressureThreshold: env.EVENTS_MEMORY_PRESSURE_THRESHOLD,
     redis: {
       port: env.PUBSUB_REDIS_PORT,
       host: env.PUBSUB_REDIS_HOST,
@@ -1343,6 +1359,47 @@ function initializeEventRepo() {
     registers: [metricsRegister],
   });
 
+  // Add metrics for flush scheduler
+  new Gauge({
+    name: "event_flush_scheduler_queued_items",
+    help: "Total number of items queued in the flush scheduler",
+    collect() {
+      const status = repo.flushSchedulerStatus;
+      this.set(status.queuedItems);
+    },
+    registers: [metricsRegister],
+  });
+
+  new Gauge({
+    name: "event_flush_scheduler_batch_queue_length",
+    help: "Number of batches waiting to be flushed",
+    collect() {
+      const status = repo.flushSchedulerStatus;
+      this.set(status.batchQueueLength);
+    },
+    registers: [metricsRegister],
+  });
+
+  new Gauge({
+    name: "event_flush_scheduler_concurrency",
+    help: "Current concurrency level of the flush scheduler",
+    collect() {
+      const status = repo.flushSchedulerStatus;
+      this.set(status.concurrency);
+    },
+    registers: [metricsRegister],
+  });
+
+  new Gauge({
+    name: "event_flush_scheduler_active_flushes",
+    help: "Number of active flush operations",
+    collect() {
+      const status = repo.flushSchedulerStatus;
+      this.set(status.activeFlushes);
+    },
+    registers: [metricsRegister],
+  });
+
   return repo;
 }
 
