diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index e8e63811c9617..6f17dfadabc02 100644 --- 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 ` 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 diff --git a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java index 39c4bd46afc30..11d5dfe478398 100644 --- a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java +++ b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java @@ -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 @@ /** * 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 getDags() { @@ -100,7 +125,9 @@ public Iterable 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) { diff --git a/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py b/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py index 8a4818923f040..ed232ced648e4 100644 --- a/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py +++ b/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py @@ -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() diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py index 821765ad21202..ee895f0fdc3e9 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py @@ -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( diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py index ec2f1408f910a..3716b298f6fc5 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py @@ -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.""" diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt index 69badbd49f65c..61f63a1e37f72 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt +++ b/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,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 + } } } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt index 5847f5b7326ef..d98f70d1a5032 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt +++ b/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,26 @@ class CoordinatorServer( startup: StartupDetails, coordinator: CoordinatorComm, ) { - val result = runTask(bundle, startup, coordinator) - coordinator.communicate(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(result) + } } private suspend fun parseDagsAndReport( diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Cancellation.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Cancellation.kt new file mode 100644 index 0000000000000..d8b287e6bea5b --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Cancellation.kt @@ -0,0 +1,115 @@ +/* + * 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 + } +} + +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() + 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 + } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt index 04757386aa46b..a6365f4dcd92c 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt @@ -223,10 +223,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 = ConcurrentLinkedDeque() + @Synchronized fun configure(channel: ByteWriteChannel) { writer = channel if (!channel.isClosedForWrite) { @@ -236,6 +239,7 @@ internal object LogSender { } } + @Synchronized fun send(message: LogMessage) { val channel = writer if (channel == null || channel.isClosedForWrite) { diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt index 54cdb2b3911d9..9adac0979cc57 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt @@ -98,6 +98,7 @@ internal object TaskRunner { bundle: Bundle, request: StartupDetails, client: Client, + cancellation: Cancellation = Cancellation(), ): Any { val taskDef = bundle.taskDef(request.ti.dagId, request.ti.taskId) @@ -142,7 +143,7 @@ internal object TaskRunner { return TaskResult.failure(request.tiContext.shouldRetry) } return try { - val context = Context.from(request, taskDef) + val context = Context.from(request, taskDef, cancellation) val decider = taskDef.decider if (decider == null) instance.execute(context, client) else decide(decider, taskDef, instance, context, client) TaskResult.success() @@ -159,10 +160,12 @@ internal fun runTask( bundle: Bundle, request: StartupDetails, comm: CoordinatorComm, -): Any = TaskRunner.runTask(bundle, request, Client(request, CoordinatorClient(comm))) + cancellation: Cancellation, +): Any = TaskRunner.runTask(bundle, request, Client(request, CoordinatorClient(comm)), cancellation) internal fun runTask( bundle: Bundle, request: StartupDetails, client: Client, -) = TaskRunner.runTask(bundle, request, client) + cancellation: Cancellation = Cancellation(), +) = TaskRunner.runTask(bundle, request, client, cancellation) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ContextTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ContextTest.kt index 2c48be1244314..3208c5f412f83 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ContextTest.kt +++ b/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") + } } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt index 7c7ab91b2ab22..3dfda1eba9c04 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt +++ b/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,50 @@ private class ServerNoopTask : Task { ) = Unit } +/** What the tasks below report to the test that runs them. */ +private val taskEvents = LinkedBlockingQueue() + +/** Releases [NeverCheckingTask]. */ +private val releaseTask = LinkedBlockingQueue() + +/** + * 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()}") + } +} + +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 +155,173 @@ 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. + */ + private fun dispatchTask( + taskClass: Class, + whileRunning: () -> Unit = {}, + beforeAck: () -> Unit = {}, + ): 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(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() + toServer.writeFrame(ackFrame(result.id)) + } + } + + runBlocking { server.dispatchTask(Bundle(listOf(DagDef("c").addTask(TaskDef("a", taskClass)))), comm) } + supervisor.join() + comm.close() + return reported.take() + } + + /** The events the runtime logged for the stop signals it got, in order. */ + private fun signalLogs(): List = 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 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() + 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 +395,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 = diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CancellationTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CancellationTest.kt new file mode 100644 index 0000000000000..cc6250d1f1ca2 --- /dev/null +++ b/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) -> Unit, +) { + val recorded = LinkedBlockingQueue() + 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, + count: Int, +): List = 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) { + 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() + 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() + 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 = emptyList(), + ): Pair> { + 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 + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/LogSenderTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/LogSenderTest.kt new file mode 100644 index 0000000000000..64da9e24b97e0 --- /dev/null +++ b/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) + } + } +}