shuffle client/server: enable reply flag in connect + cb for ep disco…

src/main/scala/org/apache/spark/shuffle/ucx/external/client/ExternalUcxClientWorker.scala

@@ -308,8 +308,7 @@ case class ExternalUcxClientWorker(val worker: UcpWorker,
   @`inline`
   def requestAddress(localServer: InetSocketAddress): Unit = {
     executor.post(() => shuffleServers.computeIfAbsent("0.0.0.0", _ => {
-      doConnect(localServer, ExternalAmId.ADDRESS,
-                UcpConstants.UCP_AM_SEND_FLAG_REPLY)._1
+      doConnect(localServer, ExternalAmId.ADDRESS)._1
     }))
   }
 
@@ -331,8 +330,7 @@ case class ExternalUcxClientWorker(val worker: UcpWorker,
   }
 
   private def doConnect(shuffleServer: InetSocketAddress,
-                        amId: Int = ExternalAmId.CONNECT,
-                        flag: Long = 0): (UcpEndpoint, UcpRequest) = {
+                        amId: Int): (UcpEndpoint, UcpRequest) = {
     val endpointParams = new UcpEndpointParams().setPeerErrorHandlingMode()
       .setSocketAddress(shuffleServer).sendClientId()
       .setErrorHandler(new UcpEndpointErrorHandler() {
@@ -353,7 +351,8 @@ case class ExternalUcxClientWorker(val worker: UcpWorker,
     val req = ep.sendAmNonBlocking(
       amId, UcxUtils.getAddress(header), header.remaining(),
       UcxUtils.getAddress(workerAddress), workerAddress.remaining(),
-      UcpConstants.UCP_AM_SEND_FLAG_EAGER | flag, new UcxCallback() {
+      UcpConstants.UCP_AM_SEND_FLAG_EAGER | UcpConstants.UCP_AM_SEND_FLAG_REPLY,
+      new UcxCallback() {
         override def onSuccess(request: UcpRequest): Unit = {
           connectNext()
         }
@@ -368,7 +367,8 @@ case class ExternalUcxClientWorker(val worker: UcpWorker,
   }
 
   private def startConnection(shuffleServer: InetSocketAddress): (UcpEndpoint, UcpRequest) = {
-    connectingServers.computeIfAbsent(shuffleServer, _ => doConnect(shuffleServer))
+    connectingServers.computeIfAbsent(shuffleServer, _ =>
+      doConnect(shuffleServer, ExternalAmId.CONNECT))
   }
 
   private def getConnection(host: String): UcpEndpoint = {

src/main/scala/org/apache/spark/shuffle/ucx/external/server/ExternalUcxServerTransport.scala

@@ -117,7 +117,7 @@ class ExternalUcxServerTransport(
   }
 
   def handleConnect(handler: ExternalUcxServerWorker,
-                    clientWorker: UcxWorkerId, address: ByteBuffer): Unit = {
+                    clientWorker: UcxWorkerId): Unit = {
     workerMap.getOrElseUpdate(clientWorker.appId, {
       new TrieMap[UcxWorkerId, Unit]
     }).getOrElseUpdate(clientWorker, Unit)

src/main/scala/org/apache/spark/shuffle/ucx/external/server/ExternalUcxServerWorker.scala

@@ -33,7 +33,8 @@ case class ExternalUcxServerWorker(val worker: UcpWorker,
   private[ucx] val executor = new UcxWorkerThread(
     worker, transport.ucxShuffleConf.useWakeup)
 
-  private val endpoints = mutable.Set.empty[UcpEndpoint]
+  private val emptyCallback = () => {}
+  private val endpoints = mutable.HashMap.empty[UcpEndpoint, () => Unit]
   private val listener = worker.newListener(
     new UcpListenerParams().setSockAddr(new InetSocketAddress("0.0.0.0", port))
     .setConnectionHandler((ucpConnectionRequest: UcpConnectionRequest) => {
@@ -43,9 +44,9 @@ case class ExternalUcxServerWorker(val worker: UcpWorker,
           new UcpEndpointParams().setConnectionRequest(ucpConnectionRequest)
             .setPeerErrorHandlingMode().setErrorHandler(errorHandler)
             .setName(s"Endpoint to $clientAddress"))
-        endpoints.add(ep)
+        endpoints.getOrElseUpdate(ep, emptyCallback)
       } catch {
-        case e: UcxException => logWarning(s"Accept $clientAddress fail: $e")
+        case e: UcxException => logError(s"Accept $clientAddress fail: $e")
       }
     }))
 
@@ -57,7 +58,7 @@ case class ExternalUcxServerWorker(val worker: UcpWorker,
       } else {
         logWarning(s"Ep $ucpEndpoint got an error: $errorString")
       }
-      endpoints.remove(ucpEndpoint)
+      endpoints.remove(ucpEndpoint).foreach(_())
       ucpEndpoint.close()
     }
   }
@@ -94,15 +95,15 @@ case class ExternalUcxServerWorker(val worker: UcpWorker,
 
   // AM to get worker address for client worker and connect server workers to it
   worker.setAmRecvHandler(ExternalAmId.CONNECT,
-    (headerAddress: Long, headerSize: Long, amData: UcpAmData, _: UcpEndpoint) => {
+    (headerAddress: Long, headerSize: Long, amData: UcpAmData, ep: UcpEndpoint) => {
     val header = UnsafeUtils.getByteBufferView(headerAddress, headerSize.toInt)
     val shuffleClient = UcxWorkerId.deserialize(header)
     val workerAddress = UnsafeUtils.getByteBufferView(amData.getDataAddress,
                                                       amData.getLength.toInt)
     val copiedAddress = ByteBuffer.allocateDirect(workerAddress.remaining)
     copiedAddress.put(workerAddress)
     connected(shuffleClient, copiedAddress)
-    transport.handleConnect(this, shuffleClient, copiedAddress)
+    endpoints.put(ep, () => doDisconnect(shuffleClient))
     UcsConstants.STATUS.UCS_OK
   }, UcpConstants.UCP_AM_FLAG_WHOLE_MSG )
   // Main RPC thread. reply with ucpAddress.
@@ -131,7 +132,7 @@ case class ExternalUcxServerWorker(val worker: UcpWorker,
     }
     if (!endpoints.isEmpty) {
       logInfo(s"$workerId closing ${endpoints.size} eps")
-      endpoints.map(
+      endpoints.keys.map(
         _.closeNonBlockingForce()).foreach(req =>
           while (!req.isCompleted){
             worker.progress()
@@ -158,20 +159,21 @@ case class ExternalUcxServerWorker(val worker: UcpWorker,
   }
 
   @`inline`
-  private def doDisconnect(workerId: UcxWorkerId): Unit = {
+  private def doDisconnect(shuffleClient: UcxWorkerId): Unit = {
     try {
-      workerReqs.remove(workerId).foreach(reqs => {
+      workerReqs.remove(shuffleClient).foreach(reqs => {
         val inCompletes = reqs.filterNot(_.isCompleted)
         if (inCompletes.nonEmpty) {
           inCompletes.foreach(worker.cancelRequest(_))
-          logInfo(s"$workerId canceled ${inCompletes.size} requests")
+          logInfo(s"Canceled ${inCompletes.size} requests to $shuffleClient")
         }
       })
-      Option(shuffleClients.remove(workerId)).foreach(ep => {
+      Option(shuffleClients.remove(shuffleClient)).foreach(ep => {
         ep.closeNonBlockingFlush()
+        logInfo(s"Disconnect $shuffleClient")
       })
     } catch {
-      case e: Throwable => logWarning(s"$workerId disconnect $e")
+      case e: Throwable => logWarning(s"Disconnect $shuffleClient: $e")
     }
   }
 
@@ -188,6 +190,7 @@ case class ExternalUcxServerWorker(val worker: UcpWorker,
         worker.newEndpoint(new UcpEndpointParams()
           .setErrorHandler(new UcpEndpointErrorHandler {
             override def onError(ucpEndpoint: UcpEndpoint, errorCode: Int, errorString: String): Unit = {
+              logInfo(s"Connection to $shuffleClient closed: $errorString")
               shuffleClients.remove(shuffleClient)
               workerReqs.remove(shuffleClient)
               ucpEndpoint.close()
@@ -197,7 +200,7 @@ case class ExternalUcxServerWorker(val worker: UcpWorker,
           .setUcpAddress(workerAddress))
       })
     } catch {
-      case e: UcxException => logWarning(s"Connect back to $shuffleClient fail: $e")
+      case e: UcxException => logWarning(s"Connection to $shuffleClient failed: $e")
     }
   }
 
@@ -220,7 +223,7 @@ case class ExternalUcxServerWorker(val worker: UcpWorker,
       new mutable.ListBuffer[UcpRequest]()
     })
 
-    while (reqs.nonEmpty && reqs(0).isCompleted) {
+    while (reqs.nonEmpty && reqs.head.isCompleted) {
       reqs.remove(0)
     }
     reqs.append(req)