+str AsyncCallbacks lost on finished stage by error

akka-remote/src/main/scala/akka/remote/artery/Control.scala

@@ -73,7 +73,7 @@ private[remote] object InboundControlJunction {
     def notify(inboundEnvelope: InboundEnvelope): Unit
   }
 
-  // messages for the CallbackWrapper
+  // messages for the stream callback
   private[InboundControlJunction] sealed trait CallbackMessage
   private[InboundControlJunction] final case class Attach(observer: ControlMessageObserver, done: Promise[Done])
     extends CallbackMessage
@@ -93,8 +93,7 @@ private[remote] class InboundControlJunction
 
   override def createLogicAndMaterializedValue(inheritedAttributes: Attributes) = {
     val stoppedPromise = Promise[Done]()
-    // FIXME see issue #20503 related to CallbackWrapper, we might implement this in a better way
-    val logic = new GraphStageLogic(shape) with CallbackWrapper[CallbackMessage] with InHandler with OutHandler {
+    val logic = new GraphStageLogic(shape) with InHandler with OutHandler with ControlMessageSubject {
 
       private var observers: Vector[ControlMessageObserver] = Vector.empty
 
@@ -106,10 +105,6 @@ private[remote] class InboundControlJunction
           observers = observers.filterNot(_ == observer)
       }
 
-      override def preStart(): Unit = {
-        initCallback(callback.invoke)
-      }
-
       override def postStop(): Unit = stoppedPromise.success(Done)
 
       // InHandler
@@ -127,24 +122,22 @@ private[remote] class InboundControlJunction
       override def onPull(): Unit = pull(in)
 
       setHandlers(in, out, this)
-    }
 
-    // materialized value
-    val controlSubject: ControlMessageSubject = new ControlMessageSubject {
+      // ControlMessageSubject impl
       override def attach(observer: ControlMessageObserver): Future[Done] = {
         val p = Promise[Done]()
-        logic.invoke(Attach(observer, p))
+        callback.invoke(Attach(observer, p))
         p.future
       }
 
       override def detach(observer: ControlMessageObserver): Unit =
-        logic.invoke(Dettach(observer))
+        callback.invoke(Dettach(observer))
 
       override def stopped: Future[Done] =
         stoppedPromise.future
     }
 
-    (logic, controlSubject)
+    (logic, logic)
   }
 }
 
@@ -169,18 +162,14 @@ private[remote] class OutboundControlJunction(
   override val shape: FlowShape[OutboundEnvelope, OutboundEnvelope] = FlowShape(in, out)
 
   override def createLogicAndMaterializedValue(inheritedAttributes: Attributes) = {
-    // FIXME see issue #20503 related to CallbackWrapper, we might implement this in a better way
-    val logic = new GraphStageLogic(shape) with CallbackWrapper[ControlMessage] with InHandler with OutHandler with StageLogging {
+
+    val logic = new GraphStageLogic(shape) with InHandler with OutHandler with StageLogging with OutboundControlIngress {
       import OutboundControlJunction._
 
-      private val sendControlMessageCallback = getAsyncCallback[ControlMessage](internalSendControlMessage)
+      val sendControlMessageCallback = getAsyncCallback[ControlMessage](internalSendControlMessage)
       private val maxControlMessageBufferSize: Int = outboundContext.settings.Advanced.OutboundControlQueueSize
       private val buffer = new ArrayDeque[OutboundEnvelope]
 
-      override def preStart(): Unit = {
-        initCallback(sendControlMessageCallback.invoke)
-      }
-
       // InHandler
       override def onPush(): Unit = {
         if (buffer.isEmpty && isAvailable(out))
@@ -212,16 +201,13 @@ private[remote] class OutboundControlJunction(
         outboundEnvelopePool.acquire().init(
           recipient = OptionVal.None, message = message, sender = OptionVal.None)
 
+      override def sendControlMessage(message: ControlMessage): Unit =
+        sendControlMessageCallback.invoke(message)
+
       setHandlers(in, out, this)
     }
 
-    // materialized value
-    val outboundControlIngress = new OutboundControlIngress {
-      override def sendControlMessage(message: ControlMessage): Unit =
-        logic.invoke(message)
-    }
-
-    (logic, outboundControlIngress)
+    (logic, logic)
   }
 
 }