This is an automated email from the ASF dual-hosted git repository. He-Pin pushed a commit to branch feat/sink-watch-termination in repository https://gitbox.apache.org/repos/asf/pekko.git
commit 9f13fb3eaab7b268af9aff057422a63b2111cbbf Author: θιΈ£ <[email protected]> AuthorDate: Wed Aug 5 00:37:35 2026 +0800 feat: add Sink.watchTermination operator Motivation: Sometimes you want to wait for a Sink to fully complete, including any cleanup work or final commit it performs in postStop, but the sink does not materialize a Future[Done]. The existing watchTermination operator is placed before the sink and therefore only signals when the upstream of the sink has terminated (see apache/pekko#2377, akka/akka-core#22546). Modification: Add Sink.watchTermination to the Scala and Java DSLs. It wraps sinks that consist of a single GraphStage with a delegating stage whose materialized Future[Done] completes only after the wrapped sink's postStop has run, fails with the upstream failure when the stream failed, and fails with an AbruptStreamTerminationException when the stream was abruptly terminated. The original materialized value, including mapMaterializedValue transforms, is preserved. Composite sinks consisting of multiple stages are rejected with an IllegalArgumentException. Implementation details: - WatchedSink rewrites the sink traversal, replacing the single terminal stage with a WatchedSinkStage and replaying the trailing materialized value composition steps. - WatchedSinkLogic delegates all port handlers and lifecycle hooks to the wrapped logic, mirroring interpreter, port wiring, stageId and attributes, and records termination causes from the delegated handlers, handler exceptions, and the connection failure slot (covering wrapped stages that swap their inlet handler after materialization). - GraphStageLogic gains an internal termination hook fired from afterPostStop so the promise also completes when the interpreter finalizes the wrapped logic directly (async-callback self-termination). Result: Users can await full sink termination, including postStop cleanup, via a materialized Future[Done] / CompletionStage<Done>. Tests: - sbt "stream-tests/Test/testOnly org.apache.pekko.stream.scaladsl.SinkWatchTerminationSpec" - 17/17 passed - sbt "stream-tests/Test/testOnly org.apache.pekko.stream.scaladsl.*Sink*" - 159 passed - sbt "stream-tests/Test/testOnly org.apache.pekko.stream.scaladsl.FlowWatchTerminationSpec org.apache.pekko.stream.scaladsl.QueueSinkSpec org.apache.pekko.stream.scaladsl.GraphStageTimersSpec org.apache.pekko.stream.impl.GraphStageLogicSpec org.apache.pekko.stream.impl.SubInletOutletSpec org.apache.pekko.stream.impl.LinearTraversalBuilderSpec org.apache.pekko.stream.DslConsistencySpec" - 124 passed - sbt "stream-tests/Test/testOnly org.apache.pekko.stream.javadsl.SinkTest" - passed - sbt stream/mimaReportBinaryIssues - no issues - sbt "++3.3.8" stream/compile - passed - sbt docs/paradox - passed - sbt headerCreateAll scalafmtAll scalafmtSbt javafmtCheckAll - passed - scalafmt --mode diff-ref=origin/main - no changes - git diff --check - clean - sbt sortImports - failed with scalafix plugin NoSuchMethodError (environment issue), imports kept consistent manually References: Fixes #2377 --- .../stream/operators/Sink/watchTermination.md | 44 ++++ docs/src/main/paradox/stream/operators/index.md | 2 + .../stream/operators/sink/WatchTermination.java | 54 +++++ .../stream/operators/sink/WatchTermination.scala | 50 ++++ .../org/apache/pekko/stream/javadsl/SinkTest.java | 14 ++ .../stream/scaladsl/SinkWatchTerminationSpec.scala | 201 ++++++++++++++++ .../pekko/stream/impl/fusing/WatchedSink.scala | 268 +++++++++++++++++++++ .../org/apache/pekko/stream/javadsl/Sink.scala | 21 ++ .../org/apache/pekko/stream/scaladsl/Sink.scala | 22 +- .../org/apache/pekko/stream/stage/GraphStage.scala | 14 ++ 10 files changed, 689 insertions(+), 1 deletion(-) 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 0000000000..871921c99b --- /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<Done>`] 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<Done>`] 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 011a1e1cf7..8a9cc4c75c 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|<a name="seq"></a>@ref[seq](Sink/seq.md)|Collect values emitted from the stream into a collection.| |Sink|<a name="source"></a>@ref[source](Sink/source.md)|A `Sink` that materializes this `Sink` itself as a `Source`, the returning `Source` can only have one subscriber.| |Sink|<a name="takelast"></a>@ref[takeLast](Sink/takeLast.md)|Collect the last `n` values emitted from the stream into a collection.| +|Sink|<a name="watchtermination"></a>@ref[watchTermination](Sink/watchTermination.md)|Wraps a sink so that in addition to the original materialized value a @scala[`Future[Done]`] @java[`CompletionStage<Done>`] 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 0000000000..f6a6a1fd62 --- /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<ByteString, CompletionStage<IOResult>> fileSink = + FileIO.toPath(Paths.get("target/watch-termination.txt")); + + // In addition to the IOResult of the file sink, materialize a CompletionStage<Done> + // that only completes once the file has been fully written and closed. + final Pair<CompletionStage<IOResult>, CompletionStage<Done>> result = + Source.single(ByteString.fromString("Hello, world!")) + .runWith(fileSink.watchTermination(Keep.both()), system); + + final CompletionStage<IOResult> ioResult = result.first(); + final CompletionStage<Done> 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 0000000000..ee7c674365 --- /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 059f6d34f8..7cded9ce24 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,18 @@ public class SinkTest extends StreamTestJupiter { .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<Integer>, CompletionStage<Done>> result = + Source.range(1, 4).runWith(Sink.<Integer>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 watchTerminationMustRejectCompositeSinks() { + assertThrows( + IllegalArgumentException.class, () -> Sink.foreach(x -> {}).watchTermination(Keep.right())); + } } 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 0000000000..1df1970a2e --- /dev/null +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/SinkWatchTerminationSpec.scala @@ -0,0 +1,201 @@ +/* + * 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 + +import scala.concurrent.{ ExecutionContext, Promise } +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] + } + + "reject composite sinks consisting of multiple stages" in { + val ex = intercept[IllegalArgumentException] { + Sink.foreach[Int](println).watchTermination(Keep.right) + } + ex.getMessage should include("single stage") + } + + "reject sinks created with Sink.combine" in { + val combined = Sink.combine(Sink.ignore, Sink.ignore)(Broadcast[Int](_)) + intercept[IllegalArgumentException] { + combined.watchTermination(Keep.right) + } + } + + "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 } + } + + "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) + } + } +} 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 0000000000..7a7f8fd864 --- /dev/null +++ b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/WatchedSink.scala @@ -0,0 +1,268 @@ +/* + * 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, NotUsed } +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 a sink that consists of a single [[GraphStageWithMaterializedValue]] + * so that, in addition to the original materialized value, a `Future[Done]` is materialized that only completes + * after the wrapped sink's `postStop` lifecycle hook has 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 + // use traversalSoFar rather than traversal, which would additionally wrap island and attribute + // steps that the builder keeps separately and re-applies on access + val steps = Vector.newBuilder[Traversal] + flatten(builder.traversalSoFar, steps) + val allSteps = steps.result() + + val moduleIndices = allSteps.indices.filter(i => allSteps(i).isInstanceOf[MaterializeAtomic]) + if (moduleIndices.size != 1) + throw new IllegalArgumentException( + s"Sink.watchTermination is only supported for sinks that consist of a single stage, but [$sink] consists " + + s"of ${moduleIndices.size} stages. Composite sinks such as those created with Sink.combine or GraphDSL " + + s"are not supported.") + + val moduleIndex = moduleIndices.head + allSteps(moduleIndex) match { + case MaterializeAtomic(module: GraphStageModule[SinkShape[In] @unchecked, Mat @unchecked], outToSlots) + if outToSlots.isEmpty => + val prefixSteps = allSteps.take(moduleIndex) + val suffixSteps = allSteps.drop(moduleIndex + 1) + + val watchedStage = new WatchedSinkStage[In, Mat, Mat2](module.stage, suffixSteps, matF) + val watchedModule = GraphStageModule(module.shape, module.attributes, watchedStage) + val newTraversal = + (prefixSteps :+ (MaterializeAtomic(watchedModule, outToSlots): Traversal)) + .foldLeft(EmptyTraversal: Traversal)((traversal, step) => traversal.concat(step)) + + new Sink(builder.copy(traversalSoFar = newTraversal), sink.shape) + case other => + throw new IllegalArgumentException( + s"Sink.watchTermination is only supported for sinks that consist of a single GraphStage, but [$sink] " + + s"contains [$other].") + } + } + + 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 + } + + /** + * Replays the materialized value composition steps that followed the wrapped stage in the original + * traversal, transforming the wrapped stage's materialized value into the materialized value the + * original sink would have produced. + */ + private[fusing] def runMatProgram(steps: Vector[Traversal], initial: Any): Any = { + val stack = new java.util.ArrayDeque[Any](4) + stack.addLast(initial) + var i = 0 + while (i < steps.length) { + steps(i) match { + case Pop => stack.removeLast() + case PushNotUsed => stack.addLast(NotUsed) + case transform: Transform => stack.addLast(transform(stack.removeLast())) + case compose: Compose => + val second = stack.removeLast() + val first = stack.removeLast() + stack.addLast(compose(first, second)) + case other => + throw new IllegalArgumentException( + s"Sink.watchTermination encountered an unexpected materialized value composition step [$other]") + } + i += 1 + } + stack.removeLast() + } +} + +/** + * INTERNAL API + */ +@InternalApi private[pekko] final class WatchedSinkStage[-In, Mat, Mat2]( + inner: GraphStageWithMaterializedValue[SinkShape[In], Mat], + trailingMatProgram: Vector[Traversal], + matF: (Mat, Future[Done]) => Mat2) + extends GraphStageWithMaterializedValue[SinkShape[In], Mat2] { + + override val shape: SinkShape[In] = inner.shape + + override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Mat2) = + logicAndMat(inheritedAttributes, null) + + private[pekko] override def createLogicAndMaterializedValue( + inheritedAttributes: Attributes, + materializer: Materializer): (GraphStageLogic, Mat2) = + logicAndMat(inheritedAttributes, materializer) + + private def logicAndMat(inheritedAttributes: Attributes, materializer: Materializer): (GraphStageLogic, Mat2) = { + val (innerLogic, innerMat) = inner.createLogicAndMaterializedValue(inheritedAttributes, materializer) + val terminationPromise = Promise[Done]() + val sinkMat = WatchedSink.runMatProgram(trailingMatProgram, innerMat).asInstanceOf[Mat] + (new WatchedSinkLogic(innerLogic, inner, terminationPromise), matF(sinkMat, terminationPromise.future)) + } + + override def toString: String = s"WatchedSink($inner)" +} + +/** + * INTERNAL API + * + * A delegating [[GraphStageLogic]] that behaves exactly as the wrapped logic while completing the + * termination promise only after the wrapped logic's `postStop` has run. The future is failed with + * the upstream failure when the stream failed, and completed with success otherwise. + */ +@InternalApi private[pekko] final class WatchedSinkLogic( + inner: GraphStageLogic, + innerStage: GraphStageWithMaterializedValue[? <: Shape, ?], + terminationPromise: Promise[Done]) + extends GraphStageLogic(inner.inCount, inner.outCount) { + + private var terminationFailure: Throwable = _ + private var terminationSignalled = false + + // Completes the promise even if the interpreter finalizes the wrapped logic directly, + // which happens when the wrapped stage terminates itself from an async callback. + inner.setTerminationHook(() => completeTermination()) + + // delegate all port handlers to the wrapped logic + System.arraycopy(inner.handlers, 0, handlers, 0, handlers.length) + + // wrap the inlet handler to record why the stream terminated + private val innerInHandler = inner.handlers(0).asInstanceOf[InHandler] + handlers(0) = new InHandler { + override def onPush(): Unit = + try innerInHandler.onPush() + catch { + case NonFatal(e) => + terminationFailure = e + throw e + } + + override def onUpstreamFinish(): Unit = { + terminationSignalled = true + try innerInHandler.onUpstreamFinish() + catch { + case NonFatal(e) => + terminationFailure = e + throw e + } + } + + override def onUpstreamFailure(ex: Throwable): Unit = { + terminationSignalled = true + terminationFailure = ex + try innerInHandler.onUpstreamFailure(ex) + catch { + case NonFatal(e) => + terminationFailure = e + throw e + } + } + + override def toString: String = s"WatchedSink($innerInHandler)" + } + + 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) + // mirror the port wiring so that the wrapped logic can interact with the interpreter + 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 completeTermination() + } + + protected[stream] override def afterPostStop(): Unit = { + inner.afterPostStop() + completeTermination() + } + + // completeTermination may be invoked more than once (from postStop, afterPostStop and the + // termination hook), the promise only completes on the first invocation + private def completeTermination(): Unit = { + val failure = terminationFailure + if (failure ne null) terminationPromise.tryFailure(failure) + else + upstreamFailureFromConnection match { + case Some(ex) => terminationPromise.tryFailure(ex) + case None => + if (!terminationSignalled && isAbruptTermination) + terminationPromise.tryFailure(new AbruptStageTerminationException(this)) + else terminationPromise.trySuccess(Done) + } + } + + // If the wrapped stage swapped its inlet handler after materialization, failures no longer pass + // through the wrapping handler above; the failure remains visible on the connection slot until + // after this stage has been finalized. + private def upstreamFailureFromConnection: Option[Throwable] = { + val connection = portToConn(0) + if (connection ne null) + connection.slot match { + case GraphInterpreter.Failed(ex, _) => Some(ex) + case _ => None + } + else None + } + + // postStop ran without any side of the inlet connection ever being closed, so no completion, + // failure or cancellation signal reached the wrapped sink + private def isAbruptTermination: Boolean = { + val connection = portToConn(0) + (connection ne null) && (connection.portState & (GraphInterpreter.InClosed | GraphInterpreter.OutClosed)) == 0 + } + + 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 fbf1becfbd..050ce0e0dc 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<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[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 7828e4b503..4f7785bc4b 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 bae071f83a..dbc84de195 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() --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
