FrankYang0529 · Oct 14, 2026 · jason810496 · Oct 20, 2026 · jason810496 · Oct 20, 2026
diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
@@ -587,6 +587,52 @@ 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``. 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,15 +19,18 @@
 
 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.*;
 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 void execute(@NotNull Context context, Client client) {
     }
   }
 
+  /**
+   * 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 Iterable<DagDef> getDags() {
     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 @@ def get_task_instances(self, dag_id: str, run_id: str):
             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_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 @@ def test_scratch_variable_deleted_by_java_task_is_gone(self, variable_write_run:
         )
 
 
+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,10 +129,49 @@ data class Context(
   /** Registration of the executing task; resolves wired data inputs. */
   internal var taskDef: TaskDef? = null
 
+  internal var cancellation = Cancellation()
+
+  /**
+   * 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.
+   *
+   * A context made with [copy] does not see the 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.
+   *
+   * @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,
       taskDef: TaskDef? = null,
+      cancellation: Cancellation = Cancellation(),
     ): Context =
       Context(
         dagRun =
@@ -147,6 +188,9 @@ data class Context(
             )
           },
         ti = with(request.ti) { TaskInstance(dagId, runId, taskId, mapIndex, tryNumber) },
-      ).also { it.taskDef = taskDef }
+      ).also {
+        it.taskDef = taskDef
+        it.cancellation = cancellation
+      }
   }
 }
@@ -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,26 @@ class CoordinatorServer(
     startup: StartupDetails,
     coordinator: CoordinatorComm,
   ) {
-    val result = runTask(bundle, startup, coordinator)
-    coordinator.communicate<Unit>(result)
+    val cancellation = Cancellation()
+    // 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, cancellation)
+      coordinator.communicate<Unit>(result)
+    }
   }
 
   private suspend fun parseDagsAndReport(