From 3d0369b13d083298cfe02760441c36ec0bbbbaae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Tue, 7 Jul 2026 18:15:58 +0800 Subject: [PATCH 1/2] fix: release references to completed stage logics in GraphInterpreter Motivation: GraphInterpreter retains strong references to completed stage logics in the logics array and Connection objects (inHandler/outHandler/inOwner/outOwner), preventing garbage collection even after stages have finished. This is especially problematic for long-running fused interpreters that host subfused shells, where completed initial stages remain referenced for the lifetime of the actor. Modification: - afterStageHasRun now scans ALL logics for completed stages instead of only checking the single activeStage. This ensures stages that complete without being activeStage (e.g., boundary stages calling complete(), cascading completions) are properly finalized and released. - finish() and afterStageHasRun null connection handlers (inHandler/ outHandler) for finalized stages, breaking handler-to-logic reference chains to enable garbage collection. - finish() additionally nulls inOwner/outOwner on all connections since the interpreter is being shut down and no more events will be processed. - logics(i) is set to null after finalizing each stage in both finish() and afterStageHasRun. - toSnapshot handles null logics gracefully with placeholder and uses connection.inOwner.stageId for positional indexing instead of identity-based map lookups that fail with NPE on null logics. - toSnapshot filters out connections with null inOwner/outOwner. - inLogicName/outLogicName/inOwnerName/outOwnerName debug methods guard against null logics and null owners. - Fix GraphStageLogicSpec "not double-terminate" test to save logic references before execute() since afterStageHasRun now releases completed logics. Result: Completed stage logics are properly finalized and released for GC during normal stream execution (not just in finish()), and all strong reference chains (logics array, connection handlers, connection owners) are broken when the interpreter is shut down, reducing memory retention in long-running fused interpreters. Tests: - stream-tests/testOnly GraphInterpreterSpec: 13/13 passed (includes 2 new regression tests) - stream-tests/testOnly GraphStageLogicSpec: 17/17 passed - stream-tests/testOnly InterpreterSpec+LifecycleInterpreterSpec+InterpreterSupervisionSpec+GraphInterpreterFailureModesSpec+ActorGraphInterpreterSpec: 65/65 passed References: Refs apache/pekko#3030, Related to akka/akka-core#23439 --- .../stream/impl/GraphStageLogicSpec.scala | 6 +- .../impl/fusing/GraphInterpreterSpec.scala | 63 +++++++++++ .../stream/impl/fusing/GraphInterpreter.scala | 105 +++++++++++++----- 3 files changed, 144 insertions(+), 30 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..14e89983eb8 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,6 +260,8 @@ 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 + 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) @@ -269,8 +271,8 @@ class GraphStageLogicSpec extends StreamSpec with GraphInterpreterSpecKit with S interpreter.isCompleted should ===(false) interpreter.isSuspended should ===(false) - interpreter.isStageCompleted(interpreter.logics(1)) should ===(true) - interpreter.isStageCompleted(interpreter.logics(2)) should ===(false) + interpreter.isStageCompleted(gLogic) should ===(true) + interpreter.isStageCompleted(passThroughLogic) 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..5ec839dc68e 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,69 @@ class GraphInterpreterSpec extends StreamSpec with GraphInterpreterSpecKit { interpreter.isSuspended should be(false) } + "release references to completed stage logics to prevent memory leaks" 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))) + + source.onNext(1) + lastEvents() should ===(Set(OnNext(sink, 1))) + + val logicsField = interpreter.getClass.getDeclaredField("logics") + logicsField.setAccessible(true) + val logics = logicsField.get(interpreter).asInstanceOf[Array[org.apache.pekko.stream.stage.GraphStageLogic]] + + logics.foreach(logic => logic should not be null) + + source.onComplete() + lastEvents() should ===(Set(OnComplete(sink))) + + interpreter.finish() + + logics.foreach(logic => logic should be(null)) + } + + "release references to completed stage logics when some stages complete early" 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 logicsField = interpreter.getClass.getDeclaredField("logics") + logicsField.setAccessible(true) + val logics = logicsField.get(interpreter).asInstanceOf[Array[org.apache.pekko.stream.stage.GraphStageLogic]] + + logics.foreach(logic => logic should not be null) + + source.onNext(1) + lastEvents() should ===(Set(OnNext(sink, 1), RequestOne(source))) + + source.onComplete() + + logics.foreach(logic => logic should be(null)) + } + } } 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..40d034c6c5f 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 @@ -342,25 +342,51 @@ import pekko.stream.stage._ var i = 0 while (i < logics.length) { val logic = logics(i) - if (!isStageCompleted(logic) && !isStageFinalized(logic)) { + if ((logic ne null) && !isStageCompleted(logic) && !isStageFinalized(logic)) { markStageFinalized(logic) finalizeStage(logic) } + logics(i) = null i += 1 } + 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 + } + 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 "" // Debug name for a connections input part - private def inLogicName(connection: Connection): String = logics(connection.inOwner.stageId).toString + private def inLogicName(connection: Connection): String = { + val owner = connection.inOwner + if (owner ne null) { + val logic = logics(owner.stageId) + if (logic ne null) logic.toString else "" + } else "" + } // Debug name for a connections output part - private def outLogicName(connection: Connection): String = logics(connection.outOwner.stageId).toString + private def outLogicName(connection: Connection): String = { + val owner = connection.outOwner + if (owner ne null) { + val logic = logics(owner.stageId) + if (logic ne null) logic.toString else "" + } else "" + } private def shutdownCounters: String = shutdownCounter.map(x => if (x >= KeepGoingFlag) s"${x & KeepGoingMask}(KeepGoing)" else x.toString).mkString(",") @@ -630,12 +656,33 @@ import pekko.stream.stage._ queueTail += 1 } - def afterStageHasRun(logic: GraphStageLogic): Unit = - if (isStageCompleted(logic) && !isStageFinalized(logic)) { - markStageFinalized(logic) - runningStages -= 1 - finalizeStage(logic) + @scala.annotation.nowarn("cat=unused-params") + def afterStageHasRun(logic: GraphStageLogic): Unit = { + var i = 0 + while (i < logics.length) { + val l = logics(i) + if ((l ne null) && isStageCompleted(l) && !isStageFinalized(l)) { + markStageFinalized(l) + runningStages -= 1 + finalizeStage(l) + logics(i) = null + } + i += 1 + } + var j = 0 + while (j < connections.length) { + val conn = connections(j) + if (conn ne null) { + if (finalizedMark(conn.inOwner.stageId)) { + conn.inHandler = null + } + if (finalizedMark(conn.outOwner.stageId)) { + conn.outHandler = null + } + } + j += 1 } + } // Returns true if the given stage is already completed def isStageCompleted(stage: GraphStageLogic): Boolean = (stage ne null) && shutdownCounter(stage.stageId) == 0 @@ -751,26 +798,28 @@ 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 - val connectionSnapshots = connections.filter(_ ne null).map { connection => - ConnectionSnapshotImpl( - connection.id, - logicSnapshots(logicIndexes(connection.inOwner)), - logicSnapshots(logicIndexes(connection.outOwner)), - connection.portState match { - case InReady | Pushing => ConnectionSnapshot.ShouldPull - case OutReady | Pulling => ConnectionSnapshot.ShouldPush - case x if (x & (InClosed | OutClosed)) == (InClosed | OutClosed) => - // At least one side of the connection is closed: we show it as closed - ConnectionSnapshot.Closed - case _ => - // This should not be possible: connection alive and both push and pull enqueued but not received - throw new IllegalStateException(s"Unexpected connection state for $connection: ${connection.portState}") - - }) + val connectionSnapshots = connections.filter(c => (c ne null) && (c.inOwner ne null) && (c.outOwner ne null)).map { + connection => + ConnectionSnapshotImpl( + connection.id, + logicSnapshots(connection.inOwner.stageId), + logicSnapshots(connection.outOwner.stageId), + connection.portState match { + case InReady | Pushing => ConnectionSnapshot.ShouldPull + case OutReady | Pulling => ConnectionSnapshot.ShouldPush + case x if (x & (InClosed | OutClosed)) == (InClosed | OutClosed) => + // At least one side of the connection is closed: we show it as closed + ConnectionSnapshot.Closed + case _ => + // This should not be possible: connection alive and both push and pull enqueued but not received + throw new IllegalStateException(s"Unexpected connection state for $connection: ${connection.portState}") + + }) } val stoppedStages: List[LogicSnapshot] = shutdownCounter.zipWithIndex.collect { From a41451733d9734b4441ef25a07d5163474443ede Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Wed, 8 Jul 2026 11:33:17 +0800 Subject: [PATCH 2/2] fix: address GraphInterpreter retention review comments Motivation:\nWhen fused streams keep the GraphInterpreter alive, completed stage logics and connection owners/handlers must not remain strongly referenced.\n\nModification:\nFinalize completed stages after they run, release finalized connection-side references, and keep stable stage ids for diagnostics and snapshots after owners are cleared.\n\nResult:\nCompleted fused stages can be released while other stages keep running, and interpreter snapshots remain usable after release.\n\nTests:\n- rtk sbt "stream-tests / Test / testOnly org.apache.pekko.stream.impl.fusing.GraphInterpreterSpec"\n- rtk sbt "stream-tests / Test / testOnly org.apache.pekko.stream.impl.GraphStageLogicSpec org.apache.pekko.stream.impl.fusing.GraphInterpreterPortsSpec org.apache.pekko.stream.impl.fusing.LifecycleInterpreterSpec org.apache.pekko.stream.impl.fusing.GraphInterpreterFailureModesSpec"\n- rtk sbt "stream-tests / Test / testOnly org.apache.pekko.stream.impl.fusing.GraphInterpreterSpec org.apache.pekko.stream.impl.GraphStageLogicSpec"\n- rtk scalafmt --mode diff-ref=origin/main\n- rtk scalafmt --list --mode diff-ref=origin/main\n- rtk git diff --check\n- qodercli stdout review: No must-fix findings\n- subAgent review: No GraphInterpreter must-fix findings\n\nReferences:\nRefs #3030 --- .../impl/fusing/GraphInterpreterSpec.scala | 76 ++++--- .../stream/impl/fusing/GraphInterpreter.scala | 194 +++++++++++------- 2 files changed, 178 insertions(+), 92 deletions(-) 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 5ec839dc68e..4dc037f36f5 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,7 +315,7 @@ class GraphInterpreterSpec extends StreamSpec with GraphInterpreterSpecKit { interpreter.isSuspended should be(false) } - "release references to completed stage logics to prevent memory leaks" in new TestSetup { + "release references to stage logics when finishing the interpreter" in new TestSetup { val source = new UpstreamProbe[Int]("source") val sink = new DownstreamProbe[Int]("sink") val identityStage = GraphStages.identity[Int] @@ -333,49 +333,79 @@ class GraphInterpreterSpec extends StreamSpec with GraphInterpreterSpecKit { source.onNext(1) lastEvents() should ===(Set(OnNext(sink, 1))) - val logicsField = interpreter.getClass.getDeclaredField("logics") - logicsField.setAccessible(true) - val logics = logicsField.get(interpreter).asInstanceOf[Array[org.apache.pekko.stream.stage.GraphStageLogic]] + val logics = interpreter.logics + val connections = interpreter.connections logics.foreach(logic => logic should not be null) - - source.onComplete() - lastEvents() should ===(Set(OnComplete(sink))) + connections.foreach { connection => + connection.inOwner should not be null + connection.outOwner should not be null + connection.inHandler should not be null + connection.outHandler should not be null + } interpreter.finish() logics.foreach(logic => logic should be(null)) + connections.foreach { connection => + connection.inOwner should be(null) + connection.outOwner should be(null) + connection.inHandler should be(null) + connection.outHandler should be(null) + } + + val snapshot = interpreter.toSnapshot + snapshot.logics.map(_.label).toSet should ===(Set("")) + snapshot.connections should have size connections.length } - "release references to completed stage logics when some stages complete early" in new TestSetup { - val source = new UpstreamProbe[Int]("source") - val detachStage = detacher[Int] - val identityStage = GraphStages.identity[Int] + "release references to completed stage logics while the interpreter keeps running" in new TestSetup { + val source1 = new UpstreamProbe[Int]("source1") + val source2 = new UpstreamProbe[Int]("source2") + val mergeStage = Merge[Int](2) val sink = new DownstreamProbe[Int]("sink") - builder(detachStage, identityStage) - .connect(source, detachStage.shape.in) - .connect(detachStage.shape.out, identityStage.in) - .connect(identityStage.out, sink) + builder(mergeStage) + .connect(source1, mergeStage.in(0)) + .connect(source2, mergeStage.in(1)) + .connect(mergeStage.out, sink) .init() lastEvents() should ===(Set.empty[TestEvent]) sink.requestOne() - lastEvents() should ===(Set(RequestOne(source))) + lastEvents() should ===(Set(RequestOne(source1), RequestOne(source2))) - val logicsField = interpreter.getClass.getDeclaredField("logics") - logicsField.setAccessible(true) - val logics = logicsField.get(interpreter).asInstanceOf[Array[org.apache.pekko.stream.stage.GraphStageLogic]] + val logics = interpreter.logics + val connections = interpreter.connections logics.foreach(logic => logic should not be null) - source.onNext(1) - lastEvents() should ===(Set(OnNext(sink, 1), RequestOne(source))) + source1.onComplete() + lastEvents() should ===(Set.empty[TestEvent]) - source.onComplete() + logics(0) should be(null) + logics(1) should not be null + logics(2) should not be null + logics(3) should not be null - logics.foreach(logic => logic should be(null)) + connections(0).outOwner should be(null) + connections(0).outHandler should be(null) + connections(0).inOwner should not be null + connections(0).inHandler should not be null + + connections(1).inOwner should not be null + connections(1).outOwner should not be null + connections(1).inHandler should not be null + connections(1).outHandler should not be null + + val snapshot = interpreter.toSnapshot + snapshot.connections should have size connections.length + snapshot.connections.head.out.label should ===("") + snapshot.connections.head.in.label should not be "" + + source2.onNext(1) + lastEvents() should ===(Set(OnNext(sink, 1), RequestOne(source2))) } } 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 40d034c6c5f..6ddf49cb703 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,25 @@ import pekko.stream.stage._ // Marks whether a stage has been finalized (finalizeStage been called) or not private val finalizedMark = Array.fill(logics.length)(false) + 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) @@ -342,51 +361,106 @@ import pekko.stream.stage._ var i = 0 while (i < logics.length) { val logic = logics(i) - if ((logic ne null) && !isStageCompleted(logic) && !isStageFinalized(logic)) { - markStageFinalized(logic) - finalizeStage(logic) + if (logic ne null) { + if (!isStageFinalized(logic)) { + markStageFinalized(logic) + finalizeStage(logic) + } + logics(i) = null + } + i += 1 + } + releaseAllConnectionReferences() + } + + private def releaseAllConnectionReferences(): Unit = { + var i = 0 + while (i < connections.length) { + val connection = connections(i) + if (connection ne null) { + connection.inHandler = null + connection.outHandler = null + connection.inOwner = null + connection.outOwner = null + } + i += 1 + } + } + + private def logicName(stageId: Int): String = + if (stageId >= 0) { + val logic = logics(stageId) + if (logic ne null) logic.toString else "" + } 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 releaseConnectionReferences(logic: GraphStageLogic): Unit = { + val stageId = logic.stageId + val portToConn = logic.portToConn + var i = 0 + while (i < portToConn.length) { + val connection = portToConn(i) + if (connection ne null) { + val connectionId = connection.id + if (connectionId >= 0 && connectionId < connectionInStageIds.length) { + if (connectionInStageIds(connectionId) == stageId) { + connection.inHandler = null + connection.inOwner = null + } + if (connectionOutStageIds(connectionId) == stageId) { + connection.outHandler = null + connection.outOwner = null + } + } } - logics(i) = null i += 1 } - 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 + } + + private def finalizeCompletedStage(logic: GraphStageLogic): Unit = { + markStageFinalized(logic) + runningStages -= 1 + finalizeStage(logic) + releaseConnectionReferences(logic) + logics(logic.stageId) = null + } + + private def finalizeCompletedStages(): Unit = { + var i = 0 + while (i < logics.length) { + val logic = logics(i) + if ((logic ne null) && isStageCompleted(logic) && !isStageFinalized(logic)) { + finalizeCompletedStage(logic) } - j += 1 + i += 1 } } // Debug name for a connections input part private def inOwnerName(connection: Connection): String = - if (connection.inOwner ne null) connection.inOwner.toString else "" + if (connection.inOwner ne null) connection.inOwner.toString else logicName(connectionInStageId(connection)) // Debug name for a connections output part private def outOwnerName(connection: Connection): String = - if (connection.outOwner ne null) connection.outOwner.toString else "" + if (connection.outOwner ne null) connection.outOwner.toString else logicName(connectionOutStageId(connection)) // Debug name for a connections input part - private def inLogicName(connection: Connection): String = { - val owner = connection.inOwner - if (owner ne null) { - val logic = logics(owner.stageId) - if (logic ne null) logic.toString else "" - } else "" - } + private def inLogicName(connection: Connection): String = + logicName(connectionInStageId(connection)) // Debug name for a connections output part - private def outLogicName(connection: Connection): String = { - val owner = connection.outOwner - if (owner ne null) { - val logic = logics(owner.stageId) - if (logic ne null) logic.toString else "" - } else "" - } + 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(",") @@ -461,7 +535,6 @@ import pekko.stream.stage._ case NonFatal(e) => reportStageError(e) } if (pendingFinalization) { - pendingFinalization = false afterStageHasRun(activeStage) } @@ -497,7 +570,6 @@ import pekko.stream.stage._ case NonFatal(e) => reportStageError(e) } if (pendingFinalization) { - pendingFinalization = false afterStageHasRun(activeStage) } } @@ -511,7 +583,6 @@ import pekko.stream.stage._ case NonFatal(e) => reportStageError(e) } if (pendingFinalization) { - pendingFinalization = false afterStageHasRun(activeStage) } } @@ -585,7 +656,7 @@ import pekko.stream.stage._ s"$Name CANCEL ${inOwnerName(connection)} -> ${outOwnerName( connection)} (${connection.outHandler}) [${outLogicName(connection)}]") connection.portState |= OutClosed - completeConnection(connection.outOwner.stageId) + completeConnection(connectionOutStageId(connection)) val cause = connection.slot.asInstanceOf[Cancelled].cause connection.slot = Empty connection.outHandler.onDownstreamFinish(cause) @@ -600,7 +671,7 @@ import pekko.stream.stage._ connection)} (${connection.inHandler}) [${inLogicName(connection)}]") connection.portState |= InClosed activeStage = connection.inOwner - completeConnection(connection.inOwner.stageId) + completeConnection(connectionInStageId(connection)) if ((connection.portState & InFailed) == 0) connection.inHandler.onUpstreamFinish() else connection.inHandler.onUpstreamFailure(connection.slot.asInstanceOf[Failed].ex) } else { @@ -656,31 +727,13 @@ import pekko.stream.stage._ queueTail += 1 } - @scala.annotation.nowarn("cat=unused-params") def afterStageHasRun(logic: GraphStageLogic): Unit = { - var i = 0 - while (i < logics.length) { - val l = logics(i) - if ((l ne null) && isStageCompleted(l) && !isStageFinalized(l)) { - markStageFinalized(l) - runningStages -= 1 - finalizeStage(l) - logics(i) = null - } - i += 1 + if ((logic ne null) && isStageCompleted(logic) && !isStageFinalized(logic)) { + finalizeCompletedStage(logic) } - var j = 0 - while (j < connections.length) { - val conn = connections(j) - if (conn ne null) { - if (finalizedMark(conn.inOwner.stageId)) { - conn.inHandler = null - } - if (finalizedMark(conn.outOwner.stageId)) { - conn.outHandler = null - } - } - j += 1 + if (pendingFinalization) { + pendingFinalization = false + finalizeCompletedStages() } } @@ -696,11 +749,13 @@ import pekko.stream.stage._ // Register that a connection in which the given stage participated has been completed and therefore the stage // itself might stop, too. private def completeConnection(stageId: Int): Unit = { - val activeConnections = shutdownCounter(stageId) - if (activeConnections > 0) { - val next = activeConnections - 1 - shutdownCounter(stageId) = next - if (next == 0) pendingFinalization = true + if (stageId >= 0) { + val activeConnections = shutdownCounter(stageId) + if (activeConnections > 0) { + val next = activeConnections - 1 + shutdownCounter(stageId) = next + if (next == 0) pendingFinalization = true + } } } @@ -748,7 +803,7 @@ 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) completeConnection(connectionOutStageId(connection)) } @InternalStableApi @@ -767,7 +822,7 @@ import pekko.stream.stage._ enqueue(connection) } } - if ((currentState & OutClosed) == 0) completeConnection(connection.outOwner.stageId) + if ((currentState & OutClosed) == 0) completeConnection(connectionOutStageId(connection)) } @InternalStableApi @@ -785,7 +840,7 @@ import pekko.stream.stage._ enqueue(connection) } } - if ((currentState & InClosed) == 0) completeConnection(connection.inOwner.stageId) + if ((currentState & InClosed) == 0) completeConnection(connectionInStageId(connection)) } /** @@ -803,12 +858,13 @@ import pekko.stream.stage._ case (_, idx) => LogicSnapshotImpl(idx, "", Attributes.none) } - val connectionSnapshots = connections.filter(c => (c ne null) && (c.inOwner ne null) && (c.outOwner ne null)).map { - connection => + val connectionSnapshots = connections.zipWithIndex.collect { + case (connection, idx) + if (connection ne null) && connectionInStageIds(idx) >= 0 && connectionOutStageIds(idx) >= 0 => ConnectionSnapshotImpl( connection.id, - logicSnapshots(connection.inOwner.stageId), - logicSnapshots(connection.outOwner.stageId), + logicSnapshots(connectionInStageIds(idx)), + logicSnapshots(connectionOutStageIds(idx)), connection.portState match { case InReady | Pushing => ConnectionSnapshot.ShouldPull case OutReady | Pulling => ConnectionSnapshot.ShouldPush