This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new c5fbdf74b688 [SPARK-58932][SS] Close the accepted socket in
TransformWithStateInPySparkStateServer
c5fbdf74b688 is described below
commit c5fbdf74b6886bd2c21a5cd4bc5c869e490abe63
Author: Jiwon Park <[email protected]>
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 <[email protected]>
Signed-off-by: Jungtaek Lim <[email protected]>
---
.../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: [email protected]
For additional commands, e-mail: [email protected]
