From 56772fcfb148e39da26f4d6d2ce9f1c12210b0a5 Mon Sep 17 00:00:00 2001 From: John Gerassimou Date: Thu, 26 Feb 2026 00:48:28 -0500 Subject: [PATCH 1/4] fix: harden parMap cancellation and async boundary liveness --- .../scala/SimpleStreamProcessor/Node.scala | 36 +++++--- .../scala/SimpleStreamProcessor/Stream.scala | 89 +++++++++++++++---- .../SimpleStreamProcessorTest.scala | 45 ++++++++++ 3 files changed, 143 insertions(+), 27 deletions(-) diff --git a/src/main/scala/SimpleStreamProcessor/Node.scala b/src/main/scala/SimpleStreamProcessor/Node.scala index 340d99d..e8e779d 100644 --- a/src/main/scala/SimpleStreamProcessor/Node.scala +++ b/src/main/scala/SimpleStreamProcessor/Node.scala @@ -4,6 +4,8 @@ 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.TimeUnit import scala.collection.mutable import scala.concurrent.ExecutionContext import scala.annotation.tailrec @@ -163,20 +165,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 +201,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 +212,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 + ")" diff --git a/src/main/scala/SimpleStreamProcessor/Stream.scala b/src/main/scala/SimpleStreamProcessor/Stream.scala index f7b3d3a..496db56 100644 --- a/src/main/scala/SimpleStreamProcessor/Stream.scala +++ b/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() diff --git a/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala b/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala index b95b754..ed3112e 100644 --- a/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala +++ b/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala @@ -8,6 +8,7 @@ import scala.concurrent.ExecutionContext import scala.concurrent.duration.DurationInt import scala.concurrent.Await import java.util.concurrent.atomic.AtomicInteger +import scala.jdk.CollectionConverters._ class SimpleStreamProcessorTest extends AnyFunSuite with BeforeAndAfterEach { @@ -420,4 +421,48 @@ class SimpleStreamProcessorTest extends AnyFunSuite with BeforeAndAfterEach { assert(processed.get() < 2000) } + test("parMap cancellation stops async execution") { + implicit val executionContext: ExecutionContext = ExecutionContext.global + + val node = Source[Int](Stream.fromList((1 to 5000).toList)) + .parMap(4) { i => + Thread.sleep(2) + i + } + + val handle = node.runToListAsync(Stream.Empty) + Thread.sleep(10) + handle.cancel() + + val outcome = Await.result(handle.outcome, 3.seconds) + assert(outcome == ExecutionCancelled) + } + + test("async boundary producer thread exits after downstream failure") { + val before = Thread.getAllStackTraces.keySet().asScala + .filter(_.getName.startsWith("simple-stream-async-boundary-")) + .map(_.getId) + .toSet + + val stream = Source[Int](Stream.fromList((1 to 10000).toList)) + .withName("boundary-leak-check") + .asyncBoundary(1) + .map { i => + if (i == 3) throw new RuntimeException("stop") + i + } + .run(Stream.Empty) + + intercept[RuntimeException](stream.toList) + + Thread.sleep(2500) + + val afterAlive = Thread.getAllStackTraces.keySet().asScala + .filter(t => t.getName.startsWith("simple-stream-async-boundary-") && t.isAlive) + .map(_.getId) + .toSet + + assert(afterAlive.diff(before).isEmpty) + } + } From c5644da5aa23f286f23797a3c2f6aa49694beea9 Mon Sep 17 00:00:00 2001 From: John Gerassimou Date: Thu, 26 Feb 2026 00:50:33 -0500 Subject: [PATCH 2/4] test: enforce managed resource failure precedence semantics --- .../scala/SimpleStreamProcessor/Node.scala | 22 +++++++-- .../SimpleStreamProcessorTest.scala | 45 +++++++++++++++++++ 2 files changed, 63 insertions(+), 4 deletions(-) diff --git a/src/main/scala/SimpleStreamProcessor/Node.scala b/src/main/scala/SimpleStreamProcessor/Node.scala index e8e779d..7df3901 100644 --- a/src/main/scala/SimpleStreamProcessor/Node.scala +++ b/src/main/scala/SimpleStreamProcessor/Node.scala @@ -5,6 +5,7 @@ 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 @@ -102,20 +103,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) } diff --git a/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala b/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala index ed3112e..3fb834b 100644 --- a/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala +++ b/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala @@ -209,6 +209,51 @@ class SimpleStreamProcessorTest extends AnyFunSuite with BeforeAndAfterEach { assert(captured.closed) } + test("Managed sink preserves processing failure and suppresses close failure") { + class BrokenResource extends AutoCloseable { + override def close(): Unit = throw new RuntimeException("close-failed") + } + + val sink = Source[Int](Stream.fromList(List(1, 0, 2))) + .map(i => 10 / i) + .toManagedSink(() => new BrokenResource)((_, _) => ()) + + val error = intercept[ArithmeticException](sink.run(Stream.Empty)) + assert(error.getSuppressed.exists(_.getMessage == "close-failed")) + assert(Metrics.snapshot().resourceCloseFailTotal == 1) + } + + test("Managed source close failure surfaces when processing succeeds") { + class BrokenResource extends AutoCloseable { + override def close(): Unit = throw new RuntimeException("close-failed") + } + + val source = ManagedSource[Int, BrokenResource]( + resourceFactory = () => new BrokenResource, + streamFactory = _ => Stream.fromList(List(1, 2, 3)) + ) + + val error = intercept[RuntimeException](source.run(Stream.Empty).toList) + assert(error.getMessage == "close-failed") + assert(Metrics.snapshot().resourceCloseFailTotal == 1) + } + + test("Managed source preserves processing failure and suppresses close failure") { + class BrokenResource extends AutoCloseable { + override def close(): Unit = throw new RuntimeException("close-failed") + } + + val source = ManagedSource[Int, BrokenResource]( + resourceFactory = () => new BrokenResource, + streamFactory = _ => Stream.Error(new RuntimeException("processing-failed")) + ) + + val error = intercept[RuntimeException](source.run(Stream.Empty).toList) + assert(error.getMessage == "processing-failed") + assert(error.getSuppressed.exists(_.getMessage == "close-failed")) + assert(Metrics.snapshot().resourceCloseFailTotal == 1) + } + test("Count windows split stream into fixed-size batches") { val result = Source[Int](Stream.fromList((1 to 7).toList)) .windowByCount(3) From a3d45ecca49a7b3a6ad0a234f4675546ae7b52e6 Mon Sep 17 00:00:00 2001 From: John Gerassimou Date: Thu, 26 Feb 2026 00:52:39 -0500 Subject: [PATCH 3/4] ci: add JDK17 test workflow and node recoverWith coverage --- .github/workflows/ci.yml | 27 +++++++++++++++++++ .../scala/SimpleStreamProcessor/Node.scala | 9 +++++++ .../SimpleStreamProcessorTest.scala | 9 +++++++ 3 files changed, 45 insertions(+) create mode 100644 .github/workflows/ci.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..d06aca1 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,27 @@ +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: Run unit tests + run: sbt -Dsbt.supershell=false test + + - name: Run backpressure stress smoke + run: sbt -Dsbt.supershell=false "Test / runMain SimpleStreamProcessor.BackpressureStressValidation 5000" diff --git a/src/main/scala/SimpleStreamProcessor/Node.scala b/src/main/scala/SimpleStreamProcessor/Node.scala index 7df3901..6fd9b4f 100644 --- a/src/main/scala/SimpleStreamProcessor/Node.scala +++ b/src/main/scala/SimpleStreamProcessor/Node.scala @@ -24,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") @@ -162,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, diff --git a/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala b/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala index 3fb834b..a028dfd 100644 --- a/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala +++ b/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala @@ -68,6 +68,15 @@ class SimpleStreamProcessorTest extends AnyFunSuite with BeforeAndAfterEach { assert(sink.run(Stream.Empty) == 10) } + test("Node recoverWith allows stream fallback") { + val sink = Source[Int](Stream.fromList(List(1, 0, 2))) + .map(i => 10 / i) + .recoverWith { case _: ArithmeticException => Stream.fromList(List(99, 100)) } + .toSink((acc: Int, i: Int) => acc + i, 0) + + assert(sink.run(Stream.Empty) == 209) + } + test("Stream parMap preserves input order") { implicit val executionContext: ExecutionContext = ExecutionContext.global From cf33b48d8b28cc41d4fdb5bfa2974de80cae50da Mon Sep 17 00:00:00 2001 From: John Gerassimou Date: Thu, 26 Feb 2026 00:54:49 -0500 Subject: [PATCH 4/4] ci: install sbt in GitHub Actions workflow --- .github/workflows/ci.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d06aca1..b4b05bd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,6 +20,9 @@ jobs: java-version: '17' cache: sbt + - name: Set up sbt + uses: sbt/setup-sbt@v1 + - name: Run unit tests run: sbt -Dsbt.supershell=false test