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()