From 6cbdfb96d27c1945a01bdb0c9db2385538692a20 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Fri, 10 Jul 2026 19:58:11 +0800 Subject: [PATCH] fix: release completed graph interpreter references Motivation: Completed graph stages could remain reachable through interpreter logic, owner, handler, and connection payload references after finalization. Modification: Finalize completed stages with a bounded traversal, release their logic and connection references, retain only payloads that remain observable through grab(), and keep snapshots stable after owners are released. Add focused lifecycle and reference-release coverage. Result: Completed stages are released promptly without breaking cascading completion, initialization ordering, snapshots, or port semantics. Tests: - sbt "stream-tests / Test / testOnly org.apache.pekko.stream.impl.fusing.GraphInterpreterPortsSpec org.apache.pekko.stream.impl.fusing.GraphInterpreterSpec org.apache.pekko.stream.impl.GraphStageLogicSpec org.apache.pekko.stream.impl.fusing.LifecycleInterpreterSpec" (106 passed) - scalafmt --list --mode diff-ref=origin/main (passed) - sbt +headerCheckAll and checkCodeStyle (passed) - Qoder stdout review (No must-fix findings) - git diff --check origin/main..HEAD (passed) - +mimaReportBinaryIssues interrupted during final sbt reset after Scala 2.13/3.3 sub-runs completed; full CI requested - validatePullRequest not run - user requested remote CI References: Refs #3030 --- .../stream/impl/GraphStageLogicSpec.scala | 33 +++- .../impl/fusing/GraphInterpreterSpec.scala | 113 ++++++++++++ .../fusing/LifecycleInterpreterSpec.scala | 25 ++- .../stream/impl/fusing/GraphInterpreter.scala | 162 +++++++++++++++--- 4 files changed, 306 insertions(+), 27 deletions(-) diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/impl/GraphStageLogicSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/impl/GraphStageLogicSpec.scala index 28b7b9a5243..82c04bcc0b4 100644 --- a/stream-tests/src/test/scala/org/apache/pekko/stream/impl/GraphStageLogicSpec.scala +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/impl/GraphStageLogicSpec.scala @@ -260,17 +260,42 @@ class GraphStageLogicSpec extends StreamSpec with GraphInterpreterSpecKit with S // note: a bit dangerous assumptions about connection and logic positions here // if anything around creating the logics and connections in the builder changes this may fail + // Save references before execute() since afterStageHasRun releases completed logics + val gLogic = interpreter.logics(1) + val passThroughLogic = interpreter.logics(2) interpreter.complete(interpreter.connections(0)) interpreter.cancel(interpreter.connections(1), SubscriptionWithCancelException.NoMoreElementsNeeded) - interpreter.execute(2) + interpreter.execute(1) expectMsg("postStop2") + interpreter.isSuspended should ===(true) + interpreter.isStageCompleted(gLogic) should ===(true) + interpreter.isStageCompleted(passThroughLogic) should ===(false) + interpreter.logics(gLogic.stageId) should be(null) + interpreter.logics(passThroughLogic.stageId) shouldBe theSameInstanceAs(passThroughLogic) + interpreter.activeStage should be(null) + + val upstreamConnection = interpreter.connections(0) + upstreamConnection.inOwner should be(null) + upstreamConnection.inHandler should be(null) + + val partiallyReleasedConnection = interpreter.connections(1) + partiallyReleasedConnection.outOwner should be(null) + partiallyReleasedConnection.outHandler should be(null) + partiallyReleasedConnection.inOwner shouldBe theSameInstanceAs(passThroughLogic) + partiallyReleasedConnection.inHandler should not be null + + val snapshot = interpreter.toSnapshot + val partiallyReleasedConnectionSnapshot = snapshot.connections(1) + partiallyReleasedConnectionSnapshot.out should ===(snapshot.logics(gLogic.stageId)) + partiallyReleasedConnectionSnapshot.in should ===(snapshot.logics(passThroughLogic.stageId)) + partiallyReleasedConnectionSnapshot.out.label should ===("") + partiallyReleasedConnectionSnapshot.in.label should not be "" + + interpreter.execute(1) expectNoMessage(Duration.Zero) - interpreter.isCompleted should ===(false) interpreter.isSuspended should ===(false) - interpreter.isStageCompleted(interpreter.logics(1)) should ===(true) - interpreter.isStageCompleted(interpreter.logics(2)) should ===(false) } "not allow push from constructor" in { diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/impl/fusing/GraphInterpreterSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/impl/fusing/GraphInterpreterSpec.scala index 10bcb91d408..5e58a44213f 100644 --- a/stream-tests/src/test/scala/org/apache/pekko/stream/impl/fusing/GraphInterpreterSpec.scala +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/impl/fusing/GraphInterpreterSpec.scala @@ -315,6 +315,119 @@ class GraphInterpreterSpec extends StreamSpec with GraphInterpreterSpecKit { interpreter.isSuspended should be(false) } + "release all stage references when finishing an active interpreter" in new TestSetup { + val source = new UpstreamProbe[Int]("source") + val sink = new DownstreamProbe[Int]("sink") + val identityStage = GraphStages.identity[Int] + + builder(identityStage) + .connect(source, identityStage.in) + .connect(identityStage.out, sink) + .init() + + lastEvents() should ===(Set.empty[TestEvent]) + + sink.requestOne() + lastEvents() should ===(Set(RequestOne(source))) + + // Leave an element queued so finish() also has to release connection payloads + source.onNext(1, eventLimit = 0) + interpreter.isSuspended should be(true) + lastEvents() should ===(Set.empty[TestEvent]) + interpreter.activeStage should not be null + interpreter.connections.filter(_ ne null).exists(_.slot != GraphInterpreter.Empty) should be(true) + + val logics = interpreter.logics + + // All logics should be non-null initially + logics.foreach(logic => logic should not be null) + + // Abort the still-running interpreter and release all stage references + interpreter.finish() + + // After finish(), all stage logics should have been released (set to null) + logics.foreach(logic => logic should be(null)) + interpreter.activeStage should be(null) + + // Connection-level references should also be nulled to fully break Connection -> GraphStageLogic chains + interpreter.connections.filter(_ ne null).foreach { conn => + conn.inOwner should be(null) + conn.outOwner should be(null) + conn.inHandler should be(null) + conn.outHandler should be(null) + conn.slot should be(GraphInterpreter.Empty) + } + + val snapshot = interpreter.toSnapshot + snapshot.logics.map(_.label).toSet should ===(Set("")) + snapshot.connections should have size interpreter.connections.count(_ ne null) + snapshot.connections.foreach { connection => + connection.in.label should ===("") + connection.out.label should ===("") + } + } + + "snapshot connections whose stage ids cannot be resolved" in new TestSetup { + val source = new UpstreamProbe[Int]("source") + val sink = new DownstreamProbe[Int]("sink") + val identityStage = GraphStages.identity[Int] + + builder(identityStage) + .connect(source, identityStage.in) + .connect(identityStage.out, sink) + .init() + + interpreter.connections.head.id = -1 + + val connectionSnapshot = interpreter.toSnapshot.connections.head + connectionSnapshot.in.label should ===("") + connectionSnapshot.out.label should ===("") + } + + "release references after cascading stage completion" in new TestSetup { + val source = new UpstreamProbe[Int]("source") + val detachStage = detacher[Int] + val identityStage = GraphStages.identity[Int] + val sink = new DownstreamProbe[Int]("sink") + + builder(detachStage, identityStage) + .connect(source, detachStage.shape.in) + .connect(detachStage.shape.out, identityStage.in) + .connect(identityStage.out, sink) + .init() + + lastEvents() should ===(Set.empty[TestEvent]) + + sink.requestOne() + lastEvents() should ===(Set(RequestOne(source))) + + val logics = interpreter.logics + + // All logics should be non-null initially + logics.foreach(logic => logic should not be null) + + source.onNext(1) + lastEvents() should ===(Set(OnNext(sink, 1), RequestOne(source))) + + // Complete the source - this triggers completion of detacher and identity stages. + // afterStageHasRun releases logics during normal stage completion (not just in finish()). + source.onComplete() + lastEvents() should ===(Set(OnComplete(sink))) + interpreter.isCompleted should be(true) + + // After normal stage completion via afterStageHasRun, all stage logics should be released + logics.foreach(logic => logic should be(null)) + + // Connection-level references should also be nulled when both sides are finalized + interpreter.connections.filter(_ ne null).foreach { conn => + conn.inOwner should be(null) + conn.outOwner should be(null) + conn.inHandler should be(null) + conn.outHandler should be(null) + conn.slot should be(GraphInterpreter.Empty) + } + } + } } diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/impl/fusing/LifecycleInterpreterSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/impl/fusing/LifecycleInterpreterSpec.scala index 6419fdc90c9..4cf3d3cd169 100644 --- a/stream-tests/src/test/scala/org/apache/pekko/stream/impl/fusing/LifecycleInterpreterSpec.scala +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/impl/fusing/LifecycleInterpreterSpec.scala @@ -16,7 +16,7 @@ package org.apache.pekko.stream.impl.fusing import scala.concurrent.duration._ import org.apache.pekko -import pekko.stream.Attributes +import pekko.stream.{ Attributes, ClosedShape } import pekko.stream.impl.fusing.GraphStages.SimpleLinearGraphStage import pekko.stream.stage._ import pekko.stream.testkit.StreamSpec @@ -90,6 +90,19 @@ class LifecycleInterpreterSpec extends StreamSpec with GraphInterpreterSpecKit { expectNoMessage(300.millis) } + "call preStart before postStop for zero-port stages" in new TestSetup { + builder( + ZeroPortStage(onStart = () => testActor ! "start-a", onStop = () => testActor ! "stop-a"), + ZeroPortStage(onStart = () => testActor ! "start-b", onStop = () => testActor ! "stop-b")).init() + + expectMsg("start-a") + expectMsg("stop-a") + expectMsg("start-b") + expectMsg("stop-b") + expectNoMessage(300.millis) + interpreter.isCompleted should be(true) + } + "onError when preStart fails" in new OneBoundedSetup[String](PreStartFailer(() => throw boom)) { lastEvents() should ===(Set(Cancel(boom), OnError(boom))) } @@ -188,6 +201,16 @@ class LifecycleInterpreterSpec extends StreamSpec with GraphInterpreterSpecKit { override def toString = "PreStartAndPostStopIdentity" } + private[pekko] case class ZeroPortStage(onStart: () => Unit, onStop: () => Unit) extends GraphStage[ClosedShape] { + override val shape: ClosedShape = ClosedShape + + override def createLogic(attributes: Attributes): GraphStageLogic = + new GraphStageLogic(shape) { + override def preStart(): Unit = onStart() + override def postStop(): Unit = onStop() + } + } + private[pekko] case class PreStartFailer[T](pleaseThrow: () => Unit) extends SimpleLinearGraphStage[T] { override def createLogic(attributes: Attributes): GraphStageLogic = diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/GraphInterpreter.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/GraphInterpreter.scala index 7c7db1bca02..10f06936c9a 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/GraphInterpreter.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/GraphInterpreter.scala @@ -244,6 +244,28 @@ import pekko.stream.stage._ // Marks whether a stage has been finalized (finalizeStage been called) or not private val finalizedMark = Array.fill(logics.length)(false) + // Number of stages whose preStart has been invoked. Zero-port stages are completed from construction, so + // finalization scans during init must not include later stages that have not been started yet. + private var initializedStages = 0 + + // Connection owners are released when their stages complete, but snapshots still need stable endpoint identities. + private val connectionInStageIds = connectionStageIds(isInput = true) + private val connectionOutStageIds = connectionStageIds(isInput = false) + + private def connectionStageIds(isInput: Boolean): Array[Int] = { + val stageIds = new Array[Int](connections.length) + var i = 0 + while (i < connections.length) { + val connection = connections(i) + if (connection ne null) { + val owner = if (isInput) connection.inOwner else connection.outOwner + stageIds(i) = if (owner ne null) owner.stageId else -1 + } else stageIds(i) = -1 + i += 1 + } + stageIds + } + private var _subFusingMaterializer: Materializer = _ private lazy val defaultErrorReportingLogLevel = LogLevels.defaultErrorLevel(materializer.system) @@ -330,6 +352,7 @@ import pekko.stream.stage._ log.error(e, "Error during preStart in [{}]: {}", logic.toString, e.getMessage) logic.failStage(e) } + initializedStages = i + 1 afterStageHasRun(logic) i += 1 } @@ -342,25 +365,61 @@ import pekko.stream.stage._ var i = 0 while (i < logics.length) { val logic = logics(i) - if (!isStageCompleted(logic) && !isStageFinalized(logic)) { - markStageFinalized(logic) - finalizeStage(logic) + if (logic ne null) { + if (!isStageFinalized(logic)) { + markStageFinalized(logic) + finalizeStage(logic) + } + releaseStage(logic) } i += 1 } + activeStage = null + pendingFinalization = false + + // Release connection references and any in-flight payloads retained by the stopped interpreter + var j = 0 + while (j < connections.length) { + val conn = connections(j) + if (conn ne null) { + conn.inHandler = null + conn.outHandler = null + conn.inOwner = null + conn.outOwner = null + conn.slot = Empty + } + j += 1 + } } // Debug name for a connections input part - private def inOwnerName(connection: Connection): String = connection.inOwner.toString + private def inOwnerName(connection: Connection): String = + if (connection.inOwner ne null) connection.inOwner.toString else "" // Debug name for a connections output part - private def outOwnerName(connection: Connection): String = connection.outOwner.toString + private def outOwnerName(connection: Connection): String = + if (connection.outOwner ne null) connection.outOwner.toString else "" + + private def stageIdFor(stageIds: Array[Int], connection: Connection): Int = { + val connectionId = connection.id + if (connectionId >= 0 && connectionId < stageIds.length) stageIds(connectionId) else -1 + } + + private def connectionInStageId(connection: Connection): Int = stageIdFor(connectionInStageIds, connection) + + private def connectionOutStageId(connection: Connection): Int = stageIdFor(connectionOutStageIds, connection) + + private def logicName(stageId: Int): String = + if (stageId >= 0 && stageId < logics.length) { + val logic = logics(stageId) + if (logic ne null) logic.toString else "" + } else "" // Debug name for a connections input part - private def inLogicName(connection: Connection): String = logics(connection.inOwner.stageId).toString + private def inLogicName(connection: Connection): String = logicName(connectionInStageId(connection)) // Debug name for a connections output part - private def outLogicName(connection: Connection): String = logics(connection.outOwner.stageId).toString + private def outLogicName(connection: Connection): String = logicName(connectionOutStageId(connection)) private def shutdownCounters: String = shutdownCounter.map(x => if (x >= KeepGoingFlag) s"${x & KeepGoingMask}(KeepGoing)" else x.toString).mkString(",") @@ -435,7 +494,6 @@ import pekko.stream.stage._ case NonFatal(e) => reportStageError(e) } if (pendingFinalization) { - pendingFinalization = false afterStageHasRun(activeStage) } @@ -471,7 +529,6 @@ import pekko.stream.stage._ case NonFatal(e) => reportStageError(e) } if (pendingFinalization) { - pendingFinalization = false afterStageHasRun(activeStage) } } @@ -485,7 +542,6 @@ import pekko.stream.stage._ case NonFatal(e) => reportStageError(e) } if (pendingFinalization) { - pendingFinalization = false afterStageHasRun(activeStage) } } @@ -630,12 +686,67 @@ import pekko.stream.stage._ queueTail += 1 } - def afterStageHasRun(logic: GraphStageLogic): Unit = - if (isStageCompleted(logic) && !isStageFinalized(logic)) { - markStageFinalized(logic) - runningStages -= 1 - finalizeStage(logic) + def afterStageHasRun(logic: GraphStageLogic): Unit = { + if ((logic ne null) && isStageCompleted(logic) && !isStageFinalized(logic)) + pendingFinalization = true + + while (pendingFinalization) { + pendingFinalization = false + var stageId = 0 + while (stageId < initializedStages) { + val completedLogic = logics(stageId) + if ((completedLogic ne null) && isStageCompleted(completedLogic) && !isStageFinalized(completedLogic)) { + markStageFinalized(completedLogic) + runningStages -= 1 + finalizeStage(completedLogic) + releaseStage(completedLogic) + } + stageId += 1 + } } + } + + private def hasAvailableElement(connection: Connection): Boolean = { + val state = connection.portState + val normalArrived = (state & (InReady | InFailed | InClosed)) == InReady + + if (normalArrived) connection.slot.asInstanceOf[AnyRef] ne Empty + else if ((state & (InReady | InClosed | InFailed)) == (InReady | InClosed)) + connection.slot match { + case Empty | _: Cancelled => false + case _ => true + } + else if ((state & (InReady | InFailed)) == (InReady | InFailed)) + connection.slot match { + case Failed(_, elem) => elem.asInstanceOf[AnyRef] ne Empty + case _ => false + } + else false + } + + private def releaseStage(logic: GraphStageLogic): Unit = { + var i = 0 + while (i < logic.portToConn.length) { + val connection = logic.portToConn(i) + if (connection ne null) { + if (connection.inOwner eq logic) { + connection.inOwner = null + connection.inHandler = null + } + if (connection.outOwner eq logic) { + connection.outOwner = null + connection.outHandler = null + } + // Completion and failure may follow an onPush before the element is grabbed. Keep that element available + // even after both stages have stopped, while still dropping payloads that can no longer be observed. + if ((connection.inOwner eq null) && (connection.outOwner eq null) && !hasAvailableElement(connection)) + connection.slot = Empty + } + i += 1 + } + if (activeStage eq logic) activeStage = null + logics(logic.stageId) = null + } // Returns true if the given stage is already completed def isStageCompleted(stage: GraphStageLogic): Boolean = (stage ne null) && shutdownCounter(stage.stageId) == 0 @@ -701,7 +812,8 @@ import pekko.stream.stage._ enqueue(connection) } else if ((currentState & (InClosed | Pushing | Pulling | OutClosed)) == 0) enqueue(connection) - if ((currentState & OutClosed) == 0) completeConnection(connection.outOwner.stageId) + if ((currentState & OutClosed) == 0 && (connection.outOwner ne null)) + completeConnection(connection.outOwner.stageId) } @InternalStableApi @@ -720,7 +832,8 @@ import pekko.stream.stage._ enqueue(connection) } } - if ((currentState & OutClosed) == 0) completeConnection(connection.outOwner.stageId) + if ((currentState & OutClosed) == 0 && (connection.outOwner ne null)) + completeConnection(connection.outOwner.stageId) } @InternalStableApi @@ -738,7 +851,7 @@ import pekko.stream.stage._ enqueue(connection) } } - if ((currentState & InClosed) == 0) completeConnection(connection.inOwner.stageId) + if ((currentState & InClosed) == 0 && (connection.inOwner ne null)) completeConnection(connection.inOwner.stageId) } /** @@ -751,15 +864,20 @@ import pekko.stream.stage._ def toSnapshot: RunningInterpreter = { val logicSnapshots = logics.zipWithIndex.map { - case (logic, idx) => + case (logic, idx) if logic ne null => LogicSnapshotImpl(idx, logic.toString, logic.attributes) + case (_, idx) => + LogicSnapshotImpl(idx, "", Attributes.none) } - val logicIndexes = logics.zipWithIndex.map { case (stage, idx) => stage -> idx }.toMap + def safeLogicSnapshot(stageId: Int): LogicSnapshotImpl = + if (stageId >= 0 && stageId < logicSnapshots.length) logicSnapshots(stageId) + else LogicSnapshotImpl(stageId, "", Attributes.none) + val connectionSnapshots = connections.filter(_ ne null).map { connection => ConnectionSnapshotImpl( connection.id, - logicSnapshots(logicIndexes(connection.inOwner)), - logicSnapshots(logicIndexes(connection.outOwner)), + safeLogicSnapshot(connectionInStageId(connection)), + safeLogicSnapshot(connectionOutStageId(connection)), connection.portState match { case InReady | Pushing => ConnectionSnapshot.ShouldPull case OutReady | Pulling => ConnectionSnapshot.ShouldPush