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>
2026-11-04 07:12:48 +00:00 · 2026-10-22 16:23:07 +08:00
1 parent c558f4a60f
commit eda782f76e
@@ -587,6 +587,53 @@ An ``InputTask`` whose type argument is not a concrete ``TaskInput`` fails when
 rather than mid-run.  Plain ``Task`` remains the right interface for a task the Dag file
 calls with no arguments.

+.. _java-sdk/cancellation:
+
+Stopping when Airflow asks
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Airflow asks a running task to stop when, for example, a user marks its task instance as failed. The
+supervisor, the Airflow process that started the task's JVM, then sends the JVM ``SIGTERM``. If the JVM
+is still running after a grace period, the supervisor sends ``SIGKILL``. When a user marks the task
+instance failed or success, or clears it, the grace period is
+:ref:`core.killed_task_cleanup_time <config:core__killed_task_cleanup_time>` seconds.
+
+A task can check ``Context.isCancelled()``, or call ``Context.awaitCancellation(timeout)`` where it
+would otherwise wait with ``Thread.sleep``. Unlike ``Thread.sleep``, ``awaitCancellation`` returns as
+soon as Airflow asks the task to stop. Once a task has called either method, the SDK catches the first
+``SIGTERM`` or ``SIGINT`` and the JVM keeps running. The task can then stop its work, such as a process
+it started, and throw. Keep checking until the task returns, because the first signal no longer stops
+the JVM:
+
+.. code-block:: java
+
+    public class ReportTask implements Task {
+      @Override
+      public void execute(Context context, Client client) throws Exception {
+        Process report = new ProcessBuilder("./build-report.sh").inheritIO().start();
+        while (!context.isCancelled()) {
+          if (report.waitFor(1, TimeUnit.SECONDS)) {
+            if (report.exitValue() != 0) {
+              throw new IllegalStateException("build-report.sh failed");
+            }
+            return;
+          }
+        }
+        report.destroy();
+        throw new IllegalStateException("Airflow asked the task to stop");
+      }
+    }
+
+Throw rather than return when the task stops early, because returning reports success for work the
+task did not finish. If a user changed the task instance's state while the task ran, for example by
+marking the task instance failed, Airflow keeps that state whatever the task reports.
+
+A task that has never checked cannot stop on its own, so for that task the JVM handles the signal as it
+does without the SDK: shutdown hooks run and the JVM exits. The JVM handles a second signal the same
+way, so sending the signal a second time stops a task that does not stop when asked. For a task that
+has checked, shutdown hooks run when the JVM exits after the task returns. The shutdown hooks do not run
+at all if the supervisor has to kill the JVM with ``SIGKILL``.
+
 .. _java-sdk/native-dags:

 Native Java Dags
@@ -19,6 +19,9 @@

 package org.apache.airflow.e2e;

+import static java.lang.System.Logger.Level.INFO;
+
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import org.apache.airflow.sdk.*;
@@ -26,8 +29,8 @@ import org.jetbrains.annotations.NotNull;

 /**
 * Bundle for the runner-behaviour E2E tests: deliberately broken task classes that exercise
- * instantiation failures, and tasks that round-trip Airflow Variables and task state through the
- * supervisor.
+ * instantiation failures, tasks that round-trip Airflow Variables and task state through the
+ * supervisor, and a task that runs until Airflow asks it to stop.
 */
 public class TestBundleBuilder implements BundleBuilder {
  public static class MissingNoArgConstructor implements Task {
@@ -90,6 +93,28 @@ public class TestBundleBuilder implements BundleBuilder {
    }
  }

+  /**
+   * Runs until Airflow asks it to stop. The E2E test makes that request by marking the task instance
+   * failed.
+   */
+  public static class WaitForStopRequest implements Task {
+    private static final System.Logger log = System.getLogger(WaitForStopRequest.class.getName());
+
+    public void execute(@NotNull Context context, Client client) throws InterruptedException {
+      // Check before writing the log line that the test waits for. The test makes Airflow send the
+      // stop signal after it sees that line, so by then the task has checked and the SDK keeps the
+      // stop request.
+      if (context.isCancelled()) {
+        throw new IllegalStateException("Airflow asked the task to stop before it started");
+      }
+      log.log(INFO, "Waiting for a stop request");
+      if (!context.awaitCancellation(Duration.ofMinutes(5))) {
+        throw new IllegalStateException("No stop request arrived within five minutes");
+      }
+      log.log(INFO, "Stopping at Airflow's request");
+    }
+  }
+
  @NotNull
  @Override
  public Iterable<DagDef> getDags() {
@@ -100,7 +125,9 @@ public class TestBundleBuilder implements BundleBuilder {
    variableWrite.addTask("write_and_delete", WriteAndDeleteVariable.class);
    var taskState = new DagDef("java_task_state");
    taskState.addTask("roundtrip_task_state", RoundtripTaskState.class);
-    return List.of(uninstantiable, variableWrite, taskState);
+    var stopRequest = new DagDef("java_stop_request");
+    stopRequest.addTask("wait_for_stop_request", WaitForStopRequest.class);
+    return List.of(uninstantiable, variableWrite, taskState, stopRequest);
  }

  public static void main(String[] args) {
@@ -60,3 +60,15 @@ def java_task_state():


 java_task_state()
+
+
+@task.stub(queue="java-test")
+def wait_for_stop_request(): ...
+
+
+@dag(dag_id="java_stop_request")
+def java_stop_request():
+    wait_for_stop_request()
+
+
+java_stop_request()
@@ -249,6 +249,14 @@ class AirflowClient:
            endpoint=f"dags/{dag_id}/dagRuns/{run_id}/taskInstances",
        )

+    def set_task_instance_state(self, dag_id: str, run_id: str, task_id: str, state: str):
+        """Set the state of a task instance, as marking it in the UI does."""
+        return self._make_request(
+            method="PATCH",
+            endpoint=f"dags/{dag_id}/dagRuns/{run_id}/taskInstances/{task_id}",
+            json={"new_state": state},
+        )
+
    def get_hitl_detail(self, dag_id: str, run_id: str, task_id: str, map_index: int = -1):
        """Get the Human-in-the-loop request of a task instance."""
        return self._make_request(
@@ -98,6 +98,8 @@ _VARIABLE_WRITE_DAG_ID = "java_variable_write"
 _VARIABLE_WRITE_DESCRIPTION = "written by the Java SDK e2e test"
 _TASK_STATE_DAG_ID = "java_task_state"
 _TASK_STATE_TASK_ID = "roundtrip_task_state"
+_STOP_REQUEST_DAG_ID = "java_stop_request"
+_STOP_REQUEST_TASK_ID = "wait_for_stop_request"


@dataclass
@@ -427,6 +429,60 @@ class TestJavaSDKVariableWrite:
        )


+class TestJavaSDKStopRequest:
+    """Verify a running Java task sees Airflow's request to stop.
+
+    Marking a running task instance failed makes the supervisor send the task's JVM SIGTERM.
+    ``WaitForStopRequest`` in the java-test-bundle fixture project checks ``Context.isCancelled``
+    and then waits with ``Context.awaitCancellation``, so the SDK keeps the JVM running. The task
+    then returns from the wait and logs that it is stopping. A JVM that exited on the signal could
+    not write that log line.
+    """
+
+    def test_running_task_stops_at_airflows_request(self):
+        client = AirflowClient()
+        run_id = client.trigger_dag(
+            _STOP_REQUEST_DAG_ID, json={"logical_date": datetime.now(UTC).isoformat()}
+        )["dag_run_id"]
+        run = _CompletedRun(client=client, dag_id=_STOP_REQUEST_DAG_ID, run_id=run_id)
+
+        # The SDK keeps the signal only for a task that has checked, so mark the task instance failed only
+        # after the task logs that it is waiting.
+        deadline = time.monotonic() + _JAVA_TASK_TIMEOUT
+        while (
+            state := _get_task_state(client, _STOP_REQUEST_DAG_ID, run_id, _STOP_REQUEST_TASK_ID)
+        ) != "running":
+            assert state in {None, "scheduled", "queued"}, f"the task ended before it ran; state: {state!r}"
+            assert time.monotonic() < deadline, f"the task never started; state: {state!r}"
+            time.sleep(3)
+        waiting, records = run.wait_for_log_record(
+            _STOP_REQUEST_TASK_ID, 1, lambda r: r.get("event") == "Waiting for a stop request"
+        )
+        assert waiting is not None, (
+            f"the task never started waiting; events seen: {[r.get('event') for r in records]}"
+        )
+
+        client.set_task_instance_state(_STOP_REQUEST_DAG_ID, run_id, _STOP_REQUEST_TASK_ID, "failed")
+
+        stopping, records = run.wait_for_log_record(
+            _STOP_REQUEST_TASK_ID, 1, lambda r: r.get("event") == "Stopping at Airflow's request"
+        )
+        assert stopping is not None, (
+            f"the task should see the request to stop; events seen: {[r.get('event') for r in records]}"
+        )
+        caught = [r.get("signal") for r in records if r.get("event") == "Asking the task to stop"]
+        assert caught == ["SIGTERM"], f"the SDK should log the SIGTERM it caught; records: {records}"
+        assert client.wait_for_dag_run(_STOP_REQUEST_DAG_ID, run_id, timeout=_JAVA_TASK_TIMEOUT) == "failed"
+        assert _get_task_state(client, _STOP_REQUEST_DAG_ID, run_id, _STOP_REQUEST_TASK_ID) == "failed", (
+            "Airflow should keep the state the test set"
+        )
+
+
+def _get_task_state(client: AirflowClient, dag_id: str, run_id: str, task_id: str) -> str | None:
+    task_instances = client.get_task_instances(dag_id, run_id).get("task_instances", [])
+    return next((ti.get("state") for ti in task_instances if ti["task_id"] == task_id), None)
+
+
 class TestJavaSDKInterfaceExample:
    """Verify tasks written against ``InputTask`` receive their bound arguments."""

@@ -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,
@@ -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(
@@ -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
+  }
@@ -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) {
@@ -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")
+  }
 }
@@ -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 =
@@ -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
+  }
+}
@@ -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)
+    }
+  }
+}