speed up deduplicating of runs by making it more efficient

ericallam · ericallam · commit 87784031f566 · 2026-01-13T10:05:06.000Z
diff --git a/apps/webapp/app/services/runsReplicationService.server.ts b/apps/webapp/app/services/runsReplicationService.server.ts
@@ -82,7 +82,7 @@ type TaskRunInsert = {
 export type RunsReplicationServiceEvents = {
   message: [{ lsn: string; message: PgoutputMessage; service: RunsReplicationService }];
   batchFlushed: [
-    { flushId: string; taskRunInserts: TaskRunInsertArray[]; payloadInserts: PayloadInsertArray[] },
+    { flushId: string; taskRunInserts: TaskRunInsertArray[]; payloadInserts: PayloadInsertArray[] }
   ];
 };
 
@@ -214,34 +214,14 @@ export class RunsReplicationService {
       flushInterval: options.flushIntervalMs ?? 100,
       maxConcurrency: options.maxFlushConcurrency ?? 100,
       callback: this.#flushBatch.bind(this),
-      // we can do some pre-merging to reduce the amount of data we need to send to clickhouse
-      mergeBatch: (existingBatch: TaskRunInsert[], newBatch: TaskRunInsert[]) => {
-        const merged = new Map<string, TaskRunInsert>();
-
-        for (const item of existingBatch) {
-          const key = `${item.event}_${item.run.id}`;
-          merged.set(key, item);
-        }
-
-        for (const item of newBatch) {
-          if (!item?.run?.id) {
-            this.logger.warn("Skipping replication event with null run", { event: item });
-            continue;
-          }
-
-          const key = `${item.event}_${item.run.id}`;
-          const existingItem = merged.get(key);
-
-          // Keep the run with the higher version (latest)
-          // and take the last occurrence for that version.
-          // Items originating from the same DB transaction have the same version.
-          if (!existingItem || item._version >= existingItem._version) {
-            merged.set(key, item);
-          }
-        }
-
-        return Array.from(merged.values());
+      // Key-based deduplication to reduce duplicates sent to ClickHouse
+      getKey: (item) => {
+        return `${item.event}_${item.run.id}`;
       },
+      // Keep the run with the higher version (latest)
+      // and take the last occurrence for that version.
+      // Items originating from the same DB transaction have the same version.
+      shouldReplace: (existing, incoming) => incoming._version >= existing._version,
       logger: new Logger("ConcurrentFlushScheduler", options.logLevel ?? "info"),
       tracer: options.tracer,
     });
@@ -504,7 +484,7 @@ export class RunsReplicationService {
       this._eventsProcessedCounter.add(1, { event_type: event.tag });
     }
 
-    this.logger.info("handle_transaction", {
+    this.logger.debug("handle_transaction", {
       transaction: {
         xid: transaction.xid,
         commitLsn: transaction.commitLsn,
@@ -974,13 +954,16 @@ export type ConcurrentFlushSchedulerConfig<T> = {
   flushInterval: number;
   maxConcurrency?: number;
   callback: (flushId: string, batch: T[]) => Promise<void>;
-  mergeBatch?: (existingBatch: T[], newBatch: T[]) => T[];
+  /** Key-based deduplication. Return null to skip the item. */
+  getKey: (item: T) => string | null;
+  /** Determine if incoming item should replace existing. */
+  shouldReplace: (existing: T, incoming: T) => boolean;
   tracer?: Tracer;
   logger?: Logger;
 };
 
 export class ConcurrentFlushScheduler<T> {
-  private currentBatch: T[];
+  private batch = new Map<string, T>();
   private readonly BATCH_SIZE: number;
   private readonly flushInterval: number;
   private readonly MAX_CONCURRENCY: number;
@@ -995,7 +978,6 @@ export class ConcurrentFlushScheduler<T> {
     this.logger = config.logger ?? new Logger("ConcurrentFlushScheduler", "info");
     this._tracer = config.tracer ?? trace.getTracer("concurrent-flush-scheduler");
 
-    this.currentBatch = [];
     this.BATCH_SIZE = config.batchSize;
     this.flushInterval = config.flushInterval;
     this.MAX_CONCURRENCY = config.maxConcurrency || 1;
@@ -1005,9 +987,17 @@ export class ConcurrentFlushScheduler<T> {
   }
 
   addToBatch(items: T[]): void {
-    this.currentBatch = this.config.mergeBatch
-      ? this.config.mergeBatch(this.currentBatch, items)
-      : this.currentBatch.concat(items);
+    for (const item of items) {
+      const key = this.config.getKey(item);
+      if (key === null) {
+        continue;
+      }
+
+      const existing = this.batch.get(key);
+      if (!existing || this.config.shouldReplace(existing, item)) {
+        this.batch.set(key, item);
+      }
+    }
 
     this.#flushNextBatchIfNeeded();
   }
@@ -1031,11 +1021,16 @@ export class ConcurrentFlushScheduler<T> {
     this.#flushNextBatchIfNeeded();
   }
 
+  #getBatchSize(): number {
+    return this.batch.size;
+  }
+
   #flushNextBatchIfNeeded(): void {
-    if (this.currentBatch.length >= this.BATCH_SIZE || this._isShutDown) {
+    const currentSize = this.#getBatchSize();
+    if (currentSize >= this.BATCH_SIZE || this._isShutDown) {
       this.logger.debug("Batch size threshold reached, initiating flush", {
         batchSize: this.BATCH_SIZE,
-        currentSize: this.currentBatch.length,
+        currentSize,
         isShutDown: this._isShutDown,
       });
 
@@ -1060,19 +1055,20 @@ export class ConcurrentFlushScheduler<T> {
   }
 
   async #checkAndFlush(): Promise<void> {
-    if (this.currentBatch.length > 0) {
+    const currentSize = this.#getBatchSize();
+    if (currentSize > 0) {
       this.logger.debug("Periodic flush check triggered", {
-        currentBatchSize: this.currentBatch.length,
+        currentBatchSize: currentSize,
       });
       await this.#flushNextBatch();
     }
   }
 
   async #flushNextBatch(): Promise<void> {
-    if (this.currentBatch.length === 0) return;
+    if (this.batch.size === 0) return;
 
-    const batch = this.currentBatch;
-    this.currentBatch = [];
+    const batch = Array.from(this.batch.values());
+    this.batch.clear();
 
     const callback = this.config.callback;
 
diff --git a/apps/webapp/test/concurrentFlushScheduler.test.ts b/apps/webapp/test/concurrentFlushScheduler.test.ts
@@ -0,0 +1,161 @@
+import { ConcurrentFlushScheduler } from "~/services/runsReplicationService.server";
+
+vi.setConfig({ testTimeout: 10_000 });
+
+type TestItem = {
+  id: string;
+  event: "insert" | "update";
+  version: number;
+};
+
+describe("ConcurrentFlushScheduler", () => {
+  it("should deduplicate items by key, keeping the latest version", async () => {
+    const flushedBatches: TestItem[][] = [];
+
+    const scheduler = new ConcurrentFlushScheduler<TestItem>({
+      batchSize: 100,
+      flushInterval: 50,
+      maxConcurrency: 1,
+      callback: async (_flushId, batch) => {
+        flushedBatches.push([...batch]);
+      },
+      getKey: (item) => `${item.event}_${item.id}`,
+      shouldReplace: (existing, incoming) => incoming.version >= existing.version,
+    });
+
+    scheduler.start();
+
+    // Add items with duplicate keys but different versions
+    scheduler.addToBatch([
+      { id: "run_1", event: "insert", version: 1 },
+      { id: "run_1", event: "update", version: 2 },
+      { id: "run_2", event: "insert", version: 1 },
+    ]);
+
+    // Add more items - should merge with existing
+    scheduler.addToBatch([
+      { id: "run_1", event: "insert", version: 3 }, // Higher version, should replace
+      { id: "run_1", event: "update", version: 1 }, // Lower version, should NOT replace
+      { id: "run_2", event: "update", version: 4 },
+    ]);
+
+    // Wait for flush
+    await new Promise((resolve) => setTimeout(resolve, 100));
+
+    scheduler.shutdown();
+
+    // Should have flushed once with deduplicated items
+    expect(flushedBatches.length).toBeGreaterThanOrEqual(1);
+
+    const allFlushed = flushedBatches.flat();
+
+    // Find items by their key
+    const insertRun1 = allFlushed.find((i) => i.id === "run_1" && i.event === "insert");
+    const updateRun1 = allFlushed.find((i) => i.id === "run_1" && i.event === "update");
+    const insertRun2 = allFlushed.find((i) => i.id === "run_2" && i.event === "insert");
+    const updateRun2 = allFlushed.find((i) => i.id === "run_2" && i.event === "update");
+
+    // Verify correct versions were kept
+    expect(insertRun1?.version).toBe(3); // Latest version for insert_run_1
+    expect(updateRun1?.version).toBe(2); // Original update_run_1 (v1 didn't replace v2)
+    expect(insertRun2?.version).toBe(1); // Only version for insert_run_2
+    expect(updateRun2?.version).toBe(4); // Only version for update_run_2
+  });
+
+  it("should skip items where getKey returns null", async () => {
+    const flushedBatches: TestItem[][] = [];
+
+    const scheduler = new ConcurrentFlushScheduler<TestItem>({
+      batchSize: 100,
+      flushInterval: 50,
+      maxConcurrency: 1,
+      callback: async (_flushId, batch) => {
+        flushedBatches.push([...batch]);
+      },
+      getKey: (item) => {
+        if (!item.id) {
+          return null;
+        }
+        return `${item.event}_${item.id}`;
+      },
+      shouldReplace: (existing, incoming) => incoming.version >= existing.version,
+    });
+
+    scheduler.start();
+
+    scheduler.addToBatch([
+      { id: "run_1", event: "insert", version: 1 },
+      { id: "", event: "insert", version: 2 }, // Should be skipped (null key)
+      { id: "run_2", event: "insert", version: 1 },
+    ]);
+
+    await new Promise((resolve) => setTimeout(resolve, 100));
+
+    scheduler.shutdown();
+
+    const allFlushed = flushedBatches.flat();
+    expect(allFlushed).toHaveLength(2);
+    expect(allFlushed.map((i) => i.id).sort()).toEqual(["run_1", "run_2"]);
+  });
+
+  it("should flush when batch size threshold is reached", async () => {
+    const flushedBatches: TestItem[][] = [];
+
+    const scheduler = new ConcurrentFlushScheduler<TestItem>({
+      batchSize: 3,
+      flushInterval: 10000, // Long interval so timer doesn't trigger
+      maxConcurrency: 1,
+      callback: async (_flushId, batch) => {
+        flushedBatches.push([...batch]);
+      },
+      getKey: (item) => `${item.event}_${item.id}`,
+      shouldReplace: (existing, incoming) => incoming.version >= existing.version,
+    });
+
+    scheduler.start();
+
+    // Add 3 unique items - should trigger flush
+    scheduler.addToBatch([
+      { id: "run_1", event: "insert", version: 1 },
+      { id: "run_2", event: "insert", version: 1 },
+      { id: "run_3", event: "insert", version: 1 },
+    ]);
+
+    await new Promise((resolve) => setTimeout(resolve, 50));
+
+    expect(flushedBatches.length).toBe(1);
+    expect(flushedBatches[0]).toHaveLength(3);
+
+    scheduler.shutdown();
+  });
+
+  it("should respect shouldReplace returning false", async () => {
+    const flushedBatches: TestItem[][] = [];
+
+    const scheduler = new ConcurrentFlushScheduler<TestItem>({
+      batchSize: 100,
+      flushInterval: 50,
+      maxConcurrency: 1,
+      callback: async (_flushId, batch) => {
+        flushedBatches.push([...batch]);
+      },
+      getKey: (item) => `${item.event}_${item.id}`,
+      // Never replace - first item wins
+      shouldReplace: () => false,
+    });
+
+    scheduler.start();
+
+    scheduler.addToBatch([{ id: "run_1", event: "insert", version: 10 }]);
+
+    scheduler.addToBatch([{ id: "run_1", event: "insert", version: 999 }]);
+
+    await new Promise((resolve) => setTimeout(resolve, 100));
+
+    scheduler.shutdown();
+
+    const allFlushed = flushedBatches.flat();
+    const insertRun1 = allFlushed.find((i) => i.id === "run_1" && i.event === "insert");
+    expect(insertRun1?.version).toBe(10); // First one wins
+  });
+});
diff --git a/apps/webapp/test/runsRepository.part2.test.ts b/apps/webapp/test/runsRepository.part2.test.ts
