(spark) branch branch-4.3 updated: [SPARK-58932][SS] Close the accepted socket in TransformWithStateInPySparkStateServer

[ka...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

HeartSaVioR pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.3 by this push:
     new 818c2c861432 [SPARK-58932][SS] Close the accepted socket in TransformWithStateInPySparkStateServer
818c2c861432 is described below

commit 818c2c8614321a670aed23d6374b7bfa6052e280
Author: Jiwon Park <jp...@outlook.kr>
AuthorDate: Wed Aug 26 16:49:54 2026 +0900

    [SPARK-58932][SS] Close the accepted socket in TransformWithStateInPySparkStateServer
    
    ### What changes were proposed in this pull request?
    
    Run the request loop of `TransformWithStateInPySparkStateServer.run()` inside `Utils.tryWithResource` so the socket returned by `stateServerSocket.accept()` is always closed.
    
    ### Why are the changes needed?
    
    `run()` accepts a connection from the Python worker but never closes it. The task completion listener closes only the listening `ServerSocketChannel`, so the accepted connection's file descriptor is held until the executor exits.
    
    `initStateServer()` binds an ephemeral port per task, so a long-running streaming query leaks one socket per task. Once the ephemeral port range is exhausted, every new connection on that executor fails with `java.net.BindException: Cannot assign requested address`, which also takes down unrelated connections such as the Kafka source.
    
    The loop has several early returns, so the close has to be tied to the scope.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Two tests added to `TransformWithStateInPySparkStateServerSuite`, covering the path where the request loop is never entered and the EOF early-return path. Both fail without the fix. The suite previously had no test calling `run()`.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Yes. Generated-by: Claude Opus 5
    
    Closes #58205 from j1wonpark/tws-state-server-socket-leak.
    
    Authored-by: Jiwon Park <jp...@outlook.kr>
    Signed-off-by: Jungtaek Lim <ka...@gmail.com>
    (cherry picked from commit c5fbdf74b6886bd2c21a5cd4bc5c869e490abe63)
    Signed-off-by: Jungtaek Lim <ka...@gmail.com>
---
 .../TransformWithStateInPySparkStateServer.scala   | 11 ++++++-
 ...ansformWithStateInPySparkStateServerSuite.scala | 34 ++++++++++++++++++++--
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala
index 61df2ff1116a..c2cea67061bc 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala
@@ -22,7 +22,8 @@ import java.nio.channels.{
   Channels,
   ClosedByInterruptException,
   ClosedChannelException,
-  ServerSocketChannel
+  ServerSocketChannel,
+  SocketChannel
 }
 import java.time.Duration
 
@@ -45,6 +46,7 @@ import org.apache.spark.sql.execution.streaming.state.StateMessage.KeyAndValuePa
 import org.apache.spark.sql.execution.streaming.state.StateMessage.StateResponseWithListGet
 import org.apache.spark.sql.streaming.{ListState, MapState, TTLConfig, ValueState}
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.util.Utils
 
 /**
  * This class is used to handle the state requests from the Python side. It runs on a separate
@@ -157,6 +159,13 @@ class TransformWithStateInPySparkStateServer(
         return
     }
 
+    // The task completion listener closes only the listening server socket, and the
+    // request loop has several early returns, so the accepted connection is closed
+    // through tryWithResource.
+    Utils.tryWithResource(listeningSocket)(serveRequests)
+  }
+
+  private def serveRequests(listeningSocket: SocketChannel): Unit = {
     // SPARK-51667: We have a pattern of sending messages continuously from one side
     // (Python -> JVM, and vice versa) before getting response from other side. Since most
     // messages we are sending are small, this triggers the bad combination of Nagle's algorithm
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala
index b1d586d2207b..ea1c8a0a5d8e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala
@@ -17,12 +17,14 @@
 package org.apache.spark.sql.execution.python.streaming
 
 import java.io.{DataOutputStream, InterruptedIOException}
-import java.net.InetSocketAddress
+import java.net.{InetSocketAddress, Socket}
+import java.nio.ByteBuffer
 import java.nio.channels.{
   AsynchronousCloseException,
   ClosedByInterruptException,
   ClosedChannelException,
-  ServerSocketChannel
+  ServerSocketChannel,
+  SocketChannel
 }
 import java.util.concurrent.atomic.AtomicReference
 
@@ -120,6 +122,34 @@ class TransformWithStateInPySparkStateServerSuite extends SparkFunSuite with Bef
       .thenReturn(Seq(getIntegerRow(1)))
   }
 
+  test("run closes the accepted socket once the request loop ends") {
+    val acceptedSocket = mock(classOf[SocketChannel])
+    when(serverSocket.accept()).thenReturn(acceptedSocket)
+    when(acceptedSocket.socket()).thenReturn(mock(classOf[Socket]))
+    // Ends the request loop right away: this test is about the socket, not the requests.
+    when(acceptedSocket.isConnected).thenReturn(false)
+
+    stateServer.run()
+
+    verify(acceptedSocket).close()
+  }
+
+  test("run closes the accepted socket when the client disconnects") {
+    val acceptedSocket = mock(classOf[SocketChannel])
+    when(serverSocket.accept()).thenReturn(acceptedSocket)
+    when(acceptedSocket.socket()).thenReturn(mock(classOf[Socket]))
+    when(acceptedSocket.isConnected).thenReturn(true)
+    // Channels.newInputStream synchronizes on this before reading.
+    when(acceptedSocket.blockingLock()).thenReturn(new Object)
+    when(acceptedSocket.isBlocking).thenReturn(true)
+    // No bytes ever arrive, so the read hits EOF and the loop returns early.
+    when(acceptedSocket.read(any(classOf[ByteBuffer]))).thenReturn(-1)
+
+    stateServer.run()
+
+    verify(acceptedSocket).close()
+  }
+
   test("set handle state") {
     val message = StatefulProcessorCall.newBuilder().setSetHandleState(
       SetHandleState.newBuilder().setState(HandleState.CREATED).build()).build()


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org