diff --git a/core/src/main/scala/ox/flow/FlowCompanionReactiveOps.scala b/core/src/main/scala/ox/flow/FlowCompanionReactiveOps.scala index 370975fd..9623e69b 100644 --- a/core/src/main/scala/ox/flow/FlowCompanionReactiveOps.scala +++ b/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,16 +18,27 @@ 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] @@ -35,8 +46,10 @@ trait FlowCompanionReactiveOps: 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 diff --git a/core/src/test/scala/ox/flow/reactive/FlowFromPublisherDiscardStressTest.scala b/core/src/test/scala/ox/flow/reactive/FlowFromPublisherDiscardStressTest.scala new file mode 100644 index 00000000..d2d86284 --- /dev/null +++ b/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 diff --git a/core/src/test/scala/ox/flow/reactive/FlowFromPublisherDiscardTest.scala b/core/src/test/scala/ox/flow/reactive/FlowFromPublisherDiscardTest.scala new file mode 100644 index 00000000..cdf4ff27 --- /dev/null +++ b/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 diff --git a/core/src/test/scala/ox/flow/reactive/ManualPublisher.scala b/core/src/test/scala/ox/flow/reactive/ManualPublisher.scala new file mode 100644 index 00000000..351ae0e1 --- /dev/null +++ b/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 diff --git a/doc/streaming/flows.md b/doc/streaming/flows.md index f6a3cacf..8af8a635 100644 --- a/doc/streaming/flows.md +++ b/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. \ No newline at end of file diff --git a/flow-reactive-streams/src/main/scala/ox/flow/reactive/flowReactiveStreamsExtension.scala b/flow-reactive-streams/src/main/scala/ox/flow/reactive/flowReactiveStreamsExtension.scala index e0741f41..df92d607 100644 --- a/flow-reactive-streams/src/main/scala/ox/flow/reactive/flowReactiveStreamsExtension.scala +++ b/flow-reactive-streams/src/main/scala/ox/flow/reactive/flowReactiveStreamsExtension.scala @@ -18,11 +18,21 @@ extension [A](flow: Flow[A]) end extension object FlowReactiveStreams: - /** This variant returns accepts an implementation of `org.reactivestreams.Publisher`, as opposed to `java.util.concurrent.Flow.Publisher` - * which is supported in the core module. + /** This variant accepts an implementation of `org.reactivestreams.Publisher`, as opposed to `java.util.concurrent.Flow.Publisher` which + * is supported in the core module. + * + * To release resources held by elements which are received but not emitted, use the overload with `onDiscard`. * * @see * [[Flow.fromPublisher]] */ def fromPublisher[T](p: Publisher[T])(using BufferCapacity): Flow[T] = Flow.fromPublisher(FlowAdapters.toFlowPublisher(p)) + + /** Same as the overload without `onDiscard`, but passes elements which are received but not emitted to `onDiscard`. + * + * @see + * [[Flow.fromPublisher]] + */ + def fromPublisher[T](p: Publisher[T], onDiscard: T => Unit)(using BufferCapacity): Flow[T] = + Flow.fromPublisher(FlowAdapters.toFlowPublisher(p), onDiscard) end FlowReactiveStreams diff --git a/flow-reactive-streams/src/test/scala/ox/flow/reactive/FlowReactiveStreamsFromPublisherTest.scala b/flow-reactive-streams/src/test/scala/ox/flow/reactive/FlowReactiveStreamsFromPublisherTest.scala new file mode 100644 index 00000000..d57214d0 --- /dev/null +++ b/flow-reactive-streams/src/test/scala/ox/flow/reactive/FlowReactiveStreamsFromPublisherTest.scala @@ -0,0 +1,53 @@ +package ox.flow.reactive + +import org.reactivestreams.Publisher +import org.reactivestreams.Subscriber +import org.reactivestreams.Subscription +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers +import ox.* + +import java.util.concurrent.ConcurrentLinkedQueue +import java.util.concurrent.CountDownLatch +import java.util.concurrent.atomic.AtomicBoolean +import scala.jdk.CollectionConverters.* + +class FlowReactiveStreamsFromPublisherTest extends AnyFunSuite with Matchers: + + /** A publisher driven from the test thread. Ignores demand; supports a single subscriber. */ + private class ManualPublisher[T] extends Publisher[T]: + @volatile var subscriber: Subscriber[? >: T] = null + val cancelled = new AtomicBoolean(false) + val subscribed = new CountDownLatch(1) + + override def subscribe(s: Subscriber[? >: T]): Unit = + subscriber = s + s.onSubscribe(new Subscription: + override def request(n: Long): Unit = () + override def cancel(): Unit = cancelled.set(true)) + subscribed.countDown() + + def next(t: T): Unit = subscriber.onNext(t) + def complete(): Unit = subscriber.onComplete() + end ManualPublisher + + test("should emit elements from the publisher"): + val publisher = new ManualPublisher[Int] + supervised: + val result = fork(FlowReactiveStreams.fromPublisher(publisher).runToList()) + publisher.subscribed.await() + (1 to 3).foreach(publisher.next) + publisher.complete() + result.join() shouldBe List(1, 2, 3) + + test("should discard undelivered elements and cancel the subscription when the consumer stops early"): + val publisher = new ManualPublisher[Int] + val discarded = new ConcurrentLinkedQueue[Int]() + supervised: + val result = fork(FlowReactiveStreams.fromPublisher(publisher, discarded.add(_).discard).take(1).runToList()) + publisher.subscribed.await() + (1 to 5).foreach(publisher.next) + result.join() shouldBe List(1) + discarded.asScala.toList.sorted shouldBe List(2, 3, 4, 5) + publisher.cancelled.get() shouldBe true +end FlowReactiveStreamsFromPublisherTest