Make Java SDK context signal-aware (#74690)

* Make Java SDK context signal-aware

Signed-off-by: PoAn Yang <payang@apache.org>

* Share the stop request with every context made during a task

Signed-off-by: PoAn Yang <payang@apache.org>

* Update java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt

---------

Signed-off-by: PoAn Yang <payang@apache.org>
Co-authored-by: Jason(Zhe-You) Liu <68415893+jason810496@users.noreply.github.com>

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt

@@ -19,8 +19,10 @@
 
 package org.apache.airflow.sdk
 
+import org.apache.airflow.sdk.execution.Cancellation
 import org.apache.airflow.sdk.execution.Logger
 import org.apache.airflow.sdk.execution.comm.StartupDetails
+import java.time.Duration
 import java.time.OffsetDateTime
 import org.apache.airflow.sdk.execution.comm.DagRun.DagRunType as CommDagRunType
 
@@ -127,6 +129,54 @@ data class Context(
   /** Registration of the executing task; resolves wired data inputs. */
   internal var taskDef: TaskDef? = null
 
+  // The value is read once, when the context is made. The server sets Cancellation.running back to
+  // null once it is done with the task, so a context that read the field on every call would stop
+  // seeing the stop request in code that runs after that, such as a JVM shutdown hook.
+  private val cancellation = Cancellation.current()
+
+  /**
+   * Whether Airflow has asked this task to stop.
+   *
+   * Airflow asks a running task to stop when, for example, a user marks its task
+   * instance as failed. The supervisor, which is the Airflow process that started
+   * the JVM, then sends the JVM SIGTERM. If the JVM is still running after a
+   * grace period, the supervisor sends SIGKILL.
+   *
+   * Once a task has called this method or [awaitCancellation], the SDK catches
+   * the first SIGTERM or SIGINT. The SDK records the stop request instead of
+   * letting the JVM exit, so the task can stop its work before SIGKILL arrives.
+   * From then on the task must keep checking until it returns.
+   *
+   * If a task has called neither method, or if a second signal arrives, the JVM
+   * handles the signal as it would without the SDK. By default the JVM runs its
+   * shutdown hooks and exits.
+   *
+   * Every context made while the task runs sees the same stop request,
+   * including a context made with `copy()`, and keeps seeing it after the task
+   * returns. A context made when the server is not running a task, for example
+   * in a unit test that calls a task directly, never sees a stop request.
+   */
+  fun isCancelled(): Boolean = cancellation.check()
+
+  /**
+   * Waits until Airflow asks this task to stop, or until [timeout] passes.
+   *
+   * Use this method in place of [Thread.sleep] in a polling loop, so the task
+   * stops waiting as soon as Airflow asks it to stop. See [isCancelled] for when
+   * Airflow asks a task to stop.
+   *
+   * A context made when the server is not running a task never sees a stop
+   * request. On such a context, this method waits for the whole [timeout] and
+   * returns false.
+   *
+   * @param timeout How long to wait at most.
+   * @return `true` if Airflow asked the task to stop, `false` if [timeout]
+   *    passed first.
+   * @throws InterruptedException if the waiting thread is interrupted.
+   */
+  @Throws(InterruptedException::class)
+  fun awaitCancellation(timeout: Duration): Boolean = cancellation.await(timeout)
+
   internal companion object {
     fun from(
       request: StartupDetails,

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt

@@ -33,6 +33,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
+import org.apache.airflow.sdk.execution.Cancellation
 import org.apache.airflow.sdk.execution.CoordinatorComm
 import org.apache.airflow.sdk.execution.LogSender
 import org.apache.airflow.sdk.execution.Logger
@@ -41,6 +42,7 @@ import org.apache.airflow.sdk.execution.comm.ErrorResponse
 import org.apache.airflow.sdk.execution.comm.StartupDetails
 import org.apache.airflow.sdk.execution.parseDags
 import org.apache.airflow.sdk.execution.runTask
+import org.apache.airflow.sdk.execution.trapStopSignals
 import java.io.File
 import kotlin.text.substringAfterLast
 import kotlin.text.substringBeforeLast
@@ -235,8 +237,31 @@ class CoordinatorServer(
     startup: StartupDetails,
     coordinator: CoordinatorComm,
   ) {
-    val result = runTask(bundle, startup, coordinator)
-    coordinator.communicate<Unit>(result)
+    val cancellation = Cancellation()
+    Cancellation.running = cancellation
+    try {
+      // The trap stays open until the result is reported. For a task that checks, a signal that
+      // arrives during the report therefore cannot stop the JVM before the result is sent.
+      trapStopSignals { signal ->
+        // A task that never checks cannot stop on its own. A second signal means the sender does not
+        // want to wait. In either case the JVM handles the signal as it would without the SDK.
+        val take = cancellation.isWatched && cancellation.claim()
+        if (take) {
+          // Log before the task can see the request. Otherwise the task may return and the JVM may
+          // exit before this line is written.
+          logger.warning("Asking the task to stop", mapOf("signal" to signal))
+          cancellation.request()
+        } else {
+          logger.info("Letting the JVM handle the signal", mapOf("signal" to signal))
+        }
+        take
+      }.use {
+        val result = runTask(bundle, startup, coordinator)
+        coordinator.communicate<Unit>(result)
+      }
+    } finally {
+      Cancellation.running = null
+    }
   }
 
   private suspend fun parseDagsAndReport(

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Cancellation.kt

@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk.execution
+
+import sun.misc.Signal
+import sun.misc.SignalHandler
+import java.time.Duration
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicReference
+
+/**
+ * Records whether Airflow has asked the running task to stop, and whether the task checks for that
+ * request.
+ */
+internal class Cancellation {
+  private val claimed = AtomicBoolean()
+  private val requested = CountDownLatch(1)
+
+  /**
+   * True once the task has checked for a stop request or waited for one, which shows that the task
+   * can stop on its own.
+   */
+  @Volatile
+  var isWatched = false
+    private set
+
+  val isRequested: Boolean
+    get() = requested.count == 0L
+
+  /** Returns true to the first caller only, so the SDK can take at most one signal for the task. */
+  fun claim(): Boolean = claimed.compareAndSet(false, true)
+
+  fun request() = requested.countDown()
+
+  /**
+   * Returns [isRequested] for the task. Because the task has now checked, this call also sets
+   * [isWatched].
+   */
+  fun check(): Boolean {
+    watch()
+    return isRequested
+  }
+
+  fun await(timeout: Duration): Boolean {
+    watch()
+    return requested.await(TimeUnit.NANOSECONDS.convert(timeout), TimeUnit.NANOSECONDS)
+  }
+
+  private fun watch() {
+    if (!isWatched) isWatched = true
+  }
+
+  companion object {
+    /**
+     * The stop request of the task that the server is running in this JVM. The server sets this field
+     * before the task starts. When the server is done with the task, it sets the field back to null.
+     * Every [org.apache.airflow.sdk.Context] made while the task runs keeps this instance, so a copy
+     * of a context made during the task also sees the same request. A single field is enough because
+     * the server runs one task at a time in a JVM, and a stop signal reaches the whole JVM, not one
+     * task.
+     */
+    @Volatile
+    var running: Cancellation? = null
+
+    /**
+     * Returns [running]. When the server is not running a task, for example in a unit test that calls
+     * a task directly, it returns a new [Cancellation] that nobody asks to stop.
+     */
+    fun current(): Cancellation = running ?: Cancellation()
+  }
+}
+
+private val logger = Logger(Cancellation::class)
+
+/**
+ * Opens a trap that passes SIGINT and SIGTERM to [onSignal] instead of the JVM. The trap stays open
+ * until the caller closes the returned handle.
+ *
+ * When [onSignal] returns false or throws, the JVM gets its own handler back and receives the
+ * signal again, so the JVM handles the signal as it would have without the trap. Only one trap can
+ * be open at a time, because closing a trap puts back the handlers that were installed when that
+ * trap opened.
+ */
+internal fun trapStopSignals(onSignal: (signal: String) -> Boolean): AutoCloseable {
+  val restores = listOf("INT", "TERM").mapNotNull { trapSignal(it, onSignal) }
+  return AutoCloseable { restores.forEach { restore -> restore() } }
+}
+
+private fun trapSignal(
+  name: String,
+  onSignal: (signal: String) -> Boolean,
+): (() -> Unit)? =
+  try {
+    val signal = Signal(name)
+    val previous = AtomicReference<SignalHandler>()
+    previous.set(
+      Signal.handle(signal) {
+        if (!runCatching { onSignal("SIG$name") }.getOrDefault(false)) {
+          // A signal can arrive before Signal.handle has returned the handler it replaced.
+          while (previous.get() == null) Thread.onSpinWait()
+          Signal.handle(signal, previous.get())
+          Signal.raise(signal)
+        }
+      },
+    )
+    ({ Signal.handle(signal, previous.get()) })
+  } catch (e: IllegalArgumentException) {
+    // The JVM refuses the signal, for example when it was started with -Xrs.
+    logger.warning("Cannot catch signal, so the task cannot see it", mapOf("signal" to "SIG$name", "error" to e))
+    null
+  } catch (e: LinkageError) {
+    // sun.misc.Signal lives in the jdk.unsupported module, which a jlink runtime image can leave out.
+    logger.warning("Cannot catch signal, so the task cannot see it", mapOf("signal" to "SIG$name", "error" to e))
+    null
+  }

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt

@@ -220,10 +220,13 @@ internal class Logger(
   }
 }
 
+// The log channel accepts one writer at a time, but several threads log: the task's own thread, the
+// thread that handles a stop signal, and the thread that reads replies from the supervisor.
 internal object LogSender {
   private var writer: ByteWriteChannel? = null
   internal val messages: ConcurrentLinkedDeque<LogMessage> = ConcurrentLinkedDeque()
 
+  @Synchronized
   fun configure(channel: ByteWriteChannel) {
     writer = channel
     if (!channel.isClosedForWrite) {
@@ -233,6 +236,7 @@ internal object LogSender {
     }
   }
 
+  @Synchronized
   fun send(message: LogMessage) {
     val channel = writer
     if (channel == null || channel.isClosedForWrite) {

java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ContextTest.kt

@@ -25,7 +25,10 @@ import org.apache.airflow.sdk.execution.comm.StartupDetails
 import org.apache.airflow.sdk.execution.comm.TIRunContext
 import org.junit.jupiter.api.Assertions
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.Timeout
+import java.time.Duration
 import java.time.OffsetDateTime
+import java.util.concurrent.TimeUnit
 
 class ContextTest {
   @Test
@@ -125,4 +128,15 @@ class ContextTest {
     Assertions.assertNull(dr.dataIntervalStart)
     Assertions.assertTrue(dr.conf.isEmpty())
   }
+
+  @Test
+  @Timeout(value = 30, unit = TimeUnit.SECONDS)
+  fun contextBuiltByHandIsNeverCancelled() {
+    val context = taskContext()
+    val start = System.nanoTime()
+
+    Assertions.assertFalse(context.isCancelled())
+    Assertions.assertFalse(context.awaitCancellation(Duration.ofMillis(200)), "the wait should end at the timeout")
+    Assertions.assertTrue(System.nanoTime() - start >= Duration.ofMillis(200).toNanos(), "the wait should last the whole timeout")
+  }
 }

java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt

@@ -27,17 +27,25 @@ import kotlinx.coroutines.runBlocking
 import org.apache.airflow.sdk.execution.CoordinatorComm
 import org.apache.airflow.sdk.execution.Frame
 import org.apache.airflow.sdk.execution.IncomingFrame
+import org.apache.airflow.sdk.execution.LogSender
 import org.apache.airflow.sdk.execution.RawFrame
+import org.apache.airflow.sdk.execution.comm.SucceedTask
 import org.apache.airflow.sdk.execution.comm.TaskState
+import org.apache.airflow.sdk.execution.pollSignals
+import org.apache.airflow.sdk.execution.withSignalRecorder
 import org.junit.jupiter.api.Assertions
 import org.junit.jupiter.api.DisplayName
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.Timeout
 import org.msgpack.core.MessagePack
 import org.msgpack.core.buffer.ArrayBufferInput
+import sun.misc.Signal
 import java.io.ByteArrayOutputStream
+import java.time.Duration
 import java.util.concurrent.ArrayBlockingQueue
+import java.util.concurrent.LinkedBlockingQueue
 import java.util.concurrent.TimeUnit
+import kotlin.concurrent.thread
 
 private class ServerNoopTask : Task {
   override fun execute(
@@ -46,6 +54,71 @@ private class ServerNoopTask : Task {
   ) = Unit
 }
 
+/** What the tasks below report to the test that runs them. */
+private val taskEvents = LinkedBlockingQueue<String>()
+
+/** Releases [NeverCheckingTask]. */
+private val releaseTask = LinkedBlockingQueue<String>()
+
+/**
+ * Waits for a stop request and records what its context showed. Tests use this task to signal the
+ * JVM while a task is running.
+ */
+class StopAwareTask : Task {
+  override fun execute(
+    context: Context,
+    client: Client,
+  ) {
+    // Check before reporting "started". The test raises SIGTERM as soon as it reads "started", so by
+    // then the task has checked and the runtime keeps the signal.
+    taskEvents.put("isCancelled() returned ${context.isCancelled()}")
+    taskEvents.put("started")
+    taskEvents.put("awaitCancellation returned ${context.awaitCancellation(Duration.ofSeconds(20))}")
+    taskEvents.put("isCancelled() returned ${context.isCancelled()}")
+  }
+}
+
+/** The copy that [CopyCheckingTask] checked, kept so the test can check it again after the task returns. */
+@Volatile
+private var checkedCopy: Context? = null
+
+/**
+ * This task works like [StopAwareTask], but it checks for a stop request only through a copy of its
+ * context that another thread made.
+ */
+class CopyCheckingTask : Task {
+  override fun execute(
+    context: Context,
+    client: Client,
+  ) {
+    thread { checkedCopy = context.copy() }.join()
+    val copy = checkNotNull(checkedCopy)
+    taskEvents.put("isCancelled() returned ${copy.isCancelled()}")
+    taskEvents.put("started")
+    taskEvents.put("awaitCancellation returned ${copy.awaitCancellation(Duration.ofSeconds(20))}")
+  }
+}
+
+class CheckOnceTask : Task {
+  override fun execute(
+    context: Context,
+    client: Client,
+  ) {
+    taskEvents.put("isCancelled() returned ${context.isCancelled()}")
+  }
+}
+
+/** Runs until the test releases it, without ever checking for a stop request. */
+class NeverCheckingTask : Task {
+  override fun execute(
+    context: Context,
+    client: Client,
+  ) {
+    taskEvents.put("started")
+    releaseTask.poll(20, TimeUnit.SECONDS)
+  }
+}
+
 class ServerTest {
   private fun hexToBytes(hex: String): ByteArray =
     hex
@@ -103,6 +176,223 @@ class ServerTest {
     comm.close()
   }
 
+  /**
+   * Runs [taskClass] as task "a" of Dag "c" through [CoordinatorServer.dispatchTask] and returns the
+   * result that the task reports. A thread that plays the supervisor sends the startup details, calls
+   * [whileRunning], reads the result, and calls [beforeAck] before it acknowledges the result. With
+   * [ack] set to false, that thread closes the connection instead, so reporting the result fails.
+   */
+  private fun dispatchTask(
+    taskClass: Class<out Task>,
+    whileRunning: () -> Unit = {},
+    beforeAck: () -> Unit = {},
+    ack: Boolean = true,
+  ): IncomingFrame {
+    taskEvents.clear()
+    LogSender.messages.clear()
+    val toServer = ByteChannel(autoFlush = true)
+    val fromServer = ByteChannel(autoFlush = true)
+    val comm = CoordinatorComm(toServer, fromServer)
+    val server = CoordinatorServer(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0))
+    val reported = ArrayBlockingQueue<IncomingFrame>(1)
+    val supervisor =
+      thread(isDaemon = true) {
+        runBlocking {
+          // Deliver the StartupDetails frame (id 2, dag_id "c", task_id "a").
+          toServer.writeFrame(hexToBytes(STARTUP_HEX))
+          whileRunning()
+          val prefix = fromServer.readByteArray(4)
+          val payload = fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt())
+          val result = CoordinatorComm.decode(payload)
+          reported.put(result)
+          beforeAck()
+          if (ack) toServer.writeFrame(ackFrame(result.id)) else comm.close()
+        }
+      }
+
+    try {
+      runBlocking { server.dispatchTask(Bundle(listOf(DagDef("c").addTask(TaskDef("a", taskClass)))), comm) }
+    } finally {
+      supervisor.join()
+      comm.close()
+    }
+    return reported.take()
+  }
+
+  /** The events the runtime logged for the stop signals it got, in order. */
+  private fun signalLogs(): List<String> = LogSender.messages.map { it.event }.filter { it in SIGNAL_EVENTS }
+
+  /**
+   * Waits until the runtime has logged [count] stop signals. By then the runtime has decided what to
+   * do with each signal.
+   */
+  private fun awaitSignalLogs(count: Int) {
+    val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10)
+    while (signalLogs().size < count && System.nanoTime() < deadline) Thread.sleep(10)
+  }
+
+  /** Waits until the running task reports [event], and drops any events the task reported earlier. */
+  private fun awaitTaskEvent(event: String) {
+    val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20)
+    while (System.nanoTime() < deadline && taskEvents.poll(100, TimeUnit.MILLISECONDS) != event) Unit
+  }
+
+  /** Waits until [count] SIGTERM handler threads are blocked, for example on a lock the test holds. */
+  private fun awaitBlockedSignalHandlers(count: Int) {
+    val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10)
+    while (System.nanoTime() < deadline) {
+      val blocked = Thread.getAllStackTraces().keys.count { it.name == "SIGTERM handler" && it.state == Thread.State.BLOCKED }
+      if (blocked >= count) return
+      Thread.sleep(10)
+    }
+  }
+
+  @Test
+  @DisplayName("Should let a task that waits for a stop request see SIGTERM, and leave a second SIGTERM to the JVM")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun runningTaskSeesStopSignal() =
+    withSignalRecorder("TERM") { recorded ->
+      val result =
+        dispatchTask(
+          StopAwareTask::class.java,
+          whileRunning = {
+            awaitTaskEvent("started")
+            Signal.raise(Signal("TERM"))
+          },
+          beforeAck = {
+            Signal.raise(Signal("TERM"))
+            awaitSignalLogs(2)
+          },
+        )
+
+      Assertions.assertEquals(
+        listOf("awaitCancellation returned true", "isCancelled() returned true"),
+        taskEvents.toList(),
+        "the task should see the first SIGTERM through its context",
+      )
+      Assertions.assertInstanceOf(SucceedTask::class.java, result.body, "the task returned, so its success is reported")
+      Assertions.assertEquals(listOf("Asking the task to stop", "Letting the JVM handle the signal"), signalLogs())
+      Assertions.assertEquals(listOf("SIGTERM"), pollSignals(recorded, 1), "the second SIGTERM should reach the JVM's handler")
+    }
+
+  @Test
+  @DisplayName(
+    "Should let a task that checks for a stop request only through a copy of its context see SIGTERM, " +
+      "and let the copy still see it after the task returns",
+  )
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun copiedContextSeesStopSignal() =
+    withSignalRecorder("TERM") { recorded ->
+      checkedCopy = null
+      dispatchTask(
+        CopyCheckingTask::class.java,
+        whileRunning = {
+          awaitTaskEvent("started")
+          Signal.raise(Signal("TERM"))
+        },
+      )
+
+      Assertions.assertEquals(listOf("awaitCancellation returned true"), taskEvents.toList(), "the copy should see the SIGTERM")
+      Assertions.assertEquals(listOf("Asking the task to stop"), signalLogs())
+      Assertions.assertNull(recorded.poll(), "the SDK should keep the SIGTERM for the task")
+      val copy = checkNotNull(checkedCopy)
+      Assertions.assertTrue(copy.isCancelled(), "the copy should still see the stop request after the task returns")
+      Assertions.assertFalse(copy.copy().isCancelled(), "a context made after the task returns should not see the stop request")
+    }
+
+  @Test
+  @DisplayName("Should not show the stop request to a context made after the task when reporting the result fails")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun stopRequestEndsWhenReportingFails() =
+    withSignalRecorder("TERM") { _ ->
+      Assertions.assertThrows(ApiError::class.java) {
+        dispatchTask(
+          StopAwareTask::class.java,
+          whileRunning = {
+            awaitTaskEvent("started")
+            Signal.raise(Signal("TERM"))
+          },
+          ack = false,
+        )
+      }
+
+      Assertions.assertEquals(listOf("Asking the task to stop"), signalLogs())
+      Assertions.assertFalse(taskContext().isCancelled(), "a context made after the failed report should not see the stop request")
+    }
+
+  @Test
+  @DisplayName("Should hold the first SIGTERM until the result is reported when the task checked for a stop request")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun signalWhileReportingIsHeld() =
+    withSignalRecorder("TERM") { recorded ->
+      val result =
+        dispatchTask(
+          CheckOnceTask::class.java,
+          beforeAck = {
+            // The runtime is still waiting for this acknowledgement, so it has not finished reporting yet.
+            Signal.raise(Signal("TERM"))
+            awaitSignalLogs(1)
+          },
+        )
+
+      Assertions.assertInstanceOf(SucceedTask::class.java, result.body)
+      Assertions.assertEquals(listOf("Asking the task to stop"), signalLogs())
+      Assertions.assertNull(recorded.poll(), "SIGTERM should reach the runtime until it has reported the result")
+      Signal.raise(Signal("TERM"))
+      Assertions.assertEquals(
+        listOf("SIGTERM"),
+        pollSignals(recorded, 1),
+        "after the result is reported, SIGTERM should go back to the handler that was there before",
+      )
+    }
+
+  @Test
+  @DisplayName("Should leave SIGTERM to the JVM when the task never checks for a stop request")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun signalGoesToTheJvmWhenTaskNeverChecks() =
+    withSignalRecorder("TERM") { recorded ->
+      val handedBack = LinkedBlockingQueue<String>()
+      dispatchTask(
+        NeverCheckingTask::class.java,
+        whileRunning = {
+          awaitTaskEvent("started")
+          Signal.raise(Signal("TERM"))
+          handedBack.put(recorded.poll(10, TimeUnit.SECONDS) ?: "nothing")
+          releaseTask.put("done")
+        },
+      )
+
+      Assertions.assertEquals("SIGTERM", handedBack.take(), "SIGTERM should reach the JVM's handler while the task runs")
+      Assertions.assertEquals(listOf("Letting the JVM handle the signal"), signalLogs())
+    }
+
+  @Test
+  @DisplayName("Should take only one of two SIGTERMs that arrive together, and leave the other to the JVM")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun takesOnlyOneOfTwoSignals() =
+    withSignalRecorder("TERM") { recorded ->
+      dispatchTask(
+        StopAwareTask::class.java,
+        whileRunning = {
+          awaitTaskEvent("started")
+          // Each handler decides whether to take its signal and then logs, which needs this lock.
+          // Holding the lock stops both handlers after they decide and before they finish.
+          synchronized(LogSender) {
+            Signal.raise(Signal("TERM"))
+            Signal.raise(Signal("TERM"))
+            awaitBlockedSignalHandlers(2)
+          }
+        },
+      )
+
+      Assertions.assertEquals(
+        listOf("Asking the task to stop", "Letting the JVM handle the signal"),
+        signalLogs().sorted(),
+        "one SIGTERM should be taken for the task and the other left to the JVM",
+      )
+      Assertions.assertEquals(listOf("SIGTERM"), pollSignals(recorded, 1), "one SIGTERM should reach the JVM's handler")
+    }
+
   @Test
   @DisplayName("Should fail when the initial frame is not startup details")
   @Timeout(value = 30, unit = TimeUnit.SECONDS)
@@ -176,6 +466,8 @@ class ServerTest {
   }
 
   private companion object {
+    val SIGNAL_EVENTS = setOf("Asking the task to stop", "Letting the JVM handle the signal")
+
     // [2, msg, null] with msg coming from
     // https://github.com/astronomer/airflow/blob/f39c8da8/task-sdk/tests/task_sdk/execution_time/test_comms.py#L73-L108
     val STARTUP_HEX =

java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CancellationTest.kt

@@ -0,0 +1,205 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk.execution
+
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.Assumptions
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.Timeout
+import sun.misc.Signal
+import sun.misc.SignalHandler
+import java.util.concurrent.BlockingQueue
+import java.util.concurrent.LinkedBlockingQueue
+import java.util.concurrent.TimeUnit
+
+/**
+ * Runs [block] while the named [signals], such as "TERM", go to a recorder instead of the JVM's own
+ * handlers. A signal that the code under test fails to catch then fails an assertion instead of
+ * stopping the test JVM.
+ */
+internal fun withSignalRecorder(
+  vararg signals: String,
+  block: (recorded: BlockingQueue<String>) -> Unit,
+) {
+  val recorded = LinkedBlockingQueue<String>()
+  val originals =
+    signals.associate { name ->
+      val signal = Signal(name)
+      signal to Signal.handle(signal) { recorded.put("SIG$name") }
+    }
+  try {
+    // The OS never delivers an ignored signal, so a test that raises one cannot pass.
+    Assumptions.assumeTrue(SignalHandler.SIG_IGN !in originals.values, "a signal this test raises is ignored in this JVM")
+    block(recorded)
+  } finally {
+    originals.forEach { (signal, handler) -> Signal.handle(signal, handler) }
+  }
+}
+
+/** Waits for the next [count] signal names that reach [queue]. */
+internal fun pollSignals(
+  queue: BlockingQueue<String>,
+  count: Int,
+): List<String?> = List(count) { queue.poll(10, TimeUnit.SECONDS) }
+
+/**
+ * Main class of the JVMs that [CancellationTest] starts.
+ *
+ * Without arguments, it prints each signal that trapStopSignals cannot catch. With `give-back`, it
+ * raises SIGTERM in a trap whose callback declines the signal. With `hold`, it raises SIGTERM in a
+ * trap whose callback takes the signal. In both modes it prints what happens next.
+ */
+object TrapProbe {
+  @JvmStatic
+  fun main(args: Array<String>) {
+    when (args.singleOrNull()) {
+      null -> {
+        trapStopSignals { false }.close()
+        LogCapture
+          .drain()
+          .filter { it.event == "Cannot catch signal, so the task cannot see it" }
+          .forEach { println("${it.arguments["signal"]} ${it.arguments["error"]?.javaClass?.name}") }
+      }
+      else -> {
+        val take = args.single() == "hold"
+        Runtime.getRuntime().addShutdownHook(Thread { println("shutdown hook ran") })
+        trapStopSignals { take }.use {
+          Signal.raise(Signal("TERM"))
+          Thread.sleep(2000)
+          println("still running")
+        }
+      }
+    }
+  }
+}
+
+class CancellationTest {
+  @Test
+  @DisplayName("Should pass SIGINT and SIGTERM to the callback until closed, then hand them back")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun trapsStopSignalsUntilClosed() =
+    withSignalRecorder("INT", "TERM") { recorded ->
+      val trapped = LinkedBlockingQueue<String>()
+      trapStopSignals {
+        trapped.put(it)
+        true
+      }.use {
+        Signal.raise(Signal("INT"))
+        Signal.raise(Signal("TERM"))
+        Assertions.assertEquals(setOf("SIGINT", "SIGTERM"), pollSignals(trapped, 2).toSet(), "the callback should get both signals")
+      }
+      Signal.raise(Signal("INT"))
+      Signal.raise(Signal("TERM"))
+      Assertions.assertEquals(
+        setOf("SIGINT", "SIGTERM"),
+        pollSignals(recorded, 2).toSet(),
+        "closing the trap should hand both signals back to the handlers that were there before",
+      )
+      Assertions.assertTrue(trapped.isEmpty(), "the callback should not get signals after the trap is closed")
+    }
+
+  @Test
+  @DisplayName("Should hand a signal the callback declines back to the previous handler, and keep later ones there")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun handsBackDeclinedSignal() =
+    withSignalRecorder("TERM") { recorded ->
+      val offered = LinkedBlockingQueue<String>()
+      trapStopSignals {
+        offered.put(it)
+        false
+      }.use {
+        Signal.raise(Signal("TERM"))
+        Assertions.assertEquals(listOf("SIGTERM"), pollSignals(recorded, 1), "a declined signal should reach the previous handler")
+        Signal.raise(Signal("TERM"))
+        Assertions.assertEquals(listOf("SIGTERM"), pollSignals(recorded, 1), "after declining, the trap should leave SIGTERM alone")
+      }
+      Assertions.assertEquals(listOf("SIGTERM"), offered.toList(), "only the first signal should reach the callback")
+    }
+
+  @Test
+  @DisplayName("Should hand a signal back to the previous handler when the callback throws")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun handsBackSignalWhenCallbackThrows() =
+    withSignalRecorder("TERM") { recorded ->
+      trapStopSignals { throw IllegalStateException("the log channel is closed") }.use {
+        Signal.raise(Signal("TERM"))
+        Assertions.assertEquals(listOf("SIGTERM"), pollSignals(recorded, 1), "a failing callback should not lose the signal")
+      }
+    }
+
+  @Test
+  @DisplayName("Should let the JVM run its shutdown hooks and exit with 143 when the callback declines SIGTERM")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun declinedSignalStopsTheJvm() {
+    val (exitCode, output) = runTrapProbe("give-back")
+
+    Assertions.assertEquals(143, exitCode, "the JVM should exit as it does on SIGTERM without a trap")
+    Assertions.assertEquals(listOf("shutdown hook ran"), output)
+  }
+
+  @Test
+  @DisplayName("Should keep the JVM running when the callback takes SIGTERM")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun takenSignalKeepsTheJvmRunning() {
+    val (exitCode, output) = runTrapProbe("hold")
+
+    Assertions.assertEquals(0, exitCode)
+    Assertions.assertEquals(listOf("still running", "shutdown hook ran"), output, "the hook should run only at the normal exit")
+  }
+
+  @Test
+  @DisplayName("Should warn and leave the signals alone when the JVM is started with -Xrs")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun warnsWhenTheJvmKeepsTheSignals() {
+    Assertions.assertEquals(
+      0 to listOf("SIGINT java.lang.IllegalArgumentException", "SIGTERM java.lang.IllegalArgumentException"),
+      runTrapProbe(jvmOptions = listOf("-Xrs")),
+    )
+  }
+
+  @Test
+  @DisplayName("Should warn and leave the signals alone when the runtime has no jdk.unsupported module")
+  @Timeout(value = 60, unit = TimeUnit.SECONDS)
+  fun warnsWhenSunMiscSignalIsMissing() {
+    Assertions.assertEquals(
+      0 to listOf("SIGINT java.lang.NoClassDefFoundError", "SIGTERM java.lang.NoClassDefFoundError"),
+      runTrapProbe(jvmOptions = listOf("--limit-modules", "java.base")),
+    )
+  }
+
+  private fun runTrapProbe(
+    vararg args: String,
+    jvmOptions: List<String> = emptyList(),
+  ): Pair<Int, List<String>> {
+    val java =
+      ProcessHandle
+        .current()
+        .info()
+        .command()
+        .orElseThrow()
+    val process =
+      ProcessBuilder(listOf(java) + jvmOptions + listOf("-cp", System.getProperty("java.class.path"), TrapProbe::class.java.name) + args)
+        .redirectErrorStream(true)
+        .start()
+    val output = process.inputStream.bufferedReader().readLines()
+    return process.waitFor() to output
+  }
+}

java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/LogSenderTest.kt

@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk.execution
+
+import io.ktor.utils.io.ByteChannel
+import io.ktor.utils.io.readLine
+import kotlinx.coroutines.runBlocking
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.jsonObject
+import kotlinx.serialization.json.jsonPrimitive
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.Timeout
+import java.util.concurrent.TimeUnit
+import kotlin.concurrent.thread
+
+class LogSenderTest {
+  @Test
+  @DisplayName("Should make a sender wait while another thread sends a message")
+  @Timeout(value = 30, unit = TimeUnit.SECONDS)
+  fun sendWaitsForAnotherSend() =
+    assertSecondSenderWaits { channel, message ->
+      LogSender.configure(channel)
+      thread { LogSender.send(message) }
+    }
+
+  @Test
+  @DisplayName("Should make a sender wait while configure writes the messages queued before it")
+  @Timeout(value = 30, unit = TimeUnit.SECONDS)
+  fun sendWaitsForConfigure() =
+    assertSecondSenderWaits { channel, message ->
+      LogSender.messages.add(message)
+      thread { LogSender.configure(channel) }
+    }
+
+  /**
+   * Checks that a sender waits while another thread writes to [LogSender]'s channel.
+   *
+   * [startFirst] starts the thread that writes first. That thread gets a message larger than the
+   * 1 MiB that a ByteChannel buffers, so the thread stops in the middle of the write until the test
+   * reads the channel.
+   */
+  private fun assertSecondSenderWaits(startFirst: (channel: ByteChannel, message: LogMessage) -> Thread) {
+    val channel = ByteChannel(autoFlush = true)
+    LogSender.messages.clear()
+    try {
+      val large = "x".repeat(2 * 1024 * 1024)
+      val first = startFirst(channel, LogMessage(large, emptyMap(), "test", Level.INFO))
+      val parked = setOf(Thread.State.WAITING, Thread.State.TIMED_WAITING)
+      while (first.state !in parked || first.stackTrace.none { it.methodName == "sendTo" }) {
+        Assertions.assertNotEquals(Thread.State.TERMINATED, first.state, "the first message no longer fills the channel")
+        Thread.sleep(10)
+      }
+      val second = thread { LogSender.send(LogMessage("small", emptyMap(), "test", Level.INFO)) }
+      while (second.state !in parked + Thread.State.BLOCKED + Thread.State.TERMINATED) Thread.sleep(10)
+
+      Assertions.assertEquals(Thread.State.BLOCKED, second.state, "the second sender should wait for the first writer")
+
+      val events =
+        runBlocking {
+          List(2) {
+            Json
+              .parseToJsonElement(channel.readLine()!!)
+              .jsonObject["event"]!!
+              .jsonPrimitive.content
+          }
+        }
+      first.join()
+      second.join()
+      Assertions.assertEquals(listOf(large.length, 5), events.map { it.length }, "both messages should arrive whole, in order")
+    } finally {
+      // Once its channel is closed, LogSender queues messages again. Later tests rely on this queuing.
+      channel.cancel(null)
+    }
+  }
+}