PR comments.

Author: suhsteve

src/csharp/Microsoft.Spark/Interop/Ipc/CallbackServer.cs

@@ -91,15 +91,6 @@ internal int RegisterCallback(ICallbackHandler callbackHandler)
             return callbackId;
         }
 
-        /// <summary>
-        /// Runs the callback server.
-        /// </summary>
-        internal void Run()
-        {
-            _listener = SocketFactory.CreateSocket();
-            Run(_listener);
-        }
-
         /// <summary>
         /// Runs the callback server.
         /// </summary>
@@ -117,10 +108,11 @@ internal void Run(ISocketWrapper listener)
 
             try
             {
-                listener.Listen();
+                _listener = listener;
+                _listener.Listen();
 
                 // Communicate with the JVM the callback server's address and port.
-                var localEndPoint = (IPEndPoint)listener.LocalEndPoint;
+                var localEndPoint = (IPEndPoint)_listener.LocalEndPoint;
                 _jvm.CallStaticJavaMethod(
                     "DotnetHandler",
                     "connectCallback",
@@ -130,7 +122,7 @@ internal void Run(ISocketWrapper listener)
                 s_logger.LogInfo($"Started CallbackServer on {localEndPoint}");
 
                 // Start accepting connections from JVM.
-                new Thread(() => StartServer(listener))
+                new Thread(() => StartServer(_listener))
                 {
                     IsBackground = true
                 }.Start();
@@ -142,6 +134,14 @@ internal void Run(ISocketWrapper listener)
             }
         }
 
+        /// <summary>
+        /// Runs the callback server.
+        /// </summary>
+        private void Run()
+        {
+            Run(SocketFactory.CreateSocket());
+        }
+
         /// <summary>
         /// Starts listening to any connection from JVM.
         /// </summary>

src/scala/microsoft-spark-2.4.x/src/main/scala/org/apache/spark/api/dotnet/CallbackClient.scala

@@ -31,12 +31,7 @@ class CallbackClient(address: String, port: Int) extends Logging {
     getOrCreateConnection() match {
       case Some(connection) =>
         try {
-          connection.send(callbackId, writeBody) match {
-            case ConnectionStatus.SUCCESS =>
-              addConnection(connection)
-            case status =>
-              throw new Exception(s"Error encountered with connection: '$status'")
-          }
+          connection.send(callbackId, writeBody)
         } catch {
           case e: Exception =>
             logError(s"Error calling callback [callback id = $callbackId].", e)

src/scala/microsoft-spark-2.4.x/src/main/scala/org/apache/spark/api/dotnet/CallbackConnection.scala

@@ -25,7 +25,7 @@ class CallbackConnection(address: String, port: Int) extends Logging {
 
   def send(
       callbackId: Int,
-      writeBody: DataOutputStream => Unit): ConnectionStatus.ConnectionStatus = {
+      writeBody: DataOutputStream => Unit): Unit = {
     logInfo(s"Calling callback [callback id = $callbackId] ...")
 
     try {
@@ -38,8 +38,7 @@ class CallbackConnection(address: String, port: Int) extends Logging {
       byteArrayOutputStream.writeTo(outputStream);
     } catch {
       case e: Exception => {
-        logError("Error writing to stream.", e)
-        return ConnectionStatus.ERROR_WRITE
+        throw new Exception("Error writing to stream.", e)
       }
     }
 
@@ -52,19 +51,16 @@ class CallbackConnection(address: String, port: Int) extends Logging {
       endOfStreamResponse match {
         case CallbackFlags.END_OF_STREAM =>
           logInfo(s"Received END_OF_STREAM signal. Calling callback [callback id = $callbackId] successful.")
-          return ConnectionStatus.SUCCESS
         case _ =>  {
-          logError(s"Error verifying end of stream. Expected: ${CallbackFlags.END_OF_STREAM}, " +
+          throw new Exception(s"Error verifying end of stream. Expected: ${CallbackFlags.END_OF_STREAM}, " +
               s"Received: $endOfStreamResponse")
         }
       }
     } catch {
       case e: Exception => {
-        logError("Error while verifying end of stream.", e)
+        throw new Exception("Error while verifying end of stream.", e)
       }
     }
-
-    ConnectionStatus.ERROR_END_OF_STREAM
   }
 
   def close(): Unit = {
@@ -113,9 +109,4 @@ class CallbackConnection(address: String, port: Int) extends Logging {
     val DOTNET_EXCEPTION_THROWN: Int = -3
     val END_OF_STREAM: Int = -4
   }
-}
-
-object ConnectionStatus extends Enumeration {
-  type ConnectionStatus = Value
-  val SUCCESS, ERROR_WRITE, ERROR_END_OF_STREAM = Value
-}
+}

src/scala/microsoft-spark-2.4.x/src/main/scala/org/apache/spark/api/dotnet/DotnetBackend.scala

@@ -83,10 +83,19 @@ class DotnetBackend extends Logging {
   }
 }
 
-object DotnetBackend {
+object DotnetBackend extends Logging {
   @volatile private[spark] var callbackClient: CallbackClient = null
 
-  private[spark] def shutdownCallbackClient(): Unit = {
+  private[spark] def setCallbackClient(address: String, port: Int) = synchronized {
+    if (DotnetBackend.callbackClient == null) {
+      logInfo(s"Connecting to a callback server at $address:$port")
+      DotnetBackend.callbackClient = new CallbackClient(address, port)
+    } else {
+      throw new Exception("Callback client already set.")
+    }
+  }
+
+  private[spark] def shutdownCallbackClient(): Unit = synchronized {
     if (callbackClient != null) {
       callbackClient.shutdown()
       callbackClient = null

src/scala/microsoft-spark-2.4.x/src/main/scala/org/apache/spark/api/dotnet/DotnetBackendHandler.scala

@@ -70,12 +70,7 @@ class DotnetBackendHandler(server: DotnetBackend)
           val address = readString(dis)
           assert(readObjectType(dis) == 'i')
           val port = readInt(dis)
-          if (DotnetBackend.callbackClient == null) {
-            logInfo(s"Connecting to a callback server at $address:$port")
-            DotnetBackend.callbackClient = new CallbackClient(address, port)
-          } else {
-            throw new Exception("Callback client already set.")
-          }
+          DotnetBackend.setCallbackClient(address, port);
           writeInt(dos, 0)
           writeType(dos, "void")
         case "closeCallback" =>