fix: release references to completed stage logics in GraphInterpreter

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 {

stream-tests/src/test/scala/org/apache/pekko/stream/impl/fusing/GraphInterpreterSpec.scala

@@ -315,6 +315,99 @@ class GraphInterpreterSpec extends StreamSpec with GraphInterpreterSpecKit {
       interpreter.isSuspended should be(false)
     }
 
+    "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]
+
+      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 logics = interpreter.logics
+      val connections = interpreter.connections
+
+      logics.foreach(logic => logic should not be null)
+      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("<completed>"))
+      snapshot.connections should have size connections.length
+    }
+
+    "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(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(source1), RequestOne(source2)))
+
+      val logics = interpreter.logics
+      val connections = interpreter.connections
+
+      logics.foreach(logic => logic should not be null)
+
+      source1.onComplete()
+      lastEvents() should ===(Set.empty[TestEvent])
+
+      logics(0) should be(null)
+      logics(1) should not be null
+      logics(2) should not be null
+      logics(3) should not 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 ===("<completed>")
+      snapshot.connections.head.in.label should not be "<completed>"
+
+      source2.onNext(1)
+      lastEvents() should ===(Set(OnNext(sink, 1), RequestOne(source2)))
+    }
+
   }
 
 }