diff --git a/docs/src/main/paradox/stream/operators/Sink/watchTermination.md b/docs/src/main/paradox/stream/operators/Sink/watchTermination.md new file mode 100644 index 00000000000..871921c99b1 --- /dev/null +++ b/docs/src/main/paradox/stream/operators/Sink/watchTermination.md @@ -0,0 +1,44 @@ +# Sink.watchTermination + +Wraps a sink so that in addition to the original materialized value a @scala[`Future[Done]`] @java[`CompletionStage`] is materialized that only completes after the wrapped sink has fully terminated, including its `postStop` lifecycle hook. + +@ref[Sink operators](../index.md#sink-operators) + +## Signature + +@apidoc[Sink.watchTermination](Sink) { scala="#watchTermination[Mat2]()(matF:(Mat,scala.concurrent.Future[org.apache.pekko.Done])=>Mat2):org.apache.pekko.stream.scaladsl.Sink[In,Mat2]" java="#watchTermination(org.apache.pekko.japi.function.Function2)" } + + +## Description + +Wraps a sink so that in addition to the original materialized value a @scala[`Future[Done]`] @java[`CompletionStage`] is materialized +that completes when the wrapped sink has fully terminated: it completes with success after the wrapped sink's `postStop` +lifecycle hook has run, or fails with the upstream failure when the stream failed. + +This differs from @ref[watchTermination](../Source-or-Flow/watchTermination.md), which is placed *before* the sink and +therefore only signals when the upstream of the sink has terminated. Because `Sink.watchTermination` wraps the sink +itself, the materialized @scala[`Future`] @java[`CompletionStage`] can be used to wait for any cleanup or final +commits the sink performs in `postStop`, for example a file sink closing the file it was writing to. + +Only sinks that consist of a single `GraphStage` are supported, for example `Sink.ignore`, `Sink.head`, +`Sink.queue`, `Sink.actorRef` or sinks created from custom graph stages. Composite sinks consisting of +multiple stages — such as `Sink.foreach`, `Sink.fold`, or sinks created with `Sink.combine` or `GraphDSL` — +are not supported and throw an @scala[`IllegalArgumentException`] @java[`IllegalArgumentException`]. + +## Examples + +Scala +: @@snip [WatchTermination.scala](/docs/src/test/scala/docs/stream/operators/sink/WatchTermination.scala) { #watchTermination } + +Java +: @@snip [WatchTermination.java](/docs/src/test/java/jdocs/stream/operators/sink/WatchTermination.java) { #watchTermination } + +## Reactive Streams semantics + +@@@div { .callout } + +**backpressures** when the wrapped sink backpressures + +**cancels** when the wrapped sink cancels + +@@@ diff --git a/docs/src/main/paradox/stream/operators/index.md b/docs/src/main/paradox/stream/operators/index.md index 011a1e1cf76..8a9cc4c75c7 100644 --- a/docs/src/main/paradox/stream/operators/index.md +++ b/docs/src/main/paradox/stream/operators/index.md @@ -87,6 +87,7 @@ These built-in sinks are available from @scala[`org.apache.pekko.stream.scaladsl |Sink|@ref[seq](Sink/seq.md)|Collect values emitted from the stream into a collection.| |Sink|@ref[source](Sink/source.md)|A `Sink` that materializes this `Sink` itself as a `Source`, the returning `Source` can only have one subscriber.| |Sink|@ref[takeLast](Sink/takeLast.md)|Collect the last `n` values emitted from the stream into a collection.| +|Sink|@ref[watchTermination](Sink/watchTermination.md)|Wraps a sink so that in addition to the original materialized value a @scala[`Future[Done]`] @java[`CompletionStage`] is materialized that only completes after the wrapped sink has fully terminated, including its `postStop` lifecycle hook.| ## Additional Sink and Source converters @@ -615,6 +616,7 @@ For more background see the @ref[Error Handling in Streams](../stream-error.md) * [UnzipWith](UnzipWith.md) * [watch](Source-or-Flow/watch.md) * [watchTermination](Source-or-Flow/watchTermination.md) +* [watchTermination](Sink/watchTermination.md) * [wireTap](Source-or-Flow/wireTap.md) * [withBackoff](RestartSource/withBackoff.md) * [withBackoff](RestartFlow/withBackoff.md) diff --git a/docs/src/test/java/jdocs/stream/operators/sink/WatchTermination.java b/docs/src/test/java/jdocs/stream/operators/sink/WatchTermination.java new file mode 100644 index 00000000000..f6a6a1fd629 --- /dev/null +++ b/docs/src/test/java/jdocs/stream/operators/sink/WatchTermination.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package jdocs.stream.operators.sink; + +import java.nio.file.Paths; +import java.util.concurrent.CompletionStage; +import org.apache.pekko.Done; +import org.apache.pekko.actor.ActorSystem; +import org.apache.pekko.japi.Pair; +import org.apache.pekko.stream.IOResult; +import org.apache.pekko.stream.javadsl.FileIO; +import org.apache.pekko.stream.javadsl.Keep; +import org.apache.pekko.stream.javadsl.Sink; +import org.apache.pekko.stream.javadsl.Source; +import org.apache.pekko.util.ByteString; + +public class WatchTermination { + + private ActorSystem system = null; + + void example() { + // #watchTermination + final Sink> fileSink = + FileIO.toPath(Paths.get("target/watch-termination.txt")); + + // In addition to the IOResult of the file sink, materialize a CompletionStage + // that only completes once the file has been fully written and closed. + final Pair, CompletionStage> result = + Source.single(ByteString.fromString("Hello, world!")) + .runWith(fileSink.watchTermination(Keep.both()), system); + + final CompletionStage ioResult = result.first(); + final CompletionStage terminated = result.second(); + + // Once `terminated` completes the sink has stopped, including its postStop + // cleanup, so the file is guaranteed to be closed at this point. + // #watchTermination + } +} diff --git a/docs/src/test/scala/docs/stream/operators/sink/WatchTermination.scala b/docs/src/test/scala/docs/stream/operators/sink/WatchTermination.scala new file mode 100644 index 00000000000..ee7c674365f --- /dev/null +++ b/docs/src/test/scala/docs/stream/operators/sink/WatchTermination.scala @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package docs.stream.operators.sink + +import java.nio.file.Paths + +import scala.concurrent.Future + +import org.apache.pekko +import pekko.Done +import pekko.actor.ActorSystem +import pekko.stream.IOResult +import pekko.stream.scaladsl.{ FileIO, Keep, Sink, Source } +import pekko.util.ByteString + +object WatchTermination { + implicit val system: ActorSystem = ??? + + def watchTerminationExample(): Unit = { + // #watchTermination + val fileSink: Sink[ByteString, Future[IOResult]] = + FileIO.toPath(Paths.get("target/watch-termination.txt")) + + // In addition to the IOResult of the file sink, materialize a Future[Done] + // that only completes once the file has been fully written and closed. + val (ioResult, terminated): (Future[IOResult], Future[Done]) = + Source + .single(ByteString("Hello, world!")) + .runWith(fileSink.watchTermination(Keep.both)) + + // Once `terminated` completes the sink has stopped, including its postStop + // cleanup, so the file is guaranteed to be closed at this point. + // #watchTermination + } +} diff --git a/stream-tests/src/test/java/org/apache/pekko/stream/javadsl/SinkTest.java b/stream-tests/src/test/java/org/apache/pekko/stream/javadsl/SinkTest.java index 059f6d34f85..b0cb1ab3f5a 100644 --- a/stream-tests/src/test/java/org/apache/pekko/stream/javadsl/SinkTest.java +++ b/stream-tests/src/test/java/org/apache/pekko/stream/javadsl/SinkTest.java @@ -278,4 +278,20 @@ public void mustBeAbleToUseSinkAsSource() throws Exception { .get(1, TimeUnit.SECONDS); assertEquals(List.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10), r); } + + @Test + public void mustBeAbleToUseWatchTermination() throws Exception { + final Pair, CompletionStage> result = + Source.range(1, 4).runWith(Sink.head().watchTermination(Keep.both()), system); + assertEquals(1, result.first().toCompletableFuture().get(1, TimeUnit.SECONDS).intValue()); + assertEquals(Done.done(), result.second().toCompletableFuture().get(1, TimeUnit.SECONDS)); + } + + @Test + public void watchTerminationMustWorkWithCompositeSinks() throws Exception { + final CompletionStage done = + Source.range(1, 4) + .runWith(Sink.foreach(x -> {}).watchTermination(Keep.right()), system); + assertEquals(Done.done(), done.toCompletableFuture().get(1, TimeUnit.SECONDS)); + } } diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/SinkWatchTerminationSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/SinkWatchTerminationSpec.scala new file mode 100644 index 00000000000..8db1c3d70c6 --- /dev/null +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/SinkWatchTerminationSpec.scala @@ -0,0 +1,320 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.pekko.stream.scaladsl + +import java.util.concurrent.{ ConcurrentLinkedQueue, CyclicBarrier } + +import scala.concurrent.{ Await, ExecutionContext, Future, Promise } +import scala.concurrent.duration._ +import scala.jdk.CollectionConverters._ +import scala.util.control.NoStackTrace + +import org.apache.pekko +import pekko.Done +import pekko.stream._ +import pekko.stream.stage.{ GraphStage, GraphStageLogic, InHandler } +import pekko.stream.testkit.StreamSpec +import pekko.stream.testkit.scaladsl.TestSource + +class SinkWatchTerminationSpec extends StreamSpec { + + "A Sink.watchTermination" must { + + "complete future with success when stream is completed" in { + val done = Source(1 to 4).runWith(Sink.ignore.watchTermination(Keep.right)) + done.futureValue should ===(Done) + } + + "complete future with success when the stream is empty" in { + val done = Source.empty[Int].runWith(Sink.ignore.watchTermination(Keep.right)) + done.futureValue should ===(Done) + } + + "complete future with success when the sink cancels itself" in { + val done = Source(1 to 4).runWith(Sink.head[Int].watchTermination(Keep.right)) + done.futureValue should ===(Done) + } + + "keep the original materialized value" in { + val (head, done) = Source(1 to 4).runWith(Sink.head[Int].watchTermination(Keep.both)) + head.futureValue should ===(1) + done.futureValue should ===(Done) + } + + "keep materialized value transformations of the wrapped sink" in { + val transformed: Sink[Int, scala.concurrent.Future[Int]] = + Sink.headOption[Int].mapMaterializedValue(_.map(_.getOrElse(0))(ExecutionContext.parasitic)) + val (head, done) = Source(1 to 4).runWith(transformed.watchTermination(Keep.both)) + head.futureValue should ===(1) + done.futureValue should ===(Done) + } + + "fail future when stream is failed" in { + val ex = new RuntimeException("Stream failed.") with NoStackTrace + val (p, done) = TestSource[Int]().toMat(Sink.ignore.watchTermination(Keep.right))(Keep.both).run() + p.sendNext(1) + p.sendError(ex) + whenReady(done.failed) { _ shouldBe ex } + } + + "complete future only after the postStop of the wrapped sink has run" in { + val events = new ConcurrentLinkedQueue[String]() + + class PostStopSignalingSink extends GraphStage[SinkShape[Int]] { + val in = Inlet[Int]("PostStopSignalingSink.in") + override val shape: SinkShape[Int] = SinkShape(in) + + override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = + new GraphStageLogic(shape) with InHandler { + override def preStart(): Unit = pull(in) + override def onPush(): Unit = pull(in) + override def postStop(): Unit = events.add("postStop") + setHandler(in, this) + } + } + + val done = Source(1 to 4).runWith(Sink.fromGraph(new PostStopSignalingSink).watchTermination(Keep.right)) + done.onComplete(_ => events.add("futureCompleted"))(ExecutionContext.parasitic) + done.futureValue should ===(Done) + events.asScala.toList should ===(List("postStop", "futureCompleted")) + } + + "fail future when stream abruptly terminated" in { + val mat = Materializer(system) + val done = TestSource[Int]().toMat(Sink.ignore.watchTermination(Keep.right))(Keep.both).run()(mat)._2 + mat.shutdown() + done.failed.futureValue shouldBe an[AbruptTerminationException] + } + + "complete future for a composite sink built with Sink.foreach" in { + val done = Source(1 to 4).runWith(Sink.foreach[Int](_ => ()).watchTermination(Keep.right)) + done.futureValue should ===(Done) + } + + "complete future for a composite sink built with Sink.combine" in { + val combined = Sink.combine(Sink.ignore, Sink.ignore)(Broadcast[Int](_)) + val done = Source(1 to 4).runWith(combined.watchTermination(Keep.right)) + done.futureValue should ===(Done) + } + + "complete future for a composite sink built with GraphDSL" in { + val composite = Sink.fromGraph(GraphDSL.create() { implicit b => + import GraphDSL.Implicits._ + val bcast = b.add(Broadcast[Int](2)) + val s1 = b.add(Sink.ignore) + val s2 = b.add(Sink.ignore) + bcast.out(0) ~> s1 + bcast.out(1) ~> s2 + SinkShape(bcast.in) + }) + val done = Source(1 to 4).runWith(composite.watchTermination(Keep.right)) + done.futureValue should ===(Done) + } + + "fail future when upstream fails on a composite sink" in { + val ex = new RuntimeException("composite fail") with NoStackTrace + val combined = Sink.combine(Sink.ignore, Sink.ignore)(Broadcast[Int](_)) + val (p, done) = TestSource[Int]().toMat(combined.watchTermination(Keep.right))(Keep.both).run() + p.sendNext(1) + p.sendError(ex) + whenReady(done.failed) { _ shouldBe ex } + } + + "complete future only after all stages' postStop have run in a composite sink" in { + val events = new ConcurrentLinkedQueue[String]() + + class SignalingSink(name: String) extends GraphStage[SinkShape[Int]] { + val in = Inlet[Int](s"$name.in") + override val shape: SinkShape[Int] = SinkShape(in) + + override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = + new GraphStageLogic(shape) with InHandler { + override def preStart(): Unit = pull(in) + override def onPush(): Unit = pull(in) + override def postStop(): Unit = events.add(name) + setHandler(in, this) + } + } + + val composite = Sink.fromGraph(GraphDSL.create() { implicit b => + import GraphDSL.Implicits._ + val bcast = b.add(Broadcast[Int](2)) + val s1 = b.add(new SignalingSink("s1")) + val s2 = b.add(new SignalingSink("s2")) + bcast.out(0) ~> s1 + bcast.out(1) ~> s2 + SinkShape(bcast.in) + }) + + val done = Source(1 to 4).runWith(composite.watchTermination(Keep.right)) + done.onComplete(_ => events.add("futureCompleted"))(ExecutionContext.parasitic) + done.futureValue should ===(Done) + val list = events.asScala.toList + list.last should ===("futureCompleted") + list.filterNot(_ == "futureCompleted").toSet should ===(Set("s1", "s2")) + } + + "keep the original materialized value of a composite sink" in { + val composite = Sink.fromGraph(GraphDSL.createGraph(Sink.queue[Int]()) { implicit b => queue => + import GraphDSL.Implicits._ + val bcast = b.add(Broadcast[Int](2)) + bcast.out(0) ~> queue + bcast.out(1) ~> Sink.ignore + SinkShape(bcast.in) + }) + val (queue, done) = Source(1 to 4).runWith(composite.watchTermination(Keep.both)) + queue.pull().futureValue should ===(Some(1)) + queue.cancel() + done.futureValue should ===(Done) + } + + "fail future when a fully fused composite stream abruptly terminated" in { + val mat = Materializer(system) + val combined = Sink.combine(Sink.ignore, Sink.ignore)(Broadcast[Int](_)) + val done = Source.maybe[Int].toMat(combined.watchTermination(Keep.right))(Keep.right).run()(mat) + mat.shutdown() + done.failed.futureValue shouldBe an[AbruptStageTerminationException] + } + + "work with Sink.queue" in { + val (queue, done) = Source(1 to 4).runWith(Sink.queue[Int]().watchTermination(Keep.both)) + queue.pull().futureValue should ===(Some(1)) + queue.pull().futureValue should ===(Some(2)) + queue.cancel() + done.futureValue should ===(Done) + } + + "signal termination once after single materialization value promise completed" in { + val terminationSignal = Promise[Done]() + + class CompletingSink extends GraphStage[SinkShape[Int]] { + val in = Inlet[Int]("CompletingSink.in") + override val shape: SinkShape[Int] = SinkShape(in) + + override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = + new GraphStageLogic(shape) with InHandler { + override def preStart(): Unit = pull(in) + override def onPush(): Unit = pull(in) + override def onUpstreamFinish(): Unit = { + terminationSignal.trySuccess(Done) + completeStage() + } + setHandler(in, this) + } + } + + val done = Source(1 to 4).runWith(Sink.fromGraph(new CompletingSink).watchTermination(Keep.right)) + terminationSignal.future.futureValue should ===(Done) + done.futureValue should ===(Done) + } + + "fail future when stream is failed after the wrapped sink swapped its inlet handler" in { + val ex = new RuntimeException("Stream failed.") with NoStackTrace + val (p, done) = TestSource[Int]() + .toMat(Sink.lazySink(() => Sink.ignore).watchTermination(Keep.right))(Keep.both) + .run() + p.sendNext(1) + p.sendError(ex) + whenReady(done.failed) { _ shouldBe ex } + } + + "complete future when stream completes after the wrapped sink swapped its inlet handler" in { + val (p, done) = TestSource[Int]() + .toMat(Sink.lazySink(() => Sink.ignore).watchTermination(Keep.right))(Keep.both) + .run() + p.sendNext(1) + p.sendNext(2) + p.sendNext(3) + p.sendComplete() + done.futureValue should ===(Done) + } + + "fail future when a handler of the wrapped sink throws" in { + class FailingSink extends GraphStage[SinkShape[Int]] { + val in = Inlet[Int]("FailingSink.in") + override val shape: SinkShape[Int] = SinkShape(in) + + override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = + new GraphStageLogic(shape) with InHandler { + override def preStart(): Unit = pull(in) + override def onPush(): Unit = throw new RuntimeException("boom") with NoStackTrace + setHandler(in, this) + } + } + + val done = Source.single(1).runWith(Sink.fromGraph(new FailingSink).watchTermination(Keep.right)) + done.failed.futureValue shouldBe a[RuntimeException] + } + + "fail future when a fully fused stream abruptly terminated" in { + val mat = Materializer(system) + val done = Source.maybe[Int].toMat(Sink.ignore.watchTermination(Keep.right))(Keep.right).run()(mat) + mat.shutdown() + done.failed.futureValue shouldBe an[AbruptStageTerminationException] + } + + "fail future when upstream of Sink.queue fails" in { + val ex = new RuntimeException("Stream failed.") with NoStackTrace + val (p, (queue, done)) = + TestSource[Int]() + .toMat(Sink.queue[Int]().watchTermination(Keep.both))(Keep.both) + .run() + p.sendNext(1) + queue.pull().futureValue should ===(Some(1)) + p.sendError(ex) + queue.pull().failed.futureValue shouldBe ex + whenReady(done.failed) { _ shouldBe ex } + } + + "work with a sink behind an async island" in { + val done = Source(1 to 4).runWith(Sink.ignore.async.watchTermination(Keep.right)) + done.futureValue should ===(Done) + } + + "produce independent futures when the same blueprint is materialized multiple times" in { + val watched = Sink.ignore.watchTermination(Keep.right) + val done1 = Source(1 to 4).runWith(watched) + val done2 = Source(5 to 8).runWith(watched) + done1.futureValue should ===(Done) + done2.futureValue should ===(Done) + (done1 should not).be(done2) + } + + "produce independent futures when the same blueprint is materialized concurrently" in { + implicit val ec: ExecutionContext = system.dispatcher + val watched = Sink.ignore.watchTermination(Keep.right) + // Regression test: Sink blueprints must be safely re-materializable from concurrent threads. + // A prior implementation shared a single mutable tracker field across all materializations of + // the same blueprint, which raced when two threads materialized it at the same time and could + // hang one of the resulting futures forever. + for (_ <- 1 to 500) { + val barrier = new CyclicBarrier(2) + val f1 = Future { + barrier.await() + Source(1 to 3).runWith(watched) + }.flatMap(identity) + val f2 = Future { + barrier.await() + Source(4 to 6).runWith(watched) + }.flatMap(identity) + Await.result(f1, 5.seconds) should ===(Done) + Await.result(f2, 5.seconds) should ===(Done) + } + } + } +} diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/WatchedSink.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/WatchedSink.scala new file mode 100644 index 00000000000..78839e50756 --- /dev/null +++ b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/WatchedSink.scala @@ -0,0 +1,298 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.pekko.stream.impl.fusing + +import scala.concurrent.{ Future, Promise } +import scala.util.control.NonFatal + +import org.apache.pekko.Done +import org.apache.pekko.annotation.InternalApi +import org.apache.pekko.stream._ +import org.apache.pekko.stream.impl._ +import org.apache.pekko.stream.scaladsl.Sink +import org.apache.pekko.stream.stage._ +import org.apache.pekko.util.OptionVal + +/** + * INTERNAL API + * + * Implements `Sink.watchTermination`: wraps every [[GraphStageWithMaterializedValue]] in the sink's + * traversal so that a shared `Future[Done]` is materialized that only completes after all wrapped + * stages' `postStop` lifecycle hooks have run. + */ +@InternalApi private[pekko] object WatchedSink { + + def apply[In, Mat, Mat2](sink: Sink[In, Mat], matF: (Mat, Future[Done]) => Mat2): Sink[In, Mat2] = { + val builder = sink.traversalBuilder + val steps = Vector.newBuilder[Traversal] + flatten(builder.traversalSoFar, steps) + val allSteps = steps.result() + + val stageCount = allSteps.count { + case MaterializeAtomic(_: GraphStageModule[?, ?], _) => true + case _ => false + } + + if (stageCount == 0) + throw new IllegalArgumentException( + s"Sink.watchTermination is only supported for sinks that contain at least one GraphStage, but [$sink] " + + s"contains none.") + + allSteps.foreach { + case MaterializeAtomic(_: GraphStageModule[?, ?], _) => + case MaterializeAtomic(other, _) => + throw new IllegalArgumentException( + s"Sink.watchTermination is only supported for sinks built from GraphStages, but [$sink] " + + s"contains [$other].") + case _ => + } + + // Per-materialization holder: the traversal walk is sequential, so the first stage + // creates the tracker and subsequent stages reuse it within the same materialization. + val holder = new TrackerHolder(stageCount) + + val newSteps: Vector[Traversal] = allSteps.map { + case MaterializeAtomic(module: GraphStageModule[?, ?], outToSlots) => + val reporterStage = new TerminationReporterStage( + module.stage.asInstanceOf[GraphStageWithMaterializedValue[Shape, Any]], holder) + MaterializeAtomic( + GraphStageModule(module.shape, module.attributes, + reporterStage.asInstanceOf[GraphStageWithMaterializedValue[Shape, Any]]), + outToSlots): Traversal + case other => other + } + + val matFStep: Traversal = + Transform(((mat: Any) => { + val t = holder.tracker + val result = matF(mat.asInstanceOf[Mat], t.future) + holder.reset() + result + }).asInstanceOf[TraversalBuilder.AnyFunction1]) + + val newTraversal = (newSteps :+ matFStep) + .foldLeft(EmptyTraversal: Traversal)((traversal, step) => traversal.concat(step)) + + new Sink(builder.copy(traversalSoFar = newTraversal), sink.shape) + } + + private def flatten(traversal: Traversal, builder: scala.collection.mutable.Builder[Traversal, Vector[Traversal]]) + : Unit = traversal match { + case EmptyTraversal => + case Concat(first, second) => + flatten(first, builder) + flatten(second, builder) + case other => builder += other + } +} + +/** + * INTERNAL API + * + * Provides a fresh [[TerminationTracker]] per materialization. A single [[TrackerHolder]] instance is + * shared by every materialization of the same `Sink` blueprint (it is captured once, when + * `Sink.watchTermination` builds the traversal), so it must tolerate concurrent materializations of that + * blueprint from different threads, which is a supported usage pattern for stream blueprints. + * + * The traversal walk for a single materialization is always sequential and confined to the thread that + * calls `materialize`/`run`/`runWith`: stages call `tracker` (lazy-creating on first access on that thread), + * and the trailing Transform step calls `reset()` after capturing the future. A plain mutable field would + * therefore race across concurrent materializations; using a [[ThreadLocal]] instead scopes the tracker to + * the materializing thread so concurrent materializations never observe each other's state. + */ +@InternalApi private[pekko] final class TrackerHolder(stageCount: Int) { + private val threadLocalTracker = new ThreadLocal[TerminationTracker] + + def tracker: TerminationTracker = { + var t = threadLocalTracker.get() + if (t eq null) { + t = new TerminationTracker(stageCount) + threadLocalTracker.set(t) + } + t + } + + def reset(): Unit = threadLocalTracker.remove() +} + +/** + * INTERNAL API + */ +@InternalApi private[pekko] final class TerminationTracker(stageCount: Int) { + private var remaining = stageCount + private var _failure: Throwable = _ + private var _sawSignal: Boolean = false + private var _anyConnectionClosed: Boolean = false + private val terminationPromise = Promise[Done]() + + val future: Future[Done] = terminationPromise.future + + def stageStopped(failure: Throwable, sawSignal: Boolean, connectionClosed: Boolean, logic: GraphStageLogic): Unit = + synchronized { + if (failure ne null) _failure = failure + if (sawSignal) _sawSignal = true + if (connectionClosed) _anyConnectionClosed = true + remaining -= 1 + if (remaining == 0) { + if (_failure ne null) terminationPromise.tryFailure(_failure) + else if (!_sawSignal && !_anyConnectionClosed) + terminationPromise.tryFailure(new AbruptStageTerminationException(logic)) + else terminationPromise.trySuccess(Done) + } + } +} + +/** + * INTERNAL API + */ +@InternalApi private[pekko] final class TerminationReporterStage( + inner: GraphStageWithMaterializedValue[Shape, Any], + holder: TrackerHolder) + extends GraphStageWithMaterializedValue[Shape, Any] { + + override val shape: Shape = inner.shape + + override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Any) = + logicAndMat(inheritedAttributes, null) + + private[pekko] override def createLogicAndMaterializedValue( + inheritedAttributes: Attributes, + materializer: Materializer): (GraphStageLogic, Any) = + logicAndMat(inheritedAttributes, materializer) + + private def logicAndMat(inheritedAttributes: Attributes, materializer: Materializer): (GraphStageLogic, Any) = { + val (innerLogic, innerMat) = + if (materializer eq null) inner.createLogicAndMaterializedValue(inheritedAttributes) + else inner.createLogicAndMaterializedValue(inheritedAttributes, materializer) + (new TerminationReporterLogic(innerLogic, inner, holder.tracker), innerMat) + } + + override def toString: String = s"WatchedSink($inner)" +} + +/** + * INTERNAL API + * + * Delegates all behavior to the wrapped logic, reporting termination to a shared + * [[TerminationTracker]] after the wrapped logic's `postStop` has run. + * + * Input handlers delegate dynamically via `inner.handlers(idx)` so that stages which swap + * handlers after materialization (e.g. LazySink) are handled correctly. The try-catch has + * zero JIT cost (exception-table only) and is required to capture failures from leaf stages + * whose handler exceptions would not otherwise appear on any connection slot. + */ +@InternalApi private[pekko] final class TerminationReporterLogic( + inner: GraphStageLogic, + innerStage: GraphStageWithMaterializedValue[? <: Shape, ?], + tracker: TerminationTracker) + extends GraphStageLogic(inner.inCount, inner.outCount) { + + private var terminationFailure: Throwable = _ + private var sawTerminationSignal = false + private var reported = false + + // Fires when the interpreter finalizes the inner logic directly (e.g. completeStage from async callback) + inner.setTerminationHook(() => reportTermination()) + + System.arraycopy(inner.handlers, 0, handlers, 0, handlers.length) + + private var i = 0 + while (i < inCount) { + val idx = i + handlers(idx) = new InHandler { + override def onPush(): Unit = + try inner.handlers(idx).asInstanceOf[InHandler].onPush() + catch { + case NonFatal(e) => terminationFailure = e; throw e + } + + override def onUpstreamFinish(): Unit = { + sawTerminationSignal = true + try inner.handlers(idx).asInstanceOf[InHandler].onUpstreamFinish() + catch { + case NonFatal(e) => terminationFailure = e; throw e + } + } + + override def onUpstreamFailure(ex: Throwable): Unit = { + sawTerminationSignal = true + if (terminationFailure eq null) terminationFailure = ex + try inner.handlers(idx).asInstanceOf[InHandler].onUpstreamFailure(ex) + catch { + case NonFatal(e) => terminationFailure = e; throw e + } + } + } + i += 1 + } + + private[stream] override def interpreter_=(gi: GraphInterpreter): Unit = { + super.interpreter_=(gi) + inner.interpreter_=(gi) + } + + protected[stream] override def beforePreStart(): Unit = { + inner.stageId = stageId + inner.attributes = attributes + inner.originalStage = OptionVal.Some(innerStage) + System.arraycopy(portToConn, 0, inner.portToConn, 0, portToConn.length) + inner.beforePreStart() + } + + override def preStart(): Unit = + try inner.preStart() + catch { + case NonFatal(e) => terminationFailure = e; throw e + } + + override def postStop(): Unit = { + try inner.postStop() + finally reportTermination() + } + + protected[stream] override def afterPostStop(): Unit = inner.afterPostStop() + + private def reportTermination(): Unit = { + if (!reported) { + reported = true + var connectionFailure: Throwable = null + var connectionClosed = false + var j = 0 + while (j < portToConn.length) { + val connection = portToConn(j) + if (connection ne null) { + connection.slot match { + case GraphInterpreter.Failed(ex, _) => if (connectionFailure eq null) connectionFailure = ex + case GraphInterpreter.Cancelled(cause) => + if ((connectionFailure eq null) && + !cause.isInstanceOf[SubscriptionWithCancelException.NonFailureCancellation]) + connectionFailure = cause + case _ => + } + if ((connection.portState & (GraphInterpreter.InClosed | GraphInterpreter.OutClosed)) != 0) + connectionClosed = true + } + j += 1 + } + val failure = if (terminationFailure ne null) terminationFailure else connectionFailure + tracker.stageStopped(failure, sawTerminationSignal, connectionClosed, this) + } + } + + override def toString: String = s"WatchedSink($inner)" +} diff --git a/stream/src/main/scala/org/apache/pekko/stream/javadsl/Sink.scala b/stream/src/main/scala/org/apache/pekko/stream/javadsl/Sink.scala index fbf1becfbde..050ce0e0dce 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/javadsl/Sink.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/javadsl/Sink.scala @@ -658,6 +658,27 @@ final class Sink[In, Mat](delegate: scaladsl.Sink[In, Mat]) extends Graph[SinkSh pekko.japi.Pair(mat, sink.asJava) } + /** + * Wraps this sink so that in addition to the original materialized value a `CompletionStage` is + * materialized that completes when this sink has fully terminated: it completes with success after this sink's + * `postStop` lifecycle hook has run, or fails with the upstream failure when the stream failed. Unlike + * [[Flow.watchTermination]], which only observes termination before the sink, this allows waiting for any + * cleanup or final commits performed by the sink itself. + * + * Only sinks that consist of a single `GraphStage` are supported, for example `Sink.ignore`, `Sink.head`, + * `Sink.queue` or sinks created from custom graph stages. Composite sinks consisting of multiple stages, + * such as `Sink.foreach`, `Sink.fold` or sinks created with `Sink.combine` or `GraphDSL`, are not supported + * and throw an [[IllegalArgumentException]]. + * + * It is recommended to use the internally optimized `Keep.left` and `Keep.right` combiners + * where appropriate instead of manually writing functions that pass through one of the values. + * + * @since 2.0.0 + */ + def watchTermination[M]( + matF: function.Function2[Mat @uncheckedVariance, CompletionStage[Done], M]): Sink[In @uncheckedVariance, M] = + new Sink(delegate.watchTermination((left, right) => matF(left, right.asJava))) + /** * Replace the attributes of this [[Sink]] with the given ones. If this Sink is a composite * of multiple graphs, new attributes on the composite will be less specific than attributes diff --git a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Sink.scala b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Sink.scala index 7828e4b503f..4f7785bc4bc 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Sink.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Sink.scala @@ -27,7 +27,7 @@ import pekko.annotation.InternalApi import pekko.stream._ import pekko.stream.impl._ import pekko.stream.impl.Stages.DefaultAttributes -import pekko.stream.impl.fusing.{ CountSink, GraphStages, SourceSink } +import pekko.stream.impl.fusing.{ CountSink, GraphStages, SourceSink, WatchedSink } import pekko.stream.stage._ import org.reactivestreams.{ Publisher, Subscriber } @@ -82,6 +82,26 @@ final class Sink[-In, +Mat](override val traversalBuilder: LinearTraversalBuilde (mat, Sink.fromSubscriber(sub)) } + /** + * Wraps this sink so that in addition to the original materialized value a `Future[Done]` is materialized + * that completes when this sink has fully terminated: it completes with success after this sink's `postStop` + * lifecycle hook has run, or fails with the upstream failure when the stream failed. Unlike + * [[Flow.watchTermination]], which only observes termination before the sink, this allows waiting for any + * cleanup or final commits performed by the sink itself. + * + * Only sinks that consist of a single [[GraphStage]] are supported, for example `Sink.ignore`, `Sink.head`, + * `Sink.queue` or sinks created from custom graph stages. Composite sinks consisting of multiple stages, + * such as `Sink.foreach`, `Sink.fold` or sinks created with `Sink.combine` or `GraphDSL`, are not supported + * and throw an [[IllegalArgumentException]]. + * + * It is recommended to use the internally optimized `Keep.left` and `Keep.right` combiners + * where appropriate instead of manually writing functions that pass through one of the values. + * + * @since 2.0.0 + */ + def watchTermination[Mat2](matF: (Mat, Future[Done]) => Mat2): Sink[In, Mat2] = + WatchedSink(this, matF) + /** * Replace the attributes of this [[Sink]] with the given ones. If this Sink is a composite * of multiple graphs, new attributes on the composite will be less specific than attributes diff --git a/stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala b/stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala index bae071f83a9..dbc84de1954 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala @@ -1487,6 +1487,10 @@ abstract class GraphStageLogic private[stream] (val inCount: Int, val outCount: new AtomicReference(ConcurrentHashMap.newKeySet()) private var _stageActor: StageActor = _ + + // INTERNAL API: fired from afterPostStop, used to observe termination of wrapped logics + private var terminationHook: () => Unit = _ + final def stageActor: StageActor = _stageActor match { case null => throw StageActorRefNotInitializedException() case ref => ref @@ -1606,8 +1610,18 @@ abstract class GraphStageLogic private[stream] (val inCount: Int, val outCount: callbacks.forEach((t: Promise[Done]) => t.tryFailure(exception)) } cleanUpSubstreams(OptionVal.None) + if (terminationHook ne null) terminationHook() } + /** + * INTERNAL API + * + * Registers a hook that is invoked after this logic's `postStop` has run and its internal + * cleanups have completed. + */ + @InternalApi + private[stream] def setTerminationHook(hook: () => Unit): Unit = terminationHook = hook + /** Called from interpreter thread by GraphInterpreter.runAsyncInput */ private[stream] def onFeedbackDispatched(promise: Promise[Done]): Unit = { val callbacks = asyncCallbacksInProgress.get()