From ec5c8653eb83cc61c47ce2ab363b009126e57cb0 Mon Sep 17 00:00:00 2001 From: John Gerassimou Date: Thu, 26 Feb 2026 01:02:00 -0500 Subject: [PATCH] feat: add scoped metrics and cancellable iterator execution --- .github/workflows/ci.yml | 4 +- .../SimpleStreamProcessor/Execution.scala | 40 +++-- .../scala/SimpleStreamProcessor/Metrics.scala | 142 +++++++++++++----- .../scala/SimpleStreamProcessor/Node.scala | 64 +++++--- .../scala/SimpleStreamProcessor/Stream.scala | 5 +- .../SimpleStreamProcessorTest.scala | 40 +++++ 6 files changed, 225 insertions(+), 70 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b4b05bd..fe61727 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,5 +26,5 @@ jobs: - 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" + - name: Run backpressure stress assertions + run: sbt -Dsbt.supershell=false "Test / runMain SimpleStreamProcessor.BackpressureStressValidation 10000" diff --git a/src/main/scala/SimpleStreamProcessor/Execution.scala b/src/main/scala/SimpleStreamProcessor/Execution.scala index 4e6ef2e..0589d9c 100644 --- a/src/main/scala/SimpleStreamProcessor/Execution.scala +++ b/src/main/scala/SimpleStreamProcessor/Execution.scala @@ -10,7 +10,18 @@ case class ExecutionCompleted[A](value: A) extends ExecutionOutcome[A] case class ExecutionFailed(error: Throwable) extends ExecutionOutcome[Nothing] case object ExecutionCancelled extends ExecutionOutcome[Nothing] -case class ExecutionHandle[+A](outcome: Future[ExecutionOutcome[A]], cancel: () => Unit) +case class ExecutionHandle[+A]( + outcome: Future[ExecutionOutcome[A]], + cancel: () => Unit, + metricsSnapshot: () => Metrics.Snapshot +) + +case class CancellableIterator[+A]( + iterator: Iterator[A], + cancel: () => Unit, + outcome: Future[ExecutionOutcome[Unit]], + metricsSnapshot: () => Metrics.Snapshot +) final class CancellationToken { private val cancelled = new AtomicBoolean(false) @@ -57,22 +68,25 @@ object RuntimeControl { def runAsync[A](compute: CancellationToken => A)(implicit executionContext: ExecutionContext): ExecutionHandle[A] = { val token = new CancellationToken + val collector = Metrics.newCollector() val future = Future { - RuntimeControl.withToken(token) { - token.registerCurrentThread() - try { - if (token.isCancelled) ExecutionCancelled - else ExecutionCompleted(compute(token)) - } catch { - case _: CancellationException => ExecutionCancelled - case _: InterruptedException if token.isCancelled => ExecutionCancelled - case e: Throwable => ExecutionFailed(e) - } finally { - token.unregisterCurrentThread() + Metrics.withCollector(collector) { + RuntimeControl.withToken(token) { + token.registerCurrentThread() + try { + if (token.isCancelled) ExecutionCancelled + else ExecutionCompleted(compute(token)) + } catch { + case _: CancellationException => ExecutionCancelled + case _: InterruptedException if token.isCancelled => ExecutionCancelled + case e: Throwable => ExecutionFailed(e) + } finally { + token.unregisterCurrentThread() + } } } } - ExecutionHandle(future, () => token.cancel()) + ExecutionHandle(future, () => token.cancel(), () => collector.snapshot()) } } diff --git a/src/main/scala/SimpleStreamProcessor/Metrics.scala b/src/main/scala/SimpleStreamProcessor/Metrics.scala index 720f160..2dd5d00 100644 --- a/src/main/scala/SimpleStreamProcessor/Metrics.scala +++ b/src/main/scala/SimpleStreamProcessor/Metrics.scala @@ -3,15 +3,6 @@ package SimpleStreamProcessor import java.util.concurrent.atomic.{AtomicInteger, AtomicLong} object Metrics { - private val parMapInFlight = new AtomicInteger(0) - private val boundaryQueueDepth = new AtomicInteger(0) - private val boundaryQueueDepthMax = new AtomicInteger(0) - private val boundaryProducerBlockedMs = new AtomicLong(0) - private val lateEventDroppedTotal = new AtomicLong(0) - private val watermarkRegressionTotal = new AtomicLong(0) - private val resourceCloseFailTotal = new AtomicLong(0) - private val unhandledErrorTotal = new AtomicLong(0) - case class Snapshot( parMapInFlight: Int, boundaryQueueDepth: Int, @@ -23,58 +14,137 @@ object Metrics { unhandledErrorTotal: Long ) + trait Collector { + def reset(): Unit + def snapshot(): Snapshot + def incParMapInFlight(): Unit + def decParMapInFlight(): Unit + def setBoundaryQueueDepth(depth: Int): Unit + def addBoundaryProducerBlockedMs(ms: Long): Unit + def incLateEventDropped(): Unit + def incWatermarkRegression(): Unit + def incResourceCloseFailure(): Unit + def incUnhandledError(): Unit + } + + private final class AtomicCollector extends Collector { + private val parMapInFlight = new AtomicInteger(0) + private val boundaryQueueDepth = new AtomicInteger(0) + private val boundaryQueueDepthMax = new AtomicInteger(0) + private val boundaryProducerBlockedMs = new AtomicLong(0) + private val lateEventDroppedTotal = new AtomicLong(0) + private val watermarkRegressionTotal = new AtomicLong(0) + private val resourceCloseFailTotal = new AtomicLong(0) + private val unhandledErrorTotal = new AtomicLong(0) + + def reset(): Unit = { + parMapInFlight.set(0) + boundaryQueueDepth.set(0) + boundaryQueueDepthMax.set(0) + boundaryProducerBlockedMs.set(0) + lateEventDroppedTotal.set(0) + watermarkRegressionTotal.set(0) + resourceCloseFailTotal.set(0) + unhandledErrorTotal.set(0) + } + + def snapshot(): Snapshot = Snapshot( + parMapInFlight = parMapInFlight.get(), + boundaryQueueDepth = boundaryQueueDepth.get(), + boundaryQueueDepthMax = boundaryQueueDepthMax.get(), + boundaryProducerBlockedMs = boundaryProducerBlockedMs.get(), + lateEventDroppedTotal = lateEventDroppedTotal.get(), + watermarkRegressionTotal = watermarkRegressionTotal.get(), + resourceCloseFailTotal = resourceCloseFailTotal.get(), + unhandledErrorTotal = unhandledErrorTotal.get() + ) + + def incParMapInFlight(): Unit = { + parMapInFlight.incrementAndGet() + } + + def decParMapInFlight(): Unit = { + parMapInFlight.decrementAndGet() + } + + def setBoundaryQueueDepth(depth: Int): Unit = { + boundaryQueueDepth.set(depth) + boundaryQueueDepthMax.getAndUpdate(prev => math.max(prev, depth)) + } + + def addBoundaryProducerBlockedMs(ms: Long): Unit = { + boundaryProducerBlockedMs.addAndGet(ms) + } + + def incLateEventDropped(): Unit = { + lateEventDroppedTotal.incrementAndGet() + } + + def incWatermarkRegression(): Unit = { + watermarkRegressionTotal.incrementAndGet() + } + + def incResourceCloseFailure(): Unit = { + resourceCloseFailTotal.incrementAndGet() + } + + def incUnhandledError(): Unit = { + unhandledErrorTotal.incrementAndGet() + } + } + + private val globalCollector: AtomicCollector = new AtomicCollector + private val scopedCollector = new ThreadLocal[Collector]() + + def newCollector(): Collector = new AtomicCollector + + def currentCollector: Collector = Option(scopedCollector.get()).getOrElse(globalCollector) + + def withCollector[A](collector: Collector)(body: => A): A = { + val previous = scopedCollector.get() + scopedCollector.set(collector) + try body + finally { + if (previous == null) scopedCollector.remove() + else scopedCollector.set(previous) + } + } + def reset(): Unit = { - parMapInFlight.set(0) - boundaryQueueDepth.set(0) - boundaryQueueDepthMax.set(0) - boundaryProducerBlockedMs.set(0) - lateEventDroppedTotal.set(0) - watermarkRegressionTotal.set(0) - resourceCloseFailTotal.set(0) - unhandledErrorTotal.set(0) + globalCollector.reset() } - def snapshot(): Snapshot = Snapshot( - parMapInFlight = parMapInFlight.get(), - boundaryQueueDepth = boundaryQueueDepth.get(), - boundaryQueueDepthMax = boundaryQueueDepthMax.get(), - boundaryProducerBlockedMs = boundaryProducerBlockedMs.get(), - lateEventDroppedTotal = lateEventDroppedTotal.get(), - watermarkRegressionTotal = watermarkRegressionTotal.get(), - resourceCloseFailTotal = resourceCloseFailTotal.get(), - unhandledErrorTotal = unhandledErrorTotal.get() - ) + def snapshot(): Snapshot = globalCollector.snapshot() def incParMapInFlight(): Unit = { - parMapInFlight.incrementAndGet() + currentCollector.incParMapInFlight() } def decParMapInFlight(): Unit = { - parMapInFlight.decrementAndGet() + currentCollector.decParMapInFlight() } def setBoundaryQueueDepth(depth: Int): Unit = { - boundaryQueueDepth.set(depth) - boundaryQueueDepthMax.getAndUpdate(prev => math.max(prev, depth)) + currentCollector.setBoundaryQueueDepth(depth) } def addBoundaryProducerBlockedMs(ms: Long): Unit = { - boundaryProducerBlockedMs.addAndGet(ms) + currentCollector.addBoundaryProducerBlockedMs(ms) } def incLateEventDropped(): Unit = { - lateEventDroppedTotal.incrementAndGet() + currentCollector.incLateEventDropped() } def incWatermarkRegression(): Unit = { - watermarkRegressionTotal.incrementAndGet() + currentCollector.incWatermarkRegression() } def incResourceCloseFailure(): Unit = { - resourceCloseFailTotal.incrementAndGet() + currentCollector.incResourceCloseFailure() } def incUnhandledError(): Unit = { - unhandledErrorTotal.incrementAndGet() + currentCollector.incUnhandledError() } } diff --git a/src/main/scala/SimpleStreamProcessor/Node.scala b/src/main/scala/SimpleStreamProcessor/Node.scala index 6fd9b4f..1a10bd8 100644 --- a/src/main/scala/SimpleStreamProcessor/Node.scala +++ b/src/main/scala/SimpleStreamProcessor/Node.scala @@ -87,6 +87,33 @@ sealed trait Node[I, O] { loop(run(input)) } + def runCancellableIterator(input: Stream[I], bufferSize: Int = 64)(implicit executionContext: ExecutionContext): CancellableIterator[O] = { + val capacity = math.max(1, bufferSize) + val queue = new ArrayBlockingQueue[QueueSignal[O]](capacity) + + val handle = runForeachAsync(input) { value => + queue.put(QueueValue(value)) + } + + def publishTerminal(signal: QueueSignal[O]): Unit = { + queue.clear() + queue.offer(signal) + } + + handle.outcome.foreach { + case ExecutionCompleted(_) => publishTerminal(QueueEnd) + case ExecutionCancelled => publishTerminal(QueueEnd) + case ExecutionFailed(error) => publishTerminal(QueueError(error)) + }(executionContext) + + CancellableIterator( + iterator = Stream.fromBlockingQueue(queue).iterator, + cancel = handle.cancel, + outcome = handle.outcome, + metricsSnapshot = handle.metricsSnapshot + ) + } + def runIterator(input: Stream[I]): Iterator[O] = run(input).iterator def withName(name: String): this.type = { @@ -188,6 +215,7 @@ case class AsyncBoundaryPipe[I, O](upstream: Node[I, O], bufferSize: Int) extend val queue = new ArrayBlockingQueue[QueueSignal[O]](bufferSize) val cancellationToken = RuntimeControl.currentToken + val collector = Metrics.currentCollector val lastConsumerSignalNs = new AtomicLong(System.nanoTime()) val stalledConsumerNs = TimeUnit.SECONDS.toNanos(2) @@ -210,24 +238,26 @@ case class AsyncBoundaryPipe[I, O](upstream: Node[I, O], bufferSize: Int) extend } val producer = new Thread(() => { - cancellationToken.foreach(_.registerCurrentThread()) - try { - upstream.run(input).foreach { o => - putOrAbort(QueueValue(o)) + Metrics.withCollector(collector) { + cancellationToken.foreach(_.registerCurrentThread()) + try { + upstream.run(input).foreach { o => + putOrAbort(QueueValue(o)) + } + putOrAbort(QueueEnd) + } catch { + case _: InterruptedException if cancellationToken.exists(_.isCancelled) => + queue.offer(QueueEnd) + Metrics.setBoundaryQueueDepth(queue.size()) + case _: CancellationException => + queue.offer(QueueEnd) + Metrics.setBoundaryQueueDepth(queue.size()) + case e: Throwable => + queue.offer(QueueError(e)) + Metrics.setBoundaryQueueDepth(queue.size()) + } finally { + cancellationToken.foreach(_.unregisterCurrentThread()) } - putOrAbort(QueueEnd) - } catch { - case _: InterruptedException if cancellationToken.exists(_.isCancelled) => - queue.offer(QueueEnd) - Metrics.setBoundaryQueueDepth(queue.size()) - case _: CancellationException => - queue.offer(QueueEnd) - Metrics.setBoundaryQueueDepth(queue.size()) - case e: Throwable => - queue.offer(QueueError(e)) - Metrics.setBoundaryQueueDepth(queue.size()) - } finally { - cancellationToken.foreach(_.unregisterCurrentThread()) } }) diff --git a/src/main/scala/SimpleStreamProcessor/Stream.scala b/src/main/scala/SimpleStreamProcessor/Stream.scala index 496db56..30bd9cf 100644 --- a/src/main/scala/SimpleStreamProcessor/Stream.scala +++ b/src/main/scala/SimpleStreamProcessor/Stream.scala @@ -115,6 +115,7 @@ sealed trait Stream[+A] { def runBatch(batch: List[A]): List[B] = { val cancellationToken = RuntimeControl.currentToken + val collector = Metrics.currentCollector val executor = Executors.newFixedThreadPool(parallelism) val completion = new ExecutorCompletionService[(Int, Either[Throwable, B])](executor) val results = Array.fill[Option[B]](batch.size)(None) @@ -127,13 +128,13 @@ sealed trait Stream[+A] { if (cancellationToken.exists(_.isCancelled)) { (index, Left(new java.util.concurrent.CancellationException("Pipeline cancelled"))) } else { - Metrics.incParMapInFlight() + collector.incParMapInFlight() try { (index, Right(f(value))) } catch { case e: Throwable => (index, Left(e)) } finally { - Metrics.decParMapInFlight() + collector.decParMapInFlight() } } } diff --git a/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala b/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala index a028dfd..faede7b 100644 --- a/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala +++ b/src/test/scala/SimpleStreamProcessor/SimpleStreamProcessorTest.scala @@ -519,4 +519,44 @@ class SimpleStreamProcessorTest extends AnyFunSuite with BeforeAndAfterEach { assert(afterAlive.diff(before).isEmpty) } + test("runAsync metrics are scoped per execution handle") { + implicit val executionContext: ExecutionContext = ExecutionContext.global + + Metrics.reset() + val timedEvents = List( + Record(Timestamped("a", 1L)), + Watermark(8L), + Record(Timestamped("late", 4L)) + ) + + val node = Source[TimedEvent[String]](Stream.fromList(timedEvents)) + .windowByEventTime(windowSizeMs = 5L) + + val handle = node.runToListAsync(Stream.Empty) + val outcome = Await.result(handle.outcome, 2.seconds) + + assert(outcome.isInstanceOf[ExecutionCompleted[_]]) + assert(Metrics.snapshot().lateEventDroppedTotal == 0) + assert(handle.metricsSnapshot().lateEventDroppedTotal == 1) + } + + test("Node runCancellableIterator supports cancellation") { + implicit val executionContext: ExecutionContext = ExecutionContext.global + + val node = Source[Int](Stream.fromList((1 to 5000).toList)) + .map { i => + Thread.sleep(1) + i + } + + val cancellable = node.runCancellableIterator(Stream.Empty, bufferSize = 8) + + assert(cancellable.iterator.next() == 1) + Thread.sleep(10) + cancellable.cancel() + + val outcome = Await.result(cancellable.outcome, 2.seconds) + assert(outcome == ExecutionCancelled) + } + }