From e97c23d145541be36e2421582c2b0c397061b171 Mon Sep 17 00:00:00 2001 From: Jiwon Park Date: Wed, 26 Aug 2026 16:49:54 +0900 Subject: [PATCH] [SPARK-58932][SS] Close the accepted socket in TransformWithStateInPySparkStateServer Run the request loop of `TransformWithStateInPySparkStateServer.run()` inside `Utils.tryWithResource` so the socket returned by `stateServerSocket.accept()` is always closed. `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. No. 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()`. Yes. Generated-by: Claude Opus 5 Closes #58205 from j1wonpark/tws-state-server-socket-leak. Authored-by: Jiwon Park Signed-off-by: Jungtaek Lim (cherry picked from commit c5fbdf74b6886bd2c21a5cd4bc5c869e490abe63) --- ...ansformWithStateInPySparkStateServer.scala | 10 +++++- ...rmWithStateInPySparkStateServerSuite.scala | 32 ++++++++++++++++++- 2 files changed, 40 insertions(+), 2 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 4fee6a6e71d30..f457ee981fbfc 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 @@ -18,7 +18,7 @@ package org.apache.spark.sql.execution.python.streaming import java.io.{BufferedInputStream, BufferedOutputStream, DataInputStream, DataOutputStream, EOFException, InterruptedIOException} -import java.nio.channels.{Channels, ClosedByInterruptException, ServerSocketChannel} +import java.nio.channels.{Channels, ClosedByInterruptException, ServerSocketChannel, SocketChannel} import java.time.Duration import scala.collection.mutable @@ -40,6 +40,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 @@ -140,6 +141,13 @@ class TransformWithStateInPySparkStateServer( def run(): Unit = { val listeningSocket = stateServerSocket.accept() + // 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 e253a6aa45c35..d6a6a56092889 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,7 +17,9 @@ package org.apache.spark.sql.execution.python.streaming import java.io.DataOutputStream -import java.nio.channels.ServerSocketChannel +import java.net.Socket +import java.nio.ByteBuffer +import java.nio.channels.{ServerSocketChannel, SocketChannel} import scala.collection.mutable @@ -110,6 +112,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()