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..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 @@ -16,15 +16,25 @@ */ package org.apache.spark.sql.execution.python.streaming -import java.io.DataOutputStream -import java.nio.channels.ServerSocketChannel +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} @@ -38,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]) @@ -637,6 +648,100 @@ 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]) + } + } + + 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, + 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) }