Add onDiscard callback to Flow.fromPublisher

core/src/main/scala/ox/flow/FlowCompanionReactiveOps.scala

@@ -1,11 +1,11 @@
 package ox.flow
 
 import ox.channels.BufferCapacity
+import ox.channels.Channel
 import ox.channels.ChannelClosed
 import ox.channels.toInt
 import ox.discard
 import ox.forkUnsupervised
-import ox.pipe
 import ox.repeatWhile
 import ox.tapException
 import ox.unsupervised
@@ -18,25 +18,38 @@ import java.util.concurrent.atomic.AtomicReference
 trait FlowCompanionReactiveOps:
   this: Flow.type =>
 
+  /** Same as the overload with `onDiscard`, except that elements received from the publisher, but not emitted, are dropped. */
+  def fromPublisher[T](p: Publisher[T])(using BufferCapacity): Flow[T] = fromPublisher(p, (_: T) => ())
+
   /** Creates a [[Flow]] from a [[Publisher]], that is, which emits the elements received by subscribing to the publisher. A new
     * subscription is created every time this flow is run.
     *
     * The data is passed from a subscription to the flow using a [[ox.channel.Channel]], with a capacity given by the [[BufferCapacity]] in
     * scope. That's also how many elements will be at most requested from the publisher at a time.
     *
+    * If the publisher signals an error, elements received before the error are emitted first, and then the error is rethrown.
+    *
+    * Every element received from the publisher is either emitted, or passed to `onDiscard` exactly once. Elements are discarded when the
+    * flow fails, is interrupted or stops early (e.g. due to `take`), including elements delivered by the publisher after the subscription
+    * is cancelled. Use this to release resources held by the elements (e.g. pooled buffers).
+    *
+    * `onDiscard` may be called from the publisher's threads, and must not throw.
+    *
     * The publisher parameter should implement the JDK 9+ `Flow.Publisher` API. To create a flow from a publisher implementing
     * `com.reactivestreams.Publisher`, use the `flow-reactive-streams` module.
     */
-  def fromPublisher[T](p: Publisher[T])(using BufferCapacity): Flow[T] = usingEmitInline: emit =>
+  def fromPublisher[T](p: Publisher[T], onDiscard: T => Unit)(using BufferCapacity): Flow[T] = usingEmitInline: emit =>
     // using an unsafe scope for efficiency
     unsupervised {
       val channel = BufferCapacity.newChannel[T]
       val capacity = summon[BufferCapacity].toInt
       val demandThreshold = math.ceil(capacity / 2.0).toInt
 
       // used to "extract" the subscription that is set in the subscription running in a fork
-      val subscriptionRef = new AtomicReference[Subscription]()
+      val state = new AtomicReference[SubscriptionState](SubscriptionState.NotSubscribed)
       var subscription: Subscription = null
+      // the error is stored, instead of erroring the channel, so that buffered elements are still received
+      val publisherError = new AtomicReference[Option[Throwable]](None)
 
       var toDemand = 0
 
@@ -45,24 +58,38 @@ trait FlowCompanionReactiveOps:
         forkUnsupervised {
           p.subscribe(new Subscriber[T]:
             def onSubscribe(s: Subscription): Unit =
-              subscriptionRef.set(s)
-              s.request(capacity)
+              val previous = state.getAndUpdate:
+                case SubscriptionState.Cancelled => SubscriptionState.Cancelled
+                case _                           => SubscriptionState.Subscribed(s)
+              // the flow might have been cancelled before the subscription was set
+              if previous == SubscriptionState.Cancelled then s.cancel() else s.request(capacity)
 
-            def onNext(t: T): Unit = channel.send(t)
-            def onError(t: Throwable): Unit = channel.error(t)
-            def onComplete(): Unit = channel.done())
+            def onNext(t: T): Unit = channel.sendOrClosed(t) match
+              case _: ChannelClosed => onDiscard(t)
+              case _                => ()
+
+            def onError(t: Throwable): Unit =
+              publisherError.set(Some(t))
+              channel.doneOrClosed().discard
+
+            def onComplete(): Unit = channel.doneOrClosed().discard)
         }.discard
 
         repeatWhile:
           val t = channel.receiveOrClosed()
           t match
-            case ChannelClosed.Done     => false
+            case ChannelClosed.Done =>
+              publisherError.get().foreach(e => throw e)
+              false
             case e: ChannelClosed.Error => throw e.toThrowable
             case t: T @unchecked        =>
               emit(t)
 
               // if we have an element, onSubscribe must have already happened; we can read the subscription and cache it for later
-              if subscription == null then subscription = subscriptionRef.get()
+              if subscription == null then
+                subscription = state.get() match
+                  case SubscriptionState.Subscribed(s) => s
+                  case other                           => throw new IllegalStateException(s"Unexpected subscription state: $other")
 
               // now that we'ver received an element from the channel, we can request more
               toDemand += 1
@@ -73,7 +100,34 @@ trait FlowCompanionReactiveOps:
 
               true
           end match
-        // exceptions might be propagated from the channel, but they might also originate from an interruption
-      }.tapException(_ => subscriptionRef.get().pipe(s => if s != null then s.cancel()))
+        // exceptions might be propagated from the publisher, from emit, but they might also originate from an interruption
+      }.tapException: _ =>
+        state.getAndSet(SubscriptionState.Cancelled) match
+          case SubscriptionState.Subscribed(s) => s.cancel()
+          case _                               => ()
+        channel.doneOrClosed().discard
+        discardBuffered(channel, onDiscard)
     }
+  end fromPublisher
 end FlowCompanionReactiveOps
+
+private enum SubscriptionState:
+  case NotSubscribed
+  case Subscribed(s: Subscription)
+  case Cancelled
+
+/** Passes all elements remaining in a done `channel` to `onDiscard`. Interruptions are deferred until the channel is drained, so that no
+  * element is lost.
+  */
+private def discardBuffered[T](channel: Channel[T], onDiscard: T => Unit): Unit =
+  var interrupted = false
+  var draining = true
+  while draining do
+    try
+      channel.receiveOrClosed() match
+        case _: ChannelClosed => draining = false
+        case t: T @unchecked  => onDiscard(t)
+    catch case _: InterruptedException => interrupted = true
+  end while
+  if interrupted then Thread.currentThread().interrupt()
+end discardBuffered

core/src/test/scala/ox/flow/reactive/FlowFromPublisherDiscardStressTest.scala

@@ -0,0 +1,95 @@
+package ox.flow.reactive
+
+import org.scalatest.funsuite.AnyFunSuite
+import org.scalatest.matchers.should.Matchers
+import ox.*
+import ox.channels.BufferCapacity
+import ox.flow.Flow
+
+import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.concurrent.Flow.Publisher
+import java.util.concurrent.Flow.Subscriber
+import java.util.concurrent.Flow.Subscription
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicLong
+import scala.collection.mutable.ArrayBuffer
+import scala.concurrent.duration.*
+import scala.jdk.CollectionConverters.*
+import scala.util.Random
+import scala.util.Try
+
+class FlowFromPublisherDiscardStressTest extends AnyFunSuite with Matchers:
+  private given BufferCapacity = BufferCapacity(4)
+  private val N = 1000
+  private val Iterations = 20
+
+  private def checkExactlyOnce(consume: Flow[Int] => Unit): Unit =
+    val publisher = new ThreadedPublisher(N)
+    val emitted = new ConcurrentLinkedQueue[Int]()
+    val discarded = new ConcurrentLinkedQueue[Int]()
+
+    consume(Flow.fromPublisher(publisher, discarded.add(_).discard).tap(emitted.add(_).discard))
+    val sent = publisher.join()
+
+    val handled = emitted.asScala.toList ++ discarded.asScala.toList
+    // `sent` has no duplicates, so this also catches elements emitted or discarded twice
+    withClue("duplicated:")(handled.diff(sent) shouldBe empty)
+    withClue("lost:")(sent.diff(handled) shouldBe empty)
+  end checkExactlyOnce
+
+  test("should emit or discard each element exactly once when stopping early"):
+    for _ <- 1 to Iterations do
+      val k = Random.nextInt(N) + 1
+      withClue(s"k = $k"):
+        checkExactlyOnce(_.take(k).runDrain())
+
+  test("should emit or discard each element exactly once when the consumer throws"):
+    for _ <- 1 to Iterations do
+      val failAt = Random.nextInt(N) + 1
+      withClue(s"failAt = $failAt"):
+        checkExactlyOnce(f => Try(f.map(x => if x == failAt then throw new RuntimeException("boom") else x).runDrain()).discard)
+
+  test("should emit or discard each element exactly once when the consumer is interrupted"):
+    for _ <- 1 to Iterations do
+      val delay = Random.nextInt(2000).micros
+      withClue(s"delay = $delay"):
+        checkExactlyOnce: f =>
+          supervised:
+            val result = forkCancellable(f.runDrain())
+            sleep(delay)
+            result.cancel().discard
+end FlowFromPublisherDiscardStressTest
+
+/** Emits `1..n` from its own thread, respecting demand, then completes. After observing cancellation, sends a few more elements, as
+  * allowed by the spec.
+  */
+private class ThreadedPublisher(n: Int) extends Publisher[Int]:
+  private val requested = new AtomicLong(0)
+  private val cancelled = new AtomicBoolean(false)
+  // only accessed by the publisher thread; read after it's joined
+  private val sent = ArrayBuffer[Int]()
+  @volatile private var thread: Thread = null
+
+  override def subscribe(s: Subscriber[? >: Int]): Unit =
+    thread = new Thread(() =>
+      var next = 1
+      def send(): Unit =
+        s.onNext(next)
+        sent += next
+        next += 1
+
+      while next <= n && !cancelled.get() do
+        if next <= requested.get() then send() else Thread.onSpinWait()
+      if cancelled.get() then (1 to 3).foreach(_ => send()) else s.onComplete()
+    )
+    s.onSubscribe(new Subscription:
+      override def request(k: Long): Unit = requested.addAndGet(k).discard
+      override def cancel(): Unit = cancelled.set(true))
+    thread.start()
+  end subscribe
+
+  /** Waits for the publisher thread to finish, and returns the elements that were sent. Must be called after the flow completes. */
+  def join(): List[Int] =
+    if thread != null then thread.join()
+    sent.toList
+end ThreadedPublisher

core/src/test/scala/ox/flow/reactive/FlowFromPublisherDiscardTest.scala

@@ -0,0 +1,106 @@
+package ox.flow.reactive
+
+import org.scalatest.funsuite.AnyFunSuite
+import org.scalatest.matchers.should.Matchers
+import ox.*
+import ox.flow.Flow
+
+import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.concurrent.Semaphore
+import scala.jdk.CollectionConverters.*
+import scala.util.Failure
+import scala.util.Try
+
+class FlowFromPublisherDiscardTest extends AnyFunSuite with Matchers:
+  private val boom = new RuntimeException("boom")
+
+  private class Fixture:
+    val publisher = new ManualPublisher[Int]
+    val emitted = new ConcurrentLinkedQueue[Int]()
+    val discarded = new ConcurrentLinkedQueue[Int]()
+    val flow: Flow[Int] = Flow.fromPublisher(publisher, discarded.add(_).discard).tap(emitted.add(_).discard)
+
+    def emittedList: List[Int] = emitted.asScala.toList
+    def discardedSorted: List[Int] = discarded.asScala.toList.sorted
+  end Fixture
+
+  test("should emit buffered elements before rethrowing a publisher error"):
+    val f = Fixture()
+    val received = new Semaphore(0)
+    val gate = new Semaphore(0)
+    supervised:
+      val result = forkUnsupervised:
+        Try:
+          f.flow.runForeach: _ =>
+            received.release()
+            gate.acquire()
+      f.publisher.subscribed.await()
+      f.publisher.next(1)
+      // the consumer is blocked, so 2 and 3 stay buffered when the error arrives
+      received.acquire()
+      f.publisher.next(2)
+      f.publisher.next(3)
+      f.publisher.error(boom)
+      gate.release(3)
+      result.join() shouldBe Failure(boom)
+    f.emittedList shouldBe List(1, 2, 3)
+    f.discarded shouldBe empty
+
+  test("should discard undelivered elements when the consumer throws"):
+    val f = Fixture()
+    supervised:
+      val result = forkUnsupervised(Try(f.flow.map(x => if x == 2 then throw boom else x).runToList()))
+      f.publisher.subscribed.await()
+      (1 to 5).foreach(f.publisher.next)
+      result.join() shouldBe Failure(boom)
+    f.emittedList shouldBe List(1, 2)
+    f.discardedSorted shouldBe List(3, 4, 5)
+    f.publisher.cancelled.get() shouldBe true
+
+  test("should discard undelivered elements when the consumer stops early"):
+    val f = Fixture()
+    supervised:
+      val result = fork(f.flow.take(2).runToList())
+      f.publisher.subscribed.await()
+      (1 to 5).foreach(f.publisher.next)
+      result.join() shouldBe List(1, 2)
+    f.discardedSorted shouldBe List(3, 4, 5)
+    f.publisher.cancelled.get() shouldBe true
+
+  test("should discard undelivered elements when the consumer is interrupted"):
+    val f = Fixture()
+    val received = new Semaphore(0)
+    val gate = new Semaphore(0)
+    supervised:
+      val result = forkCancellable:
+        f.flow.runForeach: _ =>
+          received.release()
+          gate.acquire()
+      f.publisher.subscribed.await()
+      (1 to 5).foreach(f.publisher.next)
+      received.acquire()
+      result.cancel().discard
+    f.emittedList shouldBe List(1)
+    f.discardedSorted shouldBe List(2, 3, 4, 5)
+    f.publisher.cancelled.get() shouldBe true
+
+  test("should discard elements delivered after cancellation"):
+    val f = Fixture()
+    supervised:
+      val result = fork(f.flow.take(1).runToList())
+      f.publisher.subscribed.await()
+      f.publisher.next(1)
+      result.join() shouldBe List(1)
+    f.publisher.next(2)
+    f.discardedSorted shouldBe List(2)
+
+  test("should cancel a subscription that arrives after the flow is cancelled"):
+    val publisher = new ManualPublisher[Int](deferOnSubscribe = true)
+    supervised:
+      val result = forkCancellable(Flow.fromPublisher(publisher).runDrain())
+      publisher.subscribed.await()
+      result.cancel().discard
+    publisher.sendOnSubscribe()
+    publisher.cancelled.get() shouldBe true
+    publisher.requested.get() shouldBe 0
+end FlowFromPublisherDiscardTest

core/src/test/scala/ox/flow/reactive/ManualPublisher.scala

@@ -0,0 +1,36 @@
+package ox.flow.reactive
+
+import ox.discard
+
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.Flow.Publisher
+import java.util.concurrent.Flow.Subscriber
+import java.util.concurrent.Flow.Subscription
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicLong
+
+/** A publisher driven from the test thread. Ignores demand; supports a single subscriber. With `deferOnSubscribe`, `onSubscribe` is only
+  * called by [[sendOnSubscribe]].
+  */
+class ManualPublisher[T](deferOnSubscribe: Boolean = false) extends Publisher[T]:
+  @volatile var subscriber: Subscriber[? >: T] = null
+  val requested = new AtomicLong(0)
+  val cancelled = new AtomicBoolean(false)
+  val subscribed = new CountDownLatch(1)
+
+  private val subscription = new Subscription:
+    override def request(n: Long): Unit = requested.addAndGet(n).discard
+    override def cancel(): Unit = cancelled.set(true)
+
+  override def subscribe(s: Subscriber[? >: T]): Unit =
+    subscriber = s
+    if !deferOnSubscribe then sendOnSubscribe()
+    subscribed.countDown()
+
+  def sendOnSubscribe(): Unit = subscriber.onSubscribe(subscription)
+
+  def next(t: T): Unit = subscriber.onNext(t)
+  def error(e: Throwable): Unit = subscriber.onError(e)
+  def complete(): Unit = subscriber.onComplete()
+
+end ManualPublisher

doc/streaming/flows.md

@@ -215,5 +215,13 @@ A `java.util.concurrent.Flow.Publisher` can be converted to a `Flow` using `Flow
 Internally, elements published to the subscription are buffered, using a buffer of capacity given by the 
 `BufferCapacity` in scope. That's also how many elements will be at most requested from the publisher at a time.
 
+If the publisher fails, elements already received are emitted before the error is propagated.
+
+Some elements might be received from the publisher, but never emitted downstream: e.g. when the consumer stops early 
+(through an exception, interruption or `take`), or when they arrive after cancellation. To release such elements 
+(e.g. reference-counted buffers), pass an `onDiscard: T => Unit` callback as the second argument to 
+`Flow.fromPublisher` / `FlowReactiveStreams.fromPublisher`. It is called exactly once for each such element, possibly 
+from the publisher's thread, and should not throw.
+
 To convert a `org.reactivestreams.Publisher` instance, you'll need the same dependency as above and call the
 `FlowReactiveStreams.fromPublisher` method.