feat: add Sink.watchTermination operator - #3409
Draft
He-Pin wants to merge 6 commits into
Draft
Conversation
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 #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
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
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
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
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
…nk.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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Sometimes you want to wait for a
Sinkto fully complete, including any cleanup work or finalcommit it performs in
postStop, but the sink does not materialize aFuture[Done]. The existingwatchTerminationoperator is placed before the sink and therefore only signals when the upstreamof the sink has terminated (see #2377).
Modification
Add
Sink.watchTerminationto the Scala and Java DSLs. It wraps sinks that consist of a singleGraphStagewith a delegating stage whose materializedFuture[Done]/CompletionStage<Done>:postStophas run,AbruptStageTerminationExceptionwhen the stream was abruptly terminated.The original materialized value, including
mapMaterializedValuetransforms, is preserved.Composite sinks consisting of multiple stages are rejected with an
IllegalArgumentException.Implementation details:
WatchedSinkrewrites the sink traversal, replacing the single terminal stage with aWatchedSinkStageand replaying the trailing materialized-value composition steps.WatchedSinkLogicfusesInHandlerdirectly into theGraphStageLogic(no anonymous handlerallocation), delegates all lifecycle hooks to the wrapped logic, and uses
OptionValforzero-allocation termination signal detection via the connection slot.
GraphStageLogicgains an internal termination hook fired fromafterPostStopso the promisealso completes when the interpreter finalizes the wrapped logic directly (async-callback
self-termination, e.g.
Sink.queuecancellation).Result
Users can await full sink termination, including
postStopcleanup, via a materializedFuture[Done]/CompletionStage<Done>.Tests
sbt "stream-tests/Test/testOnly org.apache.pekko.stream.scaladsl.SinkWatchTerminationSpec"- 17/17 passedsbt "stream-tests/Test/testOnly org.apache.pekko.stream.scaladsl.*Sink*"- 159 passedsbt "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 passedsbt "stream-tests/Test/testOnly org.apache.pekko.stream.javadsl.SinkTest"- passedsbt stream/mimaReportBinaryIssues- no issuessbt "++3.3.8" stream/compile- passedsbt docs/paradox- passedsbt headerCreateAll scalafmtAll scalafmtSbt javafmtCheckAll- passedscalafmt --mode diff-ref=origin/main- no changesgit diff --check- cleanReferences
Fixes #2377