From 9f13fb3eaab7b268af9aff057422a63b2111cbbf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Wed, 5 Aug 2026 00:37:35 +0800 Subject: [PATCH 1/6] 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. 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 +++ .../main/paradox/stream/operators/index.md | 2 + .../operators/sink/WatchTermination.java | 54 ++++ .../operators/sink/WatchTermination.scala | 50 ++++ .../apache/pekko/stream/javadsl/SinkTest.java | 14 + .../scaladsl/SinkWatchTerminationSpec.scala | 201 +++++++++++++ .../stream/impl/fusing/WatchedSink.scala | 268 ++++++++++++++++++ .../apache/pekko/stream/javadsl/Sink.scala | 21 ++ .../apache/pekko/stream/scaladsl/Sink.scala | 22 +- .../pekko/stream/stage/GraphStage.scala | 14 + 10 files changed, 689 insertions(+), 1 deletion(-) create mode 100644 docs/src/main/paradox/stream/operators/Sink/watchTermination.md create mode 100644 docs/src/test/java/jdocs/stream/operators/sink/WatchTermination.java create mode 100644 docs/src/test/scala/docs/stream/operators/sink/WatchTermination.scala create mode 100644 stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/SinkWatchTerminationSpec.scala create mode 100644 stream/src/main/scala/org/apache/pekko/stream/impl/fusing/WatchedSink.scala 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..7cded9ce24f 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 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 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 00000000000..1df1970a2e3 --- /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 00000000000..7a7f8fd864c --- /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 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() From df56648f25a24c545eeb8184e5be2d5d88f5cf4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Fri, 14 Aug 2026 23:08:01 +0800 Subject: [PATCH 2/6] refactor: fuse InHandler into WatchedSinkLogic and use OptionVal Motivation: WatchedSinkLogic allocated an anonymous InHandler instance per materialization and used Option[Throwable] in the termination path, both causing unnecessary heap allocations. Modification: - Fuse the InHandler directly into WatchedSinkLogic (with InHandler), setting handlers(0) = this, following the established Pekko pattern used by CountSink, BroadcastSinkLogic, PartitionSinkLogic, etc. - Replace Option[Throwable] with OptionVal[Throwable] (value class) in upstreamFailureFromConnection to eliminate boxing on the termination signal path Result: One fewer anonymous class allocation per materialization and zero heap allocation in the termination signal path. Behavior unchanged. Tests: - sbt "stream-tests/Test/testOnly org.apache.pekko.stream.scaladsl.SinkWatchTerminationSpec" - 17/17 passed References: Refs #3409 --- .../stream/impl/fusing/WatchedSink.scala | 69 +++++++++---------- 1 file changed, 34 insertions(+), 35 deletions(-) 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 index 7a7f8fd864c..a0731d22be7 100644 --- 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 @@ -150,7 +150,8 @@ import org.apache.pekko.util.OptionVal inner: GraphStageLogic, innerStage: GraphStageWithMaterializedValue[? <: Shape, ?], terminationPromise: Promise[Done]) - extends GraphStageLogic(inner.inCount, inner.outCount) { + extends GraphStageLogic(inner.inCount, inner.outCount) + with InHandler { private var terminationFailure: Throwable = _ private var terminationSignalled = false @@ -162,39 +163,37 @@ import org.apache.pekko.util.OptionVal // 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 + // fuse the inlet handler into this logic 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 - } + handlers(0) = this - override def onUpstreamFinish(): Unit = { - terminationSignalled = true - try innerInHandler.onUpstreamFinish() - catch { - case NonFatal(e) => - terminationFailure = e - throw e - } + override def onPush(): Unit = + try innerInHandler.onPush() + 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 onUpstreamFinish(): Unit = { + terminationSignalled = true + try innerInHandler.onUpstreamFinish() + catch { + case NonFatal(e) => + terminationFailure = e + throw e } + } - override def toString: String = s"WatchedSink($innerInHandler)" + override def onUpstreamFailure(ex: Throwable): Unit = { + terminationSignalled = true + terminationFailure = ex + try innerInHandler.onUpstreamFailure(ex) + catch { + case NonFatal(e) => + terminationFailure = e + throw e + } } private[stream] override def interpreter_=(gi: GraphInterpreter): Unit = { @@ -236,8 +235,8 @@ import org.apache.pekko.util.OptionVal if (failure ne null) terminationPromise.tryFailure(failure) else upstreamFailureFromConnection match { - case Some(ex) => terminationPromise.tryFailure(ex) - case None => + case OptionVal.Some(ex) => terminationPromise.tryFailure(ex) + case _ => if (!terminationSignalled && isAbruptTermination) terminationPromise.tryFailure(new AbruptStageTerminationException(this)) else terminationPromise.trySuccess(Done) @@ -245,16 +244,16 @@ import org.apache.pekko.util.OptionVal } // 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 + // through the fused handler above; the failure remains visible on the connection slot until // after this stage has been finalized. - private def upstreamFailureFromConnection: Option[Throwable] = { + private def upstreamFailureFromConnection: OptionVal[Throwable] = { val connection = portToConn(0) if (connection ne null) connection.slot match { - case GraphInterpreter.Failed(ex, _) => Some(ex) - case _ => None + case GraphInterpreter.Failed(ex, _) => OptionVal.Some(ex) + case _ => OptionVal.None } - else None + else OptionVal.None } // postStop ran without any side of the inlet connection ever being closed, so no completion, From 269a2d3d510ee3e2b9ce43a3c9206d987716c420 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Sat, 15 Aug 2026 17:30:44 +0800 Subject: [PATCH 3/6] style: scalafmt formatting for WatchedSink Motivation: CI checks "Code is formatted" and "Check / Code Style" failed because WatchedSink.scala was not formatted with scalafmt after the refactor commit that fused InHandler and introduced OptionVal. Modification: Run scalafmt on WatchedSink.scala to fix arrow alignment and line wrapping. Result: Both scalafmt CI checks pass. Tests: - scalafmt --mode diff-ref=origin/main --check - All files formatted - git diff --check - clean References: Refs #3409 --- .../apache/pekko/stream/impl/fusing/WatchedSink.scala | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) 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 index a0731d22be7..9ee4f4c4b01 100644 --- 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 @@ -61,9 +61,8 @@ import org.apache.pekko.util.OptionVal 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)) + 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 => @@ -96,7 +95,7 @@ import org.apache.pekko.util.OptionVal case Pop => stack.removeLast() case PushNotUsed => stack.addLast(NotUsed) case transform: Transform => stack.addLast(transform(stack.removeLast())) - case compose: Compose => + case compose: Compose => val second = stack.removeLast() val first = stack.removeLast() stack.addLast(compose(first, second)) @@ -236,7 +235,7 @@ import org.apache.pekko.util.OptionVal else upstreamFailureFromConnection match { case OptionVal.Some(ex) => terminationPromise.tryFailure(ex) - case _ => + case _ => if (!terminationSignalled && isAbruptTermination) terminationPromise.tryFailure(new AbruptStageTerminationException(this)) else terminationPromise.trySuccess(Done) From 4e056af7af248b88a95ec36eafe5b4826bef445f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Sat, 15 Aug 2026 19:23:50 +0800 Subject: [PATCH 4/6] fix: delegate dynamically to inner handler in WatchedSinkLogic Motivation: WatchedSinkLogic captured the inner logic's inlet handler once at construction time (val innerInHandler). Stages that swap their inlet handler after materialization (e.g. LazySink.switchTo) would have subsequent events delegated to the stale handler. For LazySink this caused the termination promise to hang forever: the old handler's onUpstreamFinish calls setKeepGoing(true) instead of completing the stage, so postStop never runs. Modification: - Replace the cached innerInHandler val with dynamic lookups via inner.handlers(0).asInstanceOf[InHandler] in onPush, onUpstreamFinish, and onUpstreamFailure. - Add a regression test that sends multiple elements through Sink.lazySink followed by completion, verifying the termination future completes successfully. Result: Handler swaps by the wrapped stage are respected. The termination promise completes correctly for all stream lifecycle events regardless of when the wrapped stage swaps its inlet handler. Tests: - sbt "stream-tests/Test/testOnly org.apache.pekko.stream.scaladsl.SinkWatchTerminationSpec" - 18/18 passed - sbt "stream-tests/Test/testOnly org.apache.pekko.stream.scaladsl.FlowWatchTerminationSpec org.apache.pekko.stream.scaladsl.QueueSinkSpec org.apache.pekko.stream.scaladsl.SinkSpec" - 68 passed - scalafmt --mode diff-ref=origin/main --check - All files formatted - git diff --check - clean References: Refs #3409 --- .../stream/scaladsl/SinkWatchTerminationSpec.scala | 11 +++++++++++ .../apache/pekko/stream/impl/fusing/WatchedSink.scala | 11 ++++++----- 2 files changed, 17 insertions(+), 5 deletions(-) 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 index 1df1970a2e3..4440f44871a 100644 --- 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 @@ -156,6 +156,17 @@ class SinkWatchTerminationSpec extends StreamSpec { 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") 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 index 9ee4f4c4b01..2890b9f0066 100644 --- 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 @@ -162,12 +162,13 @@ import org.apache.pekko.util.OptionVal // delegate all port handlers to the wrapped logic System.arraycopy(inner.handlers, 0, handlers, 0, handlers.length) - // fuse the inlet handler into this logic to record why the stream terminated - private val innerInHandler = inner.handlers(0).asInstanceOf[InHandler] + // Fuse the inlet handler into this logic to record why the stream terminated. + // Delegate dynamically via inner.handlers(0) so that stages which swap their + // inlet handler after materialization (e.g. LazySink.switchTo) are handled correctly. handlers(0) = this override def onPush(): Unit = - try innerInHandler.onPush() + try inner.handlers(0).asInstanceOf[InHandler].onPush() catch { case NonFatal(e) => terminationFailure = e @@ -176,7 +177,7 @@ import org.apache.pekko.util.OptionVal override def onUpstreamFinish(): Unit = { terminationSignalled = true - try innerInHandler.onUpstreamFinish() + try inner.handlers(0).asInstanceOf[InHandler].onUpstreamFinish() catch { case NonFatal(e) => terminationFailure = e @@ -187,7 +188,7 @@ import org.apache.pekko.util.OptionVal override def onUpstreamFailure(ex: Throwable): Unit = { terminationSignalled = true terminationFailure = ex - try innerInHandler.onUpstreamFailure(ex) + try inner.handlers(0).asInstanceOf[InHandler].onUpstreamFailure(ex) catch { case NonFatal(e) => terminationFailure = e From 5f5f9dd52ea7d55dc2176c75a1615ff273ae9cac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Sat, 15 Aug 2026 23:25:33 +0800 Subject: [PATCH 5/6] feat: support arbitrary composite sinks in Sink.watchTermination MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Motivation: Sink.watchTermination previously only worked with single-stage sinks, rejecting composite sinks built with GraphDSL, Sink.combine, etc. Akka issue #22546 requested a wrapper that materializes a Future[Done] completing only after the sink's postStop has run, for any sink shape. Modification: Rewrite WatchedSink to wrap every GraphStageModule in the sink's traversal with a TerminationReporterStage. A shared TerminationTracker counts stage completions and resolves the promise when the last stage's postStop runs. A per-materialization TrackerHolder ensures independent futures across re-materializations of the same blueprint. Connection slot scanning detects failures (including Cancelled with non-trivial cause) for stages whose handler exceptions bypass the try-catch wrapper. Result: Sink.watchTermination now works with any sink composed of GraphStages: Sink.foreach, Sink.combine, GraphDSL graphs, async islands, LazySink, Sink.queue, and multi-materialization of the same blueprint. Tests: - sbt "stream-tests / Test / testOnly org.apache.pekko.stream.scaladsl.SinkWatchTerminationSpec" → 24 passed - sbt "stream-tests / Test / testOnly org.apache.pekko.stream.javadsl.SinkTest" → 21 passed References: Refs akka/akka-core#22546 --- .../apache/pekko/stream/javadsl/SinkTest.java | 8 +- .../scaladsl/SinkWatchTerminationSpec.scala | 101 +++++- .../stream/impl/fusing/WatchedSink.scala | 316 ++++++++++-------- 3 files changed, 266 insertions(+), 159 deletions(-) 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 7cded9ce24f..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 @@ -288,8 +288,10 @@ public void mustBeAbleToUseWatchTermination() throws Exception { } @Test - public void watchTerminationMustRejectCompositeSinks() { - assertThrows( - IllegalArgumentException.class, () -> Sink.foreach(x -> {}).watchTermination(Keep.right())); + 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 index 4440f44871a..3f2a386f14f 100644 --- 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 @@ -100,18 +100,94 @@ class SinkWatchTerminationSpec extends StreamSpec { 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") + "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) } - "reject sinks created with Sink.combine" in { + "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](_)) - intercept[IllegalArgumentException] { - combined.watchTermination(Keep.right) + 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 { @@ -208,5 +284,14 @@ class SinkWatchTerminationSpec extends StreamSpec { 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) + } } } 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 index 2890b9f0066..78ba04b34d5 100644 --- 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 @@ -20,7 +20,7 @@ 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.Done import org.apache.pekko.annotation.InternalApi import org.apache.pekko.stream._ import org.apache.pekko.stream.impl._ @@ -31,45 +31,64 @@ 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. + * 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 - // 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) + 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 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 => + 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 that consist of a single GraphStage, but [$sink] " + + 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]]) @@ -80,59 +99,77 @@ import org.apache.pekko.util.OptionVal 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]") +/** + * INTERNAL API + * + * Mutable holder that provides a fresh [[TerminationTracker]] per materialization. + * The traversal walk is sequential: stages call `tracker` (lazy-creating on first access), + * and the trailing Transform step calls `reset()` after capturing the future, so the next + * materialization walk starts clean. + */ +@InternalApi private[pekko] final class TrackerHolder(stageCount: Int) { + private var _tracker: TerminationTracker = _ + + def tracker: TerminationTracker = { + if (_tracker eq null) _tracker = new TerminationTracker(stageCount) + _tracker + } + + def reset(): Unit = { _tracker = null } +} + +/** + * 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) } - 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] { +@InternalApi private[pekko] final class TerminationReporterStage( + inner: GraphStageWithMaterializedValue[Shape, Any], + holder: TrackerHolder) + extends GraphStageWithMaterializedValue[Shape, Any] { - override val shape: SinkShape[In] = inner.shape + override val shape: Shape = inner.shape - override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Mat2) = + override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Any) = logicAndMat(inheritedAttributes, null) private[pekko] override def createLogicAndMaterializedValue( inheritedAttributes: Attributes, - materializer: Materializer): (GraphStageLogic, Mat2) = + materializer: Materializer): (GraphStageLogic, Any) = 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)) + 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)" @@ -141,59 +178,57 @@ import org.apache.pekko.util.OptionVal /** * 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. + * 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 WatchedSinkLogic( +@InternalApi private[pekko] final class TerminationReporterLogic( inner: GraphStageLogic, innerStage: GraphStageWithMaterializedValue[? <: Shape, ?], - terminationPromise: Promise[Done]) - extends GraphStageLogic(inner.inCount, inner.outCount) - with InHandler { + tracker: TerminationTracker) + extends GraphStageLogic(inner.inCount, inner.outCount) { private var terminationFailure: Throwable = _ - private var terminationSignalled = false + private var sawTerminationSignal = false + private var reported = 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()) + // Fires when the interpreter finalizes the inner logic directly (e.g. completeStage from async callback) + inner.setTerminationHook(() => reportTermination()) - // delegate all port handlers to the wrapped logic System.arraycopy(inner.handlers, 0, handlers, 0, handlers.length) - // Fuse the inlet handler into this logic to record why the stream terminated. - // Delegate dynamically via inner.handlers(0) so that stages which swap their - // inlet handler after materialization (e.g. LazySink.switchTo) are handled correctly. - handlers(0) = this - - override def onPush(): Unit = - try inner.handlers(0).asInstanceOf[InHandler].onPush() - catch { - case NonFatal(e) => - terminationFailure = e - throw e - } - - override def onUpstreamFinish(): Unit = { - terminationSignalled = true - try inner.handlers(0).asInstanceOf[InHandler].onUpstreamFinish() - catch { - case NonFatal(e) => - terminationFailure = e - throw e - } - } + 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 = { - terminationSignalled = true - terminationFailure = ex - try inner.handlers(0).asInstanceOf[InHandler].onUpstreamFailure(ex) - 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 = { @@ -205,7 +240,6 @@ import org.apache.pekko.util.OptionVal 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() } @@ -213,54 +247,40 @@ import org.apache.pekko.util.OptionVal override def preStart(): Unit = try inner.preStart() catch { - case NonFatal(e) => - terminationFailure = e - throw e + 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 OptionVal.Some(ex) => terminationPromise.tryFailure(ex) - case _ => - if (!terminationSignalled && isAbruptTermination) - terminationPromise.tryFailure(new AbruptStageTerminationException(this)) - else terminationPromise.trySuccess(Done) - } + finally reportTermination() } - // If the wrapped stage swapped its inlet handler after materialization, failures no longer pass - // through the fused handler above; the failure remains visible on the connection slot until - // after this stage has been finalized. - private def upstreamFailureFromConnection: OptionVal[Throwable] = { - val connection = portToConn(0) - if (connection ne null) - connection.slot match { - case GraphInterpreter.Failed(ex, _) => OptionVal.Some(ex) - case _ => OptionVal.None + 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 } - else OptionVal.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 + val failure = if (terminationFailure ne null) terminationFailure else connectionFailure + tracker.stageStopped(failure, sawTerminationSignal, connectionClosed, this) + } } override def toString: String = s"WatchedSink($inner)" From d9127e264da84fc82c788bb80ed0237d7ded80c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Sun, 16 Aug 2026 16:18:16 +0800 Subject: [PATCH 6/6] fix: format WatchedSink and fix concurrent materialization race in Sink.watchTermination Motivation: PR #3409 CI failed on the 'Code Style' / 'Code is formatted' checks because WatchedSink.scala was not run through scalafmt. While reviewing the implementation, a real concurrency bug was also found: the internal TrackerHolder used a single plain mutable field to hand a fresh TerminationTracker to each materialization of a watched Sink blueprint. Sink/Flow blueprints are designed to be safely re-materializable from concurrent threads, but concurrently materializing the same watchTermination blueprint raced on that shared field and could corrupt the tracker state, causing one of the resulting Future[Done] to hang forever. Modification: - Ran scalafmt (mode diff-ref=origin/main) to fix formatting of WatchedSink.scala. - Replaced the plain mutable tracker field in TrackerHolder with a ThreadLocal[TerminationTracker]. Each single materialization walk is synchronous and confined to the calling thread, so scoping the tracker to that thread isolates concurrent materializations of the same blueprint from one another. - Added a regression test that materializes the same watchTermination blueprint from two threads concurrently (barrier-synchronized, 500 iterations) and asserts both futures complete with Done; this reproduces the hang against the previous implementation. Result: - CI formatting checks pass locally (scalafmt --list reports no files). - The new concurrency regression test fails (times out) against the old TrackerHolder and passes with the ThreadLocal-based fix. - Full SinkWatchTerminationSpec (25 tests) and GraphStageLogicSpec (17 tests) pass. - sbt +mimaReportBinaryIssues passes (internal API only, no public API changes). Tests: - sbt "stream-tests/testOnly org.apache.pekko.stream.scaladsl.SinkWatchTerminationSpec" - 25/25 passed - sbt "stream-tests/testOnly org.apache.pekko.stream.impl.GraphStageLogicSpec" - 17/17 passed - sbt +mimaReportBinaryIssues - passed - scalafmt --list --mode diff-ref=origin/main - no files to reformat - sbt headerCheckAll - passed References: Refs #3409 --- .../scaladsl/SinkWatchTerminationSpec.scala | 27 ++++++++++- .../stream/impl/fusing/WatchedSink.scala | 45 ++++++++++++------- 2 files changed, 53 insertions(+), 19 deletions(-) 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 index 3f2a386f14f..8db1c3d70c6 100644 --- 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 @@ -17,9 +17,10 @@ package org.apache.pekko.stream.scaladsl -import java.util.concurrent.ConcurrentLinkedQueue +import java.util.concurrent.{ ConcurrentLinkedQueue, CyclicBarrier } -import scala.concurrent.{ ExecutionContext, Promise } +import scala.concurrent.{ Await, ExecutionContext, Future, Promise } +import scala.concurrent.duration._ import scala.jdk.CollectionConverters._ import scala.util.control.NoStackTrace @@ -293,5 +294,27 @@ class SinkWatchTerminationSpec extends StreamSpec { 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 index 78ba04b34d5..78839e50756 100644 --- 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 @@ -44,7 +44,7 @@ import org.apache.pekko.util.OptionVal val allSteps = steps.result() val stageCount = allSteps.count { - case MaterializeAtomic(_: GraphStageModule[_, _], _) => true + case MaterializeAtomic(_: GraphStageModule[?, ?], _) => true case _ => false } @@ -54,7 +54,7 @@ import org.apache.pekko.util.OptionVal s"contains none.") allSteps.foreach { - case MaterializeAtomic(_: GraphStageModule[_, _], _) => + case MaterializeAtomic(_: GraphStageModule[?, ?], _) => case MaterializeAtomic(other, _) => throw new IllegalArgumentException( s"Sink.watchTermination is only supported for sinks built from GraphStages, but [$sink] " + @@ -67,7 +67,7 @@ import org.apache.pekko.util.OptionVal val holder = new TrackerHolder(stageCount) val newSteps: Vector[Traversal] = allSteps.map { - case MaterializeAtomic(module: GraphStageModule[_, _], outToSlots) => + case MaterializeAtomic(module: GraphStageModule[?, ?], outToSlots) => val reporterStage = new TerminationReporterStage( module.stage.asInstanceOf[GraphStageWithMaterializedValue[Shape, Any]], holder) MaterializeAtomic( @@ -79,11 +79,11 @@ import org.apache.pekko.util.OptionVal 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 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)) @@ -104,20 +104,30 @@ import org.apache.pekko.util.OptionVal /** * INTERNAL API * - * Mutable holder that provides a fresh [[TerminationTracker]] per materialization. - * The traversal walk is sequential: stages call `tracker` (lazy-creating on first access), - * and the trailing Transform step calls `reset()` after capturing the future, so the next - * materialization walk starts clean. + * 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 var _tracker: TerminationTracker = _ + private val threadLocalTracker = new ThreadLocal[TerminationTracker] def tracker: TerminationTracker = { - if (_tracker eq null) _tracker = new TerminationTracker(stageCount) - _tracker + var t = threadLocalTracker.get() + if (t eq null) { + t = new TerminationTracker(stageCount) + threadLocalTracker.set(t) + } + t } - def reset(): Unit = { _tracker = null } + def reset(): Unit = threadLocalTracker.remove() } /** @@ -269,7 +279,8 @@ import org.apache.pekko.util.OptionVal 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]) + if ((connectionFailure eq null) && + !cause.isInstanceOf[SubscriptionWithCancelException.NonFailureCancellation]) connectionFailure = cause case _ => }