Fix/spark 58751 listener shutdown

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 = {

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

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