fix: release references to completed stage logics in GraphInterpreter

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 ===("<completed>")
+      partiallyReleasedConnectionSnapshot.in.label should not be "<completed>"
+
+      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 {

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("<completed>"))
+      snapshot.connections should have size interpreter.connections.count(_ ne null)
+      snapshot.connections.foreach { connection =>
+        connection.in.label should ===("<completed>")
+        connection.out.label should ===("<completed>")
+      }
+    }
+
+    "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 ===("<completed>")
+      connectionSnapshot.out.label should ===("<completed>")
+    }
+
+    "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)
+      }
+    }
+
   }
 
 }

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 =