[SPARK-33143][PYTHON] Add configurable timeout to python server and client

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

@@ -85,6 +85,7 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
 
   private val conf = SparkEnv.get.conf
   protected val bufferSize: Int = conf.get(BUFFER_SIZE)
+  protected val gatewayConnectTimeout = conf.get(PYTHON_GATEWAY_CONNECT_TIMEOUT)
   private val reuseWorker = conf.get(PYTHON_WORKER_REUSE)
 
   // All the Python functions should have the same exec, version and envvars.
@@ -140,6 +141,7 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
     if (workerMemoryMb.isDefined) {
       envVars.put("PYSPARK_EXECUTOR_MEMORY_MB", workerMemoryMb.get.toString)
     }
+    envVars.put("SPARK_GATEWAY_CONNECT_TIMEOUT", gatewayConnectTimeout.toString)
     envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
     val worker: Socket = env.createPythonWorker(pythonExec, envVars.asScala.toMap)
     // Whether is the worker released into idle pool or closed. When any codes try to release or

core/src/main/scala/org/apache/spark/internal/config/Python.scala

@@ -50,4 +50,9 @@ private[spark] object Python {
     .version("2.4.0")
     .bytesConf(ByteUnit.MiB)
     .createOptional
+
+  val PYTHON_GATEWAY_CONNECT_TIMEOUT = ConfigBuilder("spark.python.gateway.connectTimeout")
+    .version("3.1.0")
+    .timeConf(TimeUnit.SECONDS)
+    .createWithDefaultString("15s")
 }

core/src/main/scala/org/apache/spark/security/SocketAuthServer.scala

@@ -25,6 +25,8 @@ import scala.concurrent.duration.Duration
 import scala.util.Try
 
 import org.apache.spark.SparkEnv
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config.Python.PYTHON_GATEWAY_CONNECT_TIMEOUT
 import org.apache.spark.network.util.JavaUtils
 import org.apache.spark.util.{ThreadUtils, Utils}
 
@@ -34,29 +36,35 @@ import org.apache.spark.util.{ThreadUtils, Utils}
  * handling one batch of data, with authentication and error handling.
  *
  * The socket server can only accept one connection, or close if no connection
- * in 15 seconds.
+ * in configurable amount of seconds (default 15).
  */
 private[spark] abstract class SocketAuthServer[T](
     authHelper: SocketAuthHelper,
-    threadName: String) {
+    threadName: String) extends Logging {
 
   def this(env: SparkEnv, threadName: String) = this(new SocketAuthHelper(env.conf), threadName)
   def this(threadName: String) = this(SparkEnv.get, threadName)
 
   private val promise = Promise[T]()
 
   private def startServer(): (Int, String) = {
+    logTrace("Creating listening socket")
     val serverSocket = new ServerSocket(0, 1, InetAddress.getByAddress(Array(127, 0, 0, 1)))
-    // Close the socket if no connection in 15 seconds
-    serverSocket.setSoTimeout(15000)
+    // Close the socket if no connection in configured seconds
+    val timeout = SparkEnv.get.conf.get(PYTHON_GATEWAY_CONNECT_TIMEOUT).toInt
+    logTrace(s"Setting timeout to $timeout sec")
+    serverSocket.setSoTimeout(timeout * 1000)
 
     new Thread(threadName) {
       setDaemon(true)
       override def run(): Unit = {
         var sock: Socket = null
         try {
+          logTrace(s"Waiting for connection on port ${serverSocket.getLocalPort}")
           sock = serverSocket.accept()
+          logTrace(s"Connection accepted from port ${sock.getLocalPort}")
           authHelper.authClient(sock)
+          logTrace("Client authenticated")
           promise.complete(Try(handleConnection(sock)))
         } finally {
           JavaUtils.closeQuietly(serverSocket)

core/src/test/scala/org/apache/spark/api/python/PythonRDDSuite.scala

@@ -30,15 +30,18 @@ import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.io.{LongWritable, Text}
 import org.apache.hadoop.mapred.TextInputFormat
 import org.apache.hadoop.mapreduce.lib.input.FileInputFormat
+import org.mockito.Mockito.mock
 
-import org.apache.spark.{LocalSparkContext, SparkConf, SparkContext, SparkFunSuite}
+import org.apache.spark.{LocalSparkContext, SparkConf, SparkContext, SparkEnv, SparkFunSuite}
 import org.apache.spark.api.java.JavaSparkContext
 import org.apache.spark.rdd.{HadoopRDD, RDD}
 import org.apache.spark.security.{SocketAuthHelper, SocketAuthServer}
 import org.apache.spark.util.Utils
 
 class PythonRDDSuite extends SparkFunSuite with LocalSparkContext {
 
+  private def doReturn(value: Any) = org.mockito.Mockito.doReturn(value, Seq.empty: _*)
+
   var tempDir: File = _
 
   override def beforeAll(): Unit = {
@@ -76,12 +79,22 @@ class PythonRDDSuite extends SparkFunSuite with LocalSparkContext {
   }
 
   test("python server error handling") {
-    val authHelper = new SocketAuthHelper(new SparkConf())
-    val errorServer = new ExceptionPythonServer(authHelper)
-    val client = new Socket(InetAddress.getLoopbackAddress(), errorServer.port)
-    authHelper.authToServer(client)
-    val ex = intercept[Exception] { errorServer.getResult(Duration(1, "second")) }
-    assert(ex.getCause().getMessage().contains("exception within handleConnection"))
+    val savedSparkEnv = SparkEnv.get
+    try {
+      val conf = new SparkConf()
+      val env = mock(classOf[SparkEnv])
+      doReturn(conf).when(env).conf
+      SparkEnv.set(env)
+
+      val authHelper = new SocketAuthHelper(conf)
+      val errorServer = new ExceptionPythonServer(authHelper)
+      val client = new Socket(InetAddress.getLoopbackAddress(), errorServer.port)
+      authHelper.authToServer(client)
+      val ex = intercept[Exception] { errorServer.getResult(Duration(1, "second")) }
+      assert(ex.getCause().getMessage().contains("exception within handleConnection"))
+    } finally {
+      SparkEnv.set(savedSparkEnv)
+    }
   }
 
   class ExceptionPythonServer(authHelper: SocketAuthHelper)

python/pyspark/java_gateway.py

@@ -201,7 +201,7 @@ def local_connect_and_auth(port, auth_secret):
         af, socktype, proto, _, sa = res
         try:
             sock = socket.socket(af, socktype, proto)
-            sock.settimeout(15)
+            sock.settimeout(int(os.environ.get("SPARK_GATEWAY_CONNECT_TIMEOUT", 15)))
             sock.connect(sa)
             sockfile = sock.makefile("rwb", int(os.environ.get("SPARK_BUFFER_SIZE", 65536)))
             _do_server_auth(sockfile, auth_secret)