This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new eda782f76ea Make Java SDK context signal-aware (#74690)
eda782f76ea is described below
commit eda782f76eaeebcd41f47869ec249a58c6db8fb4
Author: PoAn Yang <[email protected]>
AuthorDate: Thu Oct 22 16:23:07 2026 +0800
Make Java SDK context signal-aware (#74690)
* Make Java SDK context signal-aware
Signed-off-by: PoAn Yang <[email protected]>
* Share the stop request with every context made during a task
Signed-off-by: PoAn Yang <[email protected]>
* Update java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt
---------
Signed-off-by: PoAn Yang <[email protected]>
Co-authored-by: Jason(Zhe-You) Liu
<[email protected]>
---
.../language-sdks/java.rst | 47 ++++
.../org/apache/airflow/e2e/TestBundleBuilder.java | 33 ++-
.../src/resources/dags/java_test_dags.py | 12 +
.../airflow_e2e_tests/e2e_test_utils/clients.py | 8 +
.../java_sdk_tests/test_java_sdk_dag.py | 56 ++++
.../main/kotlin/org/apache/airflow/sdk/Context.kt | 50 ++++
.../main/kotlin/org/apache/airflow/sdk/Server.kt | 29 +-
.../apache/airflow/sdk/execution/Cancellation.kt | 134 ++++++++++
.../org/apache/airflow/sdk/execution/Logger.kt | 4 +
.../kotlin/org/apache/airflow/sdk/ContextTest.kt | 14 +
.../kotlin/org/apache/airflow/sdk/ServerTest.kt | 292 +++++++++++++++++++++
.../airflow/sdk/execution/CancellationTest.kt | 205 +++++++++++++++
.../apache/airflow/sdk/execution/LogSenderTest.kt | 94 +++++++
13 files changed, 973 insertions(+), 5 deletions(-)
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 93178c747cd..f1ba8c35874 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,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
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 39c4bd46afc..11d5dfe4783 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 @@ 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) {
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 8a4818923f0..ed232ced648 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()
+
+
[email protected](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 821765ad212..ee895f0fdc3 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 @@ 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(
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 ec2f1408f91..3716b298f6f 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_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."""
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 69badbd49f6..45ca21e3713 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,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,
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 5847f5b7326..c604e6e4cae 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,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(
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 00000000000..191286d4481
--- /dev/null
+++
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Cancellation.kt
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk.execution
+
+import sun.misc.Signal
+import sun.misc.SignalHandler
+import java.time.Duration
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicReference
+
+/**
+ * Records whether Airflow has asked the running task to stop, and whether the
task checks for that
+ * request.
+ */
+internal class Cancellation {
+ private val claimed = AtomicBoolean()
+ private val requested = CountDownLatch(1)
+
+ /**
+ * True once the task has checked for a stop request or waited for one,
which shows that the task
+ * can stop on its own.
+ */
+ @Volatile
+ var isWatched = false
+ private set
+
+ val isRequested: Boolean
+ get() = requested.count == 0L
+
+ /** Returns true to the first caller only, so the SDK can take at most one
signal for the task. */
+ fun claim(): Boolean = claimed.compareAndSet(false, true)
+
+ fun request() = requested.countDown()
+
+ /**
+ * Returns [isRequested] for the task. Because the task has now checked,
this call also sets
+ * [isWatched].
+ */
+ fun check(): Boolean {
+ watch()
+ return isRequested
+ }
+
+ fun await(timeout: Duration): Boolean {
+ watch()
+ return requested.await(TimeUnit.NANOSECONDS.convert(timeout),
TimeUnit.NANOSECONDS)
+ }
+
+ private fun watch() {
+ if (!isWatched) isWatched = true
+ }
+
+ companion object {
+ /**
+ * The stop request of the task that the server is running in this JVM.
The server sets this field
+ * before the task starts. When the server is done with the task, it sets
the field back to null.
+ * Every [org.apache.airflow.sdk.Context] made while the task runs keeps
this instance, so a copy
+ * of a context made during the task also sees the same request. A single
field is enough because
+ * the server runs one task at a time in a JVM, and a stop signal reaches
the whole JVM, not one
+ * task.
+ */
+ @Volatile
+ var running: Cancellation? = null
+
+ /**
+ * Returns [running]. When the server is not running a task, for example
in a unit test that calls
+ * a task directly, it returns a new [Cancellation] that nobody asks to
stop.
+ */
+ fun current(): Cancellation = running ?: Cancellation()
+ }
+}
+
+private val logger = Logger(Cancellation::class)
+
+/**
+ * Opens a trap that passes SIGINT and SIGTERM to [onSignal] instead of the
JVM. The trap stays open
+ * until the caller closes the returned handle.
+ *
+ * When [onSignal] returns false or throws, the JVM gets its own handler back
and receives the
+ * signal again, so the JVM handles the signal as it would have without the
trap. Only one trap can
+ * be open at a time, because closing a trap puts back the handlers that were
installed when that
+ * trap opened.
+ */
+internal fun trapStopSignals(onSignal: (signal: String) -> Boolean):
AutoCloseable {
+ val restores = listOf("INT", "TERM").mapNotNull { trapSignal(it, onSignal) }
+ return AutoCloseable { restores.forEach { restore -> restore() } }
+}
+
+private fun trapSignal(
+ name: String,
+ onSignal: (signal: String) -> Boolean,
+): (() -> Unit)? =
+ try {
+ val signal = Signal(name)
+ val previous = AtomicReference<SignalHandler>()
+ previous.set(
+ Signal.handle(signal) {
+ if (!runCatching { onSignal("SIG$name") }.getOrDefault(false)) {
+ // A signal can arrive before Signal.handle has returned the handler
it replaced.
+ while (previous.get() == null) Thread.onSpinWait()
+ Signal.handle(signal, previous.get())
+ Signal.raise(signal)
+ }
+ },
+ )
+ ({ Signal.handle(signal, previous.get()) })
+ } catch (e: IllegalArgumentException) {
+ // The JVM refuses the signal, for example when it was started with -Xrs.
+ logger.warning("Cannot catch signal, so the task cannot see it",
mapOf("signal" to "SIG$name", "error" to e))
+ null
+ } catch (e: LinkageError) {
+ // sun.misc.Signal lives in the jdk.unsupported module, which a jlink
runtime image can leave out.
+ logger.warning("Cannot catch signal, so the task cannot see it",
mapOf("signal" to "SIG$name", "error" to e))
+ null
+ }
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 975cc5ac768..c9063fc9125 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
@@ -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) {
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 2c48be12443..3208c5f412f 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 7c7ab91b2ab..b66ffa87e18 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,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 =
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 00000000000..cc6250d1f1c
--- /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<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
+ }
+}
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 00000000000..64da9e24b97
--- /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)
+ }
+ }
+}
