zozo123 commented on code in PR #59232:
URL: https://github.com/apache/spark/pull/59232#discussion_r4233891915


##########
udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoGrpcWorkerMain.scala:
##########
@@ -125,7 +128,11 @@ object EchoGrpcWorkerMain {
 
       server.start()
       stderrln(s"listening on $socketPath (transport=${transport.name})")
-      // Park until the shutdown hook fires.
+      // SIGKILL skips JVM shutdown hooks. Register before parking so a parent
+      // that has already exited still runs cleanup and unblocks the wait 
below.
+      if (launcher.isPresent) {
+        launcher.get().onExit().thenRun(() => cleanup())

Review Comment:
   Done in `fee9de6e`. The class doc now says the worker blocks until the JVM 
receives SIGTERM or its parent process exits.



##########
udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherIntegrationSuite.scala:
##########
@@ -160,6 +162,91 @@ class DirectGrpcDispatcherIntegrationSuite
     assert(!socketDir.exists(), "dispatcher close should remove its socket 
directory")
   }
 
+  test("a launcher command that forwards the connection serves sessions") {
+    val launcher = ProcessCallable.newBuilder()
+      .addCommand(javaExecutable)
+      .addCommand("-cp")
+      .addCommand(javaClasspath)
+      .addCommand(classOf[ForwardingWorkerMain.type].getName.stripSuffix("$"))
+      .addAllCommand(echoRunner.getCommandList)
+      .build()
+    dispatcher = new DirectGrpcDispatcher(workerSpec(launcher))
+    val sessions = Seq.fill(2)(dispatcher.createSession(None))
+    val processes = sessions.map(workerProcess)
+    val workers = processes.flatMap(_.process.descendants().iterator().asScala)
+    val payloads = Seq(
+      ByteString.copyFrom(Array.fill[Byte](5 * 1024 * 1024)(3)),
+      ByteString.copyFromUtf8("small"))
+
+    try {
+      assert(workers.size == 2, s"each launcher should start one worker, got 
$workers")
+      sessions.zip(payloads).foreach { case (session, payload) =>
+        session.init(basicInit)
+        val input = DataRequest.newBuilder().setData(payload).build()
+        assert(session.process(Iterator.single(input), 
emptyFinish).map(_.getData).toList ==
+          List(payload))
+      }
+    } finally {
+      sessions.foreach(_.close(emptyCancel))
+    }
+
+    assert(processes.forall(!_.process.isAlive), "session close should 
terminate the launcher")
+    workers.foreach { worker =>
+      assertWorkerStopped(worker, "session close should stop the launched 
worker")
+    }
+  }
+
+  test("SIGKILL of the forwarding launcher stops the inner worker") {
+    val launcher = ProcessCallable.newBuilder()
+      .addCommand(javaExecutable)
+      .addCommand("-cp")
+      .addCommand(javaClasspath)
+      .addCommand(classOf[ForwardingWorkerMain.type].getName.stripSuffix("$"))
+      .addAllCommand(echoRunner.getCommandList)
+      .build()
+    dispatcher = new DirectGrpcDispatcher(workerSpec(launcher))
+    val session = dispatcher.createSession(None)
+    val process = workerProcess(session)
+    val workers = process.process.descendants().iterator().asScala.toList
+    assert(workers.size == 1, s"the launcher should start one worker, got 
$workers")
+    val worker = workers.head
+    try {
+      session.init(basicInit)
+      process.process.destroyForcibly()
+      assertWorkerStopped(worker, "SIGKILL of the launcher should stop worker")
+    } finally {
+      session.close(emptyCancel)
+    }
+  }
+
+  // An orphaned worker is reparented, and an unreaped zombie stays isAlive
+  // until PID 1 reaps it. State Z in /proc means the worker has already 
exited.
+  private def assertWorkerStopped(worker: ProcessHandle, what: String): Unit = 
{

Review Comment:
   Done in `fee9de6e`. `assertWorkerStopped`, `hasExited`, and `linuxState` sit 
with the other private helpers, above the first test.



##########
udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherIntegrationSuite.scala:
##########
@@ -160,6 +162,91 @@ class DirectGrpcDispatcherIntegrationSuite
     assert(!socketDir.exists(), "dispatcher close should remove its socket 
directory")
   }
 
+  test("a launcher command that forwards the connection serves sessions") {
+    val launcher = ProcessCallable.newBuilder()
+      .addCommand(javaExecutable)
+      .addCommand("-cp")
+      .addCommand(javaClasspath)
+      .addCommand(classOf[ForwardingWorkerMain.type].getName.stripSuffix("$"))
+      .addAllCommand(echoRunner.getCommandList)
+      .build()
+    dispatcher = new DirectGrpcDispatcher(workerSpec(launcher))
+    val sessions = Seq.fill(2)(dispatcher.createSession(None))
+    val processes = sessions.map(workerProcess)
+    val workers = processes.flatMap(_.process.descendants().iterator().asScala)
+    val payloads = Seq(
+      ByteString.copyFrom(Array.fill[Byte](5 * 1024 * 1024)(3)),
+      ByteString.copyFromUtf8("small"))
+
+    try {
+      assert(workers.size == 2, s"each launcher should start one worker, got 
$workers")
+      sessions.zip(payloads).foreach { case (session, payload) =>
+        session.init(basicInit)
+        val input = DataRequest.newBuilder().setData(payload).build()
+        assert(session.process(Iterator.single(input), 
emptyFinish).map(_.getData).toList ==
+          List(payload))
+      }
+    } finally {
+      sessions.foreach(_.close(emptyCancel))
+    }
+
+    assert(processes.forall(!_.process.isAlive), "session close should 
terminate the launcher")
+    workers.foreach { worker =>
+      assertWorkerStopped(worker, "session close should stop the launched 
worker")
+    }
+  }
+
+  test("SIGKILL of the forwarding launcher stops the inner worker") {
+    val launcher = ProcessCallable.newBuilder()
+      .addCommand(javaExecutable)
+      .addCommand("-cp")
+      .addCommand(javaClasspath)
+      .addCommand(classOf[ForwardingWorkerMain.type].getName.stripSuffix("$"))
+      .addAllCommand(echoRunner.getCommandList)
+      .build()
+    dispatcher = new DirectGrpcDispatcher(workerSpec(launcher))
+    val session = dispatcher.createSession(None)
+    val process = workerProcess(session)
+    val workers = process.process.descendants().iterator().asScala.toList
+    assert(workers.size == 1, s"the launcher should start one worker, got 
$workers")
+    val worker = workers.head
+    try {
+      session.init(basicInit)
+      process.process.destroyForcibly()
+      assertWorkerStopped(worker, "SIGKILL of the launcher should stop worker")
+    } finally {
+      session.close(emptyCancel)
+    }
+  }
+
+  // An orphaned worker is reparented, and an unreaped zombie stays isAlive
+  // until PID 1 reaps it. State Z in /proc means the worker has already 
exited.
+  private def assertWorkerStopped(worker: ProcessHandle, what: String): Unit = 
{
+    val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10)
+    while (System.nanoTime() < deadline && !hasExited(worker)) {
+      Thread.sleep(50)
+    }
+    if (!hasExited(worker)) {
+      worker.destroyForcibly()
+      fail(s"$what ${worker.pid}")
+    }
+  }
+
+  private def hasExited(worker: ProcessHandle): Boolean = {
+    !worker.isAlive || linuxState(worker.pid).contains('Z')
+  }
+
+  private def linuxState(pid: Long): Option[Char] = {
+    val stat = Paths.get(s"/proc/$pid/stat")
+    if (!Files.isRegularFile(stat)) {
+      None
+    } else {
+      val text = Files.readString(stat)

Review Comment:
   Done in `fee9de6e`. A failed `/proc` read returns no state, so a reap 
between the checks does not fail the test. Both launcher tests passed locally.



##########
udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherIntegrationSuite.scala:
##########
@@ -160,6 +162,65 @@ class DirectGrpcDispatcherIntegrationSuite
     assert(!socketDir.exists(), "dispatcher close should remove its socket 
directory")
   }
 
+  test("a launcher command that forwards the connection serves sessions") {
+    val launcher = ProcessCallable.newBuilder()
+      .addCommand(javaExecutable)
+      .addCommand("-cp")
+      .addCommand(javaClasspath)
+      .addCommand(classOf[ForwardingWorkerMain.type].getName.stripSuffix("$"))
+      .addAllCommand(echoRunner.getCommandList)
+      .build()
+    dispatcher = new DirectGrpcDispatcher(workerSpec(launcher))
+    val sessions = Seq.fill(2)(dispatcher.createSession(None))
+    val processes = sessions.map(workerProcess)
+    val workers = processes.flatMap(_.process.descendants().iterator().asScala)
+    val payloads = Seq(
+      ByteString.copyFrom(Array.fill[Byte](5 * 1024 * 1024)(3)),
+      ByteString.copyFromUtf8("small"))
+
+    try {
+      assert(workers.size == 2, s"each launcher should start one worker, got 
$workers")
+      sessions.zip(payloads).foreach { case (session, payload) =>
+        session.init(basicInit)
+        val input = DataRequest.newBuilder().setData(payload).build()
+        assert(session.process(Iterator.single(input), 
emptyFinish).map(_.getData).toList ==
+          List(payload))
+      }
+    } finally {
+      sessions.foreach(_.close(emptyCancel))
+    }
+
+    assert(processes.forall(!_.process.isAlive), "session close should 
terminate the launcher")
+    workers.foreach { worker =>
+      worker.onExit().get(10, TimeUnit.SECONDS)
+      assert(!worker.isAlive, s"session close should stop the launched worker 
${worker.pid}")
+    }
+  }
+
+  test("SIGKILL of the forwarding launcher stops the inner worker") {
+    val launcher = ProcessCallable.newBuilder()
+      .addCommand(javaExecutable)
+      .addCommand("-cp")
+      .addCommand(javaClasspath)
+      .addCommand(classOf[ForwardingWorkerMain.type].getName.stripSuffix("$"))
+      .addAllCommand(echoRunner.getCommandList)
+      .build()
+    dispatcher = new DirectGrpcDispatcher(workerSpec(launcher))
+    val session = dispatcher.createSession(None)
+    val process = workerProcess(session)
+    val workers = process.process.descendants().iterator().asScala.toList
+    assert(workers.size == 1, s"the launcher should start one worker, got 
$workers")
+    val worker = workers.head
+    try {
+      session.init(basicInit)
+      process.process.destroyForcibly()
+      worker.onExit().get(10, TimeUnit.SECONDS)

Review Comment:
   Done in `fee9de6e`. `linuxState` treats an `IOException` from that read as 
no state.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to