feat: add scoped metrics and cancellable iterator execution

.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"

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())
   }
 }

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()
   }
 }

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())
       }
     })
 

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()
                   }
                 }
               }