do retry on client side

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

@@ -613,16 +613,7 @@ private[spark] object PythonRDD extends Logging {
       setDaemon(true)
       override def run() {
         try {
-          var sock: Socket = null
-          try {
-            sock = serverSocket.accept()
-          } catch {
-            case e: SocketTimeoutException =>
-              // there is a small chance that the client had connected, so retry
-              logWarning("Timed out after 4 seconds, retry once")
-              serverSocket.setSoTimeout(10)
-              sock = serverSocket.accept()
-          }
+          val sock = serverSocket.accept()
           val out = new DataOutputStream(new BufferedOutputStream(sock.getOutputStream))
           try {
             writeIteratorToStream(items, out)

python/pyspark/rdd.py

@@ -113,11 +113,23 @@ def _parse_memory(s):
 
 def _load_from_socket(port, serializer):
     sock = socket.socket()
-    sock.settimeout(5)
+    sock.settimeout(1)
     try:
         sock.connect(("localhost", port))
         rf = sock.makefile("rb", 65536)
-        for item in serializer.load_stream(rf):
+        iter = serializer.load_stream(rf)
+        try:
+            yield next(iter)
+        except socket.timeout as e:
+            # the connection is not acknowledged by JVM, retry
+            # server will be closed after 3 seconds, then it will be refused
+            for v in _load_from_socket(port, serializer):
+                yield v
+            return
+
+        # increase the timeout, because the server side may be slowed down by GC
+        sock.settimeout(10)
+        for item in iter:
             yield item
     finally:
         sock.close()