From a0fd8e54fc702edf2ef5e4b75c02c28f23dad4ca Mon Sep 17 00:00:00 2001 From: Jonathan Gao Date: Thu, 20 Aug 2026 21:34:48 +0000 Subject: [PATCH 1/5] [SPARK-58751][SS][PYTHON] Handle state server shutdown before Python connects --- ...nsformWithStateInPySparkPythonRunner.scala | 2 +- ...ansformWithStateInPySparkStateServer.scala | 21 ++++++- ...rmWithStateInPySparkStateServerSuite.scala | 61 ++++++++++++++++++- 3 files changed, 79 insertions(+), 5 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala index 488ac1a0bd232..0b7e6b938689f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala @@ -333,10 +333,10 @@ class TransformWithStateInPySparkPythonPreInitRunner( override def stop(): Unit = { super.stop() - closeServerSocketChannelSilently(stateServerSocket) if (daemonThread != null) { daemonThread.interrupt() } + closeServerSocketChannelSilently(stateServerSocket) } private def startStateServer(): Unit = { 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..61df2ff1116ad 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,12 @@ 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, + ClosedChannelException, + ServerSocketChannel +} import java.time.Duration import scala.collection.mutable @@ -138,7 +143,19 @@ class TransformWithStateInPySparkStateServer( } else new mutable.HashMap[String, Iterator[Long]]() def run(): Unit = { - val listeningSocket = stateServerSocket.accept() + val listeningSocket = try { + stateServerSocket.accept() + } catch { + case _: InterruptedException | _: InterruptedIOException | _: ClosedByInterruptException => + logInfo(log"State server listener interrupted before the Python worker connected") + Thread.currentThread().interrupt() + statefulProcessorHandle.setHandleState(StatefulProcessorHandleState.CLOSED) + return + case _: ClosedChannelException => + logInfo(log"State server socket closed before the Python worker connected") + statefulProcessorHandle.setHandleState(StatefulProcessorHandleState.CLOSED) + return + } // 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 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..0375e9d41a789 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 @@ -16,13 +16,19 @@ */ package org.apache.spark.sql.execution.python.streaming -import java.io.DataOutputStream -import java.nio.channels.ServerSocketChannel +import java.io.{DataOutputStream, InterruptedIOException} +import java.nio.channels.{ + AsynchronousCloseException, + ClosedByInterruptException, + ClosedChannelException, + ServerSocketChannel +} import scala.collection.mutable import com.google.protobuf.ByteString import org.mockito.ArgumentMatchers.{any, argThat} +import org.mockito.invocation.InvocationOnMock import org.mockito.Mockito.{mock, times, verify, when} import org.scalatest.BeforeAndAfterEach @@ -637,6 +643,57 @@ class TransformWithStateInPySparkStateServerSuite extends SparkFunSuite with Bef verify(outputStream).writeInt(argThat((x: Int) => x > 0)) } + Seq( + ("InterruptedException", () => new InterruptedException()), + ("InterruptedIOException", () => new InterruptedIOException()), + ("ClosedByInterruptException", () => new ClosedByInterruptException()) + ).foreach { case (name, newException) => + test(s"run handles $name while waiting for the Python worker") { + Thread.interrupted() + val socket = mock(classOf[ServerSocketChannel]) + when(socket.accept()) + .thenAnswer((_: InvocationOnMock) => throw newException()) + + try { + newStateServer(socket).run() + assert(Thread.currentThread().isInterrupted) + } finally { + Thread.interrupted() + } + + verify(statefulProcessorHandle).setHandleState(StatefulProcessorHandleState.CLOSED) + verify(outputStream, times(0)).writeInt(any[Int]) + } + } + + Seq( + ("AsynchronousCloseException", () => new AsynchronousCloseException()), + ("ClosedChannelException", () => new ClosedChannelException()) + ).foreach { case (name, newException) => + test(s"run handles $name while waiting for the Python worker") { + Thread.interrupted() + val socket = mock(classOf[ServerSocketChannel]) + when(socket.accept()) + .thenAnswer((_: InvocationOnMock) => throw newException()) + + newStateServer(socket).run() + + assert(!Thread.currentThread().isInterrupted) + verify(statefulProcessorHandle).setHandleState(StatefulProcessorHandleState.CLOSED) + verify(outputStream, times(0)).writeInt(any[Int]) + } + } + + private def newStateServer( + socket: ServerSocketChannel): TransformWithStateInPySparkStateServer = { + new TransformWithStateInPySparkStateServer(socket, + statefulProcessorHandle, groupingKeySchema, 2, + batchTimestampMs, eventTimeWatermarkForEviction, + outputStream, valueStateMap, transformWithStateInPySparkDeserializer, + listStateMap, mutable.HashMap[String, Iterator[Row]](), mapStateMap, + mutable.HashMap[String, Iterator[(Row, Row)]](), expiryTimerIter, listTimerMap) + } + private def getIntegerRow(value: Int): Row = { new GenericRowWithSchema(Array(value), stateSchema) } From 57abae5072d2e54f1be021b88389b221022a5496 Mon Sep 17 00:00:00 2001 From: Jonathan Gao Date: Fri, 21 Aug 2026 03:10:17 +0000 Subject: [PATCH 2/5] [SPARK-58751][SS][PYTHON] Test state server shutdown with real channel --- ...rmWithStateInPySparkStateServerSuite.scala | 50 ++++++++++++++++++- 1 file changed, 49 insertions(+), 1 deletion(-) 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 0375e9d41a789..6da372d1fd77e 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,20 +17,24 @@ package org.apache.spark.sql.execution.python.streaming import java.io.{DataOutputStream, InterruptedIOException} +import java.net.InetSocketAddress import java.nio.channels.{ AsynchronousCloseException, ClosedByInterruptException, ClosedChannelException, ServerSocketChannel } +import java.util.concurrent.atomic.AtomicReference import scala.collection.mutable +import scala.concurrent.duration._ import com.google.protobuf.ByteString import org.mockito.ArgumentMatchers.{any, argThat} import org.mockito.invocation.InvocationOnMock import org.mockito.Mockito.{mock, times, verify, when} import org.scalatest.BeforeAndAfterEach +import org.scalatest.concurrent.Eventually import org.apache.spark.SparkFunSuite import org.apache.spark.sql.{Encoder, Row} @@ -44,7 +48,8 @@ import org.apache.spark.sql.types.{IntegerType, StructField, StructType} import org.apache.spark.tags.SlowSQLTest @SlowSQLTest -class TransformWithStateInPySparkStateServerSuite extends SparkFunSuite with BeforeAndAfterEach { +class TransformWithStateInPySparkStateServerSuite + extends SparkFunSuite with BeforeAndAfterEach with Eventually { val stateName = "test" val iteratorId = "testId" val serverSocket: ServerSocketChannel = mock(classOf[ServerSocketChannel]) @@ -684,6 +689,49 @@ class TransformWithStateInPySparkStateServerSuite extends SparkFunSuite with Bef } } + Seq( + ("before accept", true), + ("while blocked in accept", false) + ).foreach { case (name, interruptBeforeRun) => + test(s"run handles real channel shutdown $name") { + val socket = ServerSocketChannel.open() + socket.bind(new InetSocketAddress("127.0.0.1", 0)) + val failure = new AtomicReference[Throwable]() + val listener = new Thread(() => { + if (interruptBeforeRun) { + Thread.currentThread().interrupt() + } + try { + newStateServer(socket).run() + } catch { + case t: Throwable => failure.set(t) + } + }) + + try { + listener.start() + if (!interruptBeforeRun) { + eventually(timeout(10.seconds)) { + assert(listener.getStackTrace.exists(_.getMethodName == "accept")) + } + listener.interrupt() + } + socket.close() + listener.join(10000) + + assert(!listener.isAlive) + assert(failure.get() == null) + assert(!socket.isOpen) + verify(statefulProcessorHandle).setHandleState(StatefulProcessorHandleState.CLOSED) + verify(outputStream, times(0)).writeInt(any[Int]) + } finally { + listener.interrupt() + socket.close() + listener.join(10000) + } + } + } + private def newStateServer( socket: ServerSocketChannel): TransformWithStateInPySparkStateServer = { new TransformWithStateInPySparkStateServer(socket, From 61f32b849a713b88f4ead076cb33e24f15caeb9a Mon Sep 17 00:00:00 2001 From: Jonathan Gao Date: Tue, 25 Aug 2026 00:34:34 +0000 Subject: [PATCH 3/5] [SPARK-58977][CI] Trigger CI From fdade8ac6230af4e95269b420e50f9fc27bf4a3e Mon Sep 17 00:00:00 2001 From: Jonathan Gao Date: Tue, 25 Aug 2026 01:59:03 +0000 Subject: [PATCH 4/5] [SPARK-58977][SS][PYTHON] Fix test import order --- .../streaming/TransformWithStateInPySparkStateServerSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 6da372d1fd77e..2ecfd404136d1 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 @@ -31,8 +31,8 @@ import scala.concurrent.duration._ import com.google.protobuf.ByteString import org.mockito.ArgumentMatchers.{any, argThat} -import org.mockito.invocation.InvocationOnMock import org.mockito.Mockito.{mock, times, verify, when} +import org.mockito.invocation.InvocationOnMock import org.scalatest.BeforeAndAfterEach import org.scalatest.concurrent.Eventually From 074b991251937e3aaa6d8043e712896dd452be94 Mon Sep 17 00:00:00 2001 From: Jonathan Gao Date: Tue, 25 Aug 2026 02:51:04 +0000 Subject: [PATCH 5/5] [SPARK-58977][TESTS] Use Eventually companion imports --- .../TransformWithStateInPySparkStateServerSuite.scala | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 2ecfd404136d1..b1d586d2207b0 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 @@ -34,7 +34,7 @@ import org.mockito.ArgumentMatchers.{any, argThat} import org.mockito.Mockito.{mock, times, verify, when} import org.mockito.invocation.InvocationOnMock import org.scalatest.BeforeAndAfterEach -import org.scalatest.concurrent.Eventually +import org.scalatest.concurrent.Eventually.{eventually, timeout} import org.apache.spark.SparkFunSuite import org.apache.spark.sql.{Encoder, Row} @@ -48,8 +48,7 @@ import org.apache.spark.sql.types.{IntegerType, StructField, StructType} import org.apache.spark.tags.SlowSQLTest @SlowSQLTest -class TransformWithStateInPySparkStateServerSuite - extends SparkFunSuite with BeforeAndAfterEach with Eventually { +class TransformWithStateInPySparkStateServerSuite extends SparkFunSuite with BeforeAndAfterEach { val stateName = "test" val iteratorId = "testId" val serverSocket: ServerSocketChannel = mock(classOf[ServerSocketChannel])