only reuse completed and not interrupted worker

Author: davies

core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala

@@ -75,7 +75,16 @@ private[spark] class PythonRDD(
 
     context.addTaskCompletionListener { context =>
       writerThread.shutdownOnTaskCompletion()
-      env.releasePythonWorker(pythonExec, envVars.toMap, worker)
+      if (!context.isInterrupted) {
+        env.releasePythonWorker(pythonExec, envVars.toMap, worker)
+      } else {
+        try {
+          worker.close()
+        } catch {
+          case e: Exception =>
+            logWarning("Failed to close worker socket", e)
+        }
+      }
     }
 
     writerThread.start()

core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala

@@ -239,6 +239,15 @@ private[spark] class PythonWorkerFactory(pythonExec: String, envVars: Map[String
   private def stopDaemon() {
     synchronized {
       if (useDaemon) {
+        while (idleWorkers.length > 0) {
+          val worker = idleWorkers.dequeue()
+          try {
+            worker.close()
+          } catch {
+            case e: Exception =>
+              logWarning("Failed to close worker socket", e)
+          }
+        }
         // Request shutdown of existing daemon by sending SIGTERM
         if (daemon != null) {
           daemon.destroy()