fix: harden parMap cancellation and async boundary liveness

.github/workflows/ci.yml

@@ -0,0 +1,30 @@
+name: CI
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up JDK 17
+        uses: actions/setup-java@v4
+        with:
+          distribution: temurin
+          java-version: '17'
+          cache: sbt
+
+      - name: Set up sbt
+        uses: sbt/setup-sbt@v1
+
+      - name: Run unit tests
+        run: sbt -Dsbt.supershell=false test
+
+      - name: Run backpressure stress smoke
+        run: sbt -Dsbt.supershell=false "Test / runMain SimpleStreamProcessor.BackpressureStressValidation 5000"

src/main/scala/SimpleStreamProcessor/Node.scala

@@ -4,6 +4,9 @@ import SimpleStreamProcessor.Stream.{QueueEnd, QueueError, QueueSignal, QueueVal
 
 import java.util.concurrent.CancellationException
 import java.util.concurrent.ArrayBlockingQueue
+import java.util.concurrent.atomic.AtomicLong
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.TimeUnit
 import scala.collection.mutable
 import scala.concurrent.ExecutionContext
 import scala.annotation.tailrec
@@ -21,6 +24,9 @@ sealed trait Node[I, O] {
 
   def recover(f: PartialFunction[Throwable, O]): Node[I, O] = RecoverPipe(this, f).withName(this.nodeName + ".recover")
 
+  def recoverWith(f: PartialFunction[Throwable, Stream[O]]): Node[I, O] =
+    RecoverWithPipe(this, f).withName(this.nodeName + ".recoverWith")
+
   def parMap[O2](parallelism: Int)(f: O => O2)(implicit executionContext: ExecutionContext): Node[I, O2] =
     ParMapPipe(this, parallelism, f, executionContext).withName(this.nodeName + ".parMap")
 
@@ -100,20 +106,33 @@ case class Source[I](stream: Stream[I]) extends Node[Unit, I] {
 case class ManagedSource[I, R <: AutoCloseable](resourceFactory: () => R, streamFactory: R => Stream[I]) extends Node[Unit, I] {
   def run(input: Stream[Unit]): Stream[I] = {
     val resource = resourceFactory()
+    val closed = new AtomicBoolean(false)
+
+    def closeResourceOnce(): Unit = {
+      if (closed.compareAndSet(false, true)) {
+        try resource.close()
+        catch {
+          case closeError: Throwable =>
+            Metrics.incResourceCloseFailure()
+            throw closeError
+        }
+      }
+    }
+
     try {
       val stream = streamFactory(resource)
       val cancellableStream = RuntimeControl.currentToken match {
         case Some(token) => stream.takeUntilCancelled(token)
         case None => stream
       }
-      cancellableStream.ensuring(() => resource.close())
+      cancellableStream.ensuring(() => closeResourceOnce())
     } catch {
       case e: Throwable =>
-        try resource.close()
+        try closeResourceOnce()
         catch {
-          case closeError: Throwable =>
-            Metrics.incResourceCloseFailure()
+          case closeError: Throwable if closeError ne e =>
             e.addSuppressed(closeError)
+          case _: Throwable =>
         }
         Stream.Error(e)
     }
@@ -146,6 +165,12 @@ case class RecoverPipe[I, O](upstream: Node[I, O], f: PartialFunction[Throwable,
   override def toString: String = super.toString + "(" + upstream + ")"
 }
 
+case class RecoverWithPipe[I, O](upstream: Node[I, O], f: PartialFunction[Throwable, Stream[O]]) extends Node[I, O] {
+  def run(input: Stream[I]): Stream[O] = upstream.run(input).recoverWith(f)
+
+  override def toString: String = super.toString + "(" + upstream + ")"
+}
+
 case class ParMapPipe[I, O, O2](
   upstream: Node[I, O],
   parallelism: Int,
@@ -163,20 +188,34 @@ case class AsyncBoundaryPipe[I, O](upstream: Node[I, O], bufferSize: Int) extend
 
     val queue = new ArrayBlockingQueue[QueueSignal[O]](bufferSize)
     val cancellationToken = RuntimeControl.currentToken
+    val lastConsumerSignalNs = new AtomicLong(System.nanoTime())
+    val stalledConsumerNs = TimeUnit.SECONDS.toNanos(2)
+
+    def putOrAbort(signal: QueueSignal[O]): Unit = {
+      var offered = false
+      while (!offered) {
+        if (cancellationToken.exists(_.isCancelled)) throw new CancellationException("Pipeline cancelled")
+
+        offered = queue.offer(signal, 100, TimeUnit.MILLISECONDS)
+        if (offered) {
+          Metrics.setBoundaryQueueDepth(queue.size())
+        } else {
+          Metrics.addBoundaryProducerBlockedMs(100)
+          val stalledNs = System.nanoTime() - lastConsumerSignalNs.get()
+          if (stalledNs >= stalledConsumerNs) {
+            throw new IllegalStateException("Async boundary consumer stalled")
+          }
+        }
+      }
+    }
 
     val producer = new Thread(() => {
       cancellationToken.foreach(_.registerCurrentThread())
       try {
         upstream.run(input).foreach { o =>
-          if (cancellationToken.exists(_.isCancelled)) throw new CancellationException("Pipeline cancelled")
-          val startedAtNs = System.nanoTime()
-          queue.put(QueueValue(o))
-          val blockedMs = (System.nanoTime() - startedAtNs) / 1000000
-          Metrics.addBoundaryProducerBlockedMs(blockedMs)
-          Metrics.setBoundaryQueueDepth(queue.size())
+          putOrAbort(QueueValue(o))
         }
-        queue.put(QueueEnd)
-        Metrics.setBoundaryQueueDepth(queue.size())
+        putOrAbort(QueueEnd)
       } catch {
         case _: InterruptedException if cancellationToken.exists(_.isCancelled) =>
           queue.offer(QueueEnd)
@@ -185,7 +224,7 @@ case class AsyncBoundaryPipe[I, O](upstream: Node[I, O], bufferSize: Int) extend
           queue.offer(QueueEnd)
           Metrics.setBoundaryQueueDepth(queue.size())
         case e: Throwable =>
-          queue.put(QueueError(e))
+          queue.offer(QueueError(e))
           Metrics.setBoundaryQueueDepth(queue.size())
       } finally {
         cancellationToken.foreach(_.unregisterCurrentThread())
@@ -196,7 +235,7 @@ case class AsyncBoundaryPipe[I, O](upstream: Node[I, O], bufferSize: Int) extend
     producer.setDaemon(true)
     producer.start()
 
-    Stream.fromBlockingQueue(queue)
+    Stream.fromBlockingQueue(queue, (_: QueueSignal[O]) => lastConsumerSignalNs.set(System.nanoTime()))
   }
 
   override def toString: String = super.toString + "(" + upstream + ")"

src/main/scala/SimpleStreamProcessor/Stream.scala

@@ -2,8 +2,8 @@ package SimpleStreamProcessor
 
 import java.util.concurrent.BlockingQueue
 import java.util.concurrent.atomic.AtomicBoolean
-import scala.concurrent.duration.Duration
-import scala.concurrent.{Await, ExecutionContext, Future}
+import java.util.concurrent.{Callable, ExecutorCompletionService, Executors, TimeUnit}
+import scala.concurrent.ExecutionContext
 
 sealed trait Stream[+A] {
 
@@ -114,11 +114,53 @@ sealed trait Stream[+A] {
     }
 
     def runBatch(batch: List[A]): List[B] = {
-      val batchFutures = batch.map { a =>
-        Metrics.incParMapInFlight()
-        Future(f(a)).andThen { case _ => Metrics.decParMapInFlight() }
+      val cancellationToken = RuntimeControl.currentToken
+      val executor = Executors.newFixedThreadPool(parallelism)
+      val completion = new ExecutorCompletionService[(Int, Either[Throwable, B])](executor)
+      val results = Array.fill[Option[B]](batch.size)(None)
+
+      try {
+        batch.zipWithIndex.foreach {
+          case (value, index) =>
+            completion.submit(new Callable[(Int, Either[Throwable, B])] {
+              override def call(): (Int, Either[Throwable, B]) = {
+                if (cancellationToken.exists(_.isCancelled)) {
+                  (index, Left(new java.util.concurrent.CancellationException("Pipeline cancelled")))
+                } else {
+                  Metrics.incParMapInFlight()
+                  try {
+                    (index, Right(f(value)))
+                  } catch {
+                    case e: Throwable => (index, Left(e))
+                  } finally {
+                    Metrics.decParMapInFlight()
+                  }
+                }
+              }
+            })
+        }
+
+        var completed = 0
+        while (completed < batch.size) {
+          if (cancellationToken.exists(_.isCancelled)) {
+            throw new java.util.concurrent.CancellationException("Pipeline cancelled")
+          }
+
+          val future = completion.poll(100, TimeUnit.MILLISECONDS)
+          if (future != null) {
+            val (index, result) = future.get()
+            result match {
+              case Right(mapped) => results(index) = Some(mapped)
+              case Left(error) => throw error
+            }
+            completed += 1
+          }
+        }
+
+        results.iterator.collect { case Some(value) => value }.toList
+      } finally {
+        executor.shutdownNow()
       }
-      Await.result(Future.sequence(batchFutures), Duration.Inf)
     }
 
     def loop(s: Stream[A]): Stream[B] = s match {
@@ -145,8 +187,15 @@ sealed trait Stream[+A] {
   def ensuring(finalizer: () => Unit): Stream[A] = {
     val closed = new AtomicBoolean(false)
 
-    def closeOnce(): Unit = {
-      if (closed.compareAndSet(false, true)) finalizer()
+    def closeOnce(): Option[Throwable] = {
+      if (closed.compareAndSet(false, true)) {
+        try {
+          finalizer()
+          None
+        } catch {
+          case e: Throwable => Some(e)
+        }
+      } else None
     }
 
     def go(s: Stream[A]): Stream[A] = s match {
@@ -155,18 +204,22 @@ sealed trait Stream[+A] {
           try go(next())
           catch {
             case e: Throwable =>
-              closeOnce()
+              closeOnce().foreach(e.addSuppressed)
               Error(e)
           }
         })
       case Halt() =>
-        closeOnce()
-        Halt()
+        closeOnce() match {
+          case Some(closeError) => Error(closeError)
+          case None => Halt()
+        }
       case Empty =>
-        closeOnce()
-        Empty
+        closeOnce() match {
+          case Some(closeError) => Error(closeError)
+          case None => Empty
+        }
       case Error(e) =>
-        closeOnce()
+        closeOnce().foreach(e.addSuppressed)
         Error(e)
     }
 
@@ -273,12 +326,14 @@ object Stream {
     else Emit(queue.poll(), () => fromQueue(queue))
   }
 
-  def fromBlockingQueue[A](queue: BlockingQueue[QueueSignal[A]]): Stream[A] = {
+  def fromBlockingQueue[A](queue: BlockingQueue[QueueSignal[A]], onSignal: QueueSignal[A] => Unit = (_: QueueSignal[A]) => ()): Stream[A] = {
     try {
-      queue.take() match {
+      val signal = queue.take()
+      onSignal(signal)
+      signal match {
         case QueueValue(value) =>
           Metrics.setBoundaryQueueDepth(queue.size())
-          Emit(value, () => fromBlockingQueue(queue))
+          Emit(value, () => fromBlockingQueue(queue, onSignal))
         case QueueEnd =>
           Metrics.setBoundaryQueueDepth(queue.size())
           Halt()