improve streams sink

packages/dynamo-store/source/src/DynamoStoreSource.mts

@@ -305,14 +305,14 @@ export class DynamoStoreSource {
   private client: DynamoStoreSourceClient
 
   constructor(
-    private readonly index: AppendsIndex.Reader,
-    private epochs: AppendsEpoch.Reader.Service,
+    index: AppendsIndex.Reader,
+    epochs: AppendsEpoch.Reader.Service,
     private readonly options: Omit<CreateOptions, "context">,
   ) {
     if (!this.options.categories && !this.options.streamFilter) {
       throw new Error("Either categories or categoryFilter must be specified")
     }
-    const sink = new StreamsSink(
+    const sink = StreamsSink.create(
       options.handler,
       options.maxConcurrentStreams,
       options.maxConcurrentBatches,

packages/message-db/message-db-consumer/src/lib/MessageDbSource.mts

@@ -91,7 +91,7 @@ export class MessageDbSource {
 
   static create(options: Options & { pool: Pool }) {
     const client = new MessageDbCategoryReader(options.pool)
-    const sink = new StreamsSink(
+    const sink = StreamsSink.create(
       options.handler,
       options.maxConcurrentStreams,
       options.maxConcurrentBatches ?? 10,

packages/propeller/src/StreamsSink.mts

@@ -21,16 +21,82 @@ class Stream {
   }
 }
 
+class QueueWorker {
+  private stopped = true
+  constructor(
+    private readonly tryGetNext: (signal: AbortSignal) => Promise<Stream>,
+    private readonly handle: EventHandler<string>,
+    private readonly tracingAttrs: Attributes,
+    private readonly onHandled: (stream: Stream) => void,
+  ) {}
+
+  stop() {
+    this.stopped = true
+  }
+
+  async run(signal: AbortSignal): Promise<void> {
+    this.stopped = false
+    while (!signal.aborted && !this.stopped) {
+      const stream = await this.tryGetNext(signal)
+      await traceHandler(tracer, this.tracingAttrs, stream.name, stream.events, this.handle)
+      this.onHandled(stream)
+    }
+  }
+}
+
+export class BatchLimiter {
+  private inProgressBatches = 0
+  private onReady?: () => void
+  private onEmpty?: () => void
+  constructor(private maxConcurrentBatches: number) {}
+
+  waitForCapacity(signal: AbortSignal) {
+    if (this.inProgressBatches < this.maxConcurrentBatches) return
+    return new Promise<void>((resolve, reject) => {
+      const abort = () => {
+        signal.removeEventListener("abort", abort)
+        reject(new Error("Aborted"))
+      }
+      signal.addEventListener("abort", abort)
+      this.onReady = () => {
+        signal.removeEventListener("abort", abort)
+        resolve()
+      }
+    })
+  }
+
+  waitForEmpty() {
+    if (this.inProgressBatches === 0) return
+    return new Promise<void>((resolve) => {
+      this.onEmpty = () => {
+        resolve()
+      }
+    })
+  }
+
+  startBatch() {
+    this.inProgressBatches++
+  }
+
+  endBatch() {
+    this.inProgressBatches--
+    this.onReady?.()
+    delete this.onReady
+    if (this.inProgressBatches === 0) {
+      this.onEmpty?.()
+      delete this.onEmpty
+    }
+  }
+}
+
 export class StreamsSink implements Sink {
   private queue = new AsyncQueue<Stream>()
   private streams = new Map<StreamName, Stream>()
   private batchStreams = new Map<IngesterBatch, Set<Stream>>()
-  private onReady?: () => void
-  private inProgressBatches = 0
   constructor(
     private readonly handle: EventHandler<string>,
     private maxConcurrentStreams: number,
-    private maxConcurrentBatches: number,
+    private readonly batchLimiter: BatchLimiter,
     private readonly tracingAttrs: Attributes = {},
   ) {
     this.addTracingAttrs({ "eqx.max_concurrent_streams": maxConcurrentStreams })
@@ -40,59 +106,52 @@ export class StreamsSink implements Sink {
     Object.assign(this.tracingAttrs, attrs)
   }
 
-  start(signal: AbortSignal) {
+  private handleStreamCompletion(stream: Stream) {
+    for (const [batch, streams] of this.batchStreams) {
+      streams.delete(stream)
+      if (streams.size === 0) {
+        batch.onComplete()
+        this.batchStreams.delete(batch)
+        this.batchLimiter.endBatch()
+      }
+    }
+  }
+
+  async start(signal: AbortSignal) {
     // set up a linked abort controller to avoid attaching too many event listeners to the provided signal
     // node has a default limit of 11 before it emits a warning to the console. We have entirely legitimate reasons
     // for attaching more than 11 listeners
     const ctrl = new AbortController()
     signal.addEventListener("abort", () => ctrl.abort())
     const _signal = ctrl.signal
     // Starts N workers that will process streams in parallel
-    return new Promise<void>((_resolve, reject) => {
-      let stopped = false
-      const aux = async (): Promise<void> => {
-        if (_signal.aborted || stopped) return
-        try {
-          const stream = await this.queue.tryGetAsync(_signal)
-          this.streams.delete(stream.name)
-          await traceHandler(tracer, this.tracingAttrs, stream.name, stream.events, this.handle)
-          for (const [batch, streams] of this.batchStreams) {
-            streams.delete(stream)
-            if (streams.size === 0) {
-              batch.onComplete()
-              this.batchStreams.delete(batch)
-              this.inProgressBatches--
-              this.onReady?.()
-              delete this.onReady
-            }
-          }
-          void aux()
-        } catch (err) {
-          stopped = true
-          reject(err)
-        }
-      }
 
-      for (let i = 0; i < this.maxConcurrentStreams; ++i) aux()
-    })
-  }
+    const workers = new Array<QueueWorker>(this.maxConcurrentStreams)
 
-  private waitForCapacity(signal: AbortSignal) {
-    if (this.inProgressBatches < this.maxConcurrentBatches) return
-    return new Promise<void>((resolve, reject) => {
-      const abort = () => reject(new Error("Aborted"))
-      signal.addEventListener("abort", abort)
-      this.onReady = () => {
-        signal.removeEventListener("abort", abort)
-        resolve()
-      }
-    })
+    for (let i = 0; i < this.maxConcurrentStreams; ++i) {
+      workers[i] = new QueueWorker(
+        async (signal) => {
+          const stream = await this.queue.tryGetAsync(signal)
+          this.streams.delete(stream.name)
+          return stream
+        },
+        this.handle,
+        this.tracingAttrs,
+        (stream) => this.handleStreamCompletion(stream),
+      )
+    }
+
+    try {
+      await Promise.all(workers.map((w) => w.run(_signal)))
+    } catch (err) {
+      workers.forEach((w) => w.stop())
+      throw err
+    }
   }
 
-  pump(batch: IngesterBatch, signal: AbortSignal): Promise<void> | void {
-    const p = this.waitForCapacity(signal)
-    if (p != null) return p.then(() => this.pump(batch, signal))
-    this.inProgressBatches++
+  async pump(batch: IngesterBatch, signal: AbortSignal): Promise<void> {
+    await this.batchLimiter.waitForCapacity(signal)
+    this.batchLimiter.startBatch()
     const streamsInBatch = new Set<Stream>()
     this.batchStreams.set(batch, streamsInBatch)
     for (const [sn, event] of batch.items) {
@@ -106,4 +165,14 @@ export class StreamsSink implements Sink {
       }
     }
   }
+
+  static create(
+    handle: EventHandler<string>,
+    maxConcurrentStreams: number,
+    maxConcurrentBatches: number,
+    tracingAttrs: Attributes = {},
+  ) {
+    const limiter = new BatchLimiter(maxConcurrentBatches)
+    return new StreamsSink(handle, maxConcurrentStreams, limiter, tracingAttrs)
+  }
 }

packages/propeller/src/StreamsSink.test.mts

@@ -1,5 +1,5 @@
 import { describe, test, expect, vi } from "vitest"
-import { StreamsSink } from "./StreamsSink.mjs"
+import { BatchLimiter, StreamsSink } from "./StreamsSink.mjs"
 import { ITimelineEvent, StreamId, StreamName } from "@equinox-js/core"
 import { sleep } from "./Sleep.js"
 import { IngesterBatch } from "./Types.js"
@@ -36,7 +36,8 @@ describe("Concurrency", () => {
         await new Promise(setImmediate)
         active.delete(stream)
       }
-      const sink = new StreamsSink(handler, concurrency, 1)
+      const limiter = new BatchLimiter(1)
+      const sink = new StreamsSink(handler, concurrency, limiter)
 
       sink.start(ctrl.signal).catch(() => {})
 
@@ -45,7 +46,7 @@ describe("Concurrency", () => {
         ctrl.signal,
       )
 
-      await new Promise(setImmediate)
+      await limiter.waitForEmpty()
       ctrl.abort()
 
       expect(maxActive).toBe(concurrency)
@@ -60,7 +61,8 @@ test("Correctly merges batches", async () => {
     invocations++
     await new Promise(setImmediate)
   }
-  const sink = new StreamsSink(handler, 10, 3)
+  const limiter = new BatchLimiter(3)
+  const sink = new StreamsSink(handler, 10, limiter)
 
   const checkpoint = vi.fn().mockResolvedValue(undefined)
 
@@ -70,8 +72,7 @@ test("Correctly merges batches", async () => {
 
   sink.start(ctrl.signal).catch(() => {})
 
-  // sleep triggers after setImmediate
-  await sleep(0, ctrl.signal)
+  await limiter.waitForEmpty()
   ctrl.abort()
 
   expect(invocations).toBe(10)
@@ -95,7 +96,8 @@ test(" Correctly limits in-flight batches", async () => {
     invocations++
     await new Promise((res) => setTimeout(res, 10))
   }
-  const sink = new StreamsSink(handler, 1, 3)
+  const limiter = new BatchLimiter(3)
+  const sink = new StreamsSink(handler, 1, limiter)
 
   sink.start(ctrl.signal).catch(() => {})
   const completed = vi.fn().mockResolvedValue(undefined)
@@ -107,12 +109,12 @@ test(" Correctly limits in-flight batches", async () => {
   await sink.pump(mkSingleBatch(complete, 3n), ctrl.signal)
   await sink.pump(mkSingleBatch(complete, 4n), ctrl.signal)
 
-  await sleep(10, ctrl.signal)
+  await limiter.waitForEmpty()
   ctrl.abort()
 
   expect(invocations).toBe(3)
 
   // onComplete is called in order and for every batch
-  expect(completed).toHaveBeenCalledTimes(5)
   expect(completed.mock.calls).toEqual([[0n], [1n], [2n], [3n], [4n]])
+  expect(completed).toHaveBeenCalledTimes(5)
 })

packages/propeller/src/index.mts

@@ -1,5 +1,5 @@
 export * from "./Checkpoints.js"
 export * from "./Tracing.js"
-export * from "./StreamsSink.mjs"
+export { StreamsSink } from "./StreamsSink.mjs"
 export * from "./Types.js"
 export * from "./FeedSource.mjs"