jason810496 commented on code in PR #75538: URL: https://github.com/apache/airflow/pull/75538#discussion_r4534588258
########## airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst: ########## @@ -748,6 +748,117 @@ the defaults of the Dag's params, and the conf of a scheduled run is empty. For accepts and the cases where it throws, see ``DagRun.getConf`` in the `Java SDK API reference <https://airflow.apache.org/docs/java-sdk/stable/>`__. +.. _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. + +By default the JVM handles the signal as it does without the SDK: shutdown hooks run and the JVM exits. +A process the task started gets the same signals, because the supervisor signals the JVM's whole process +group. So a task does not need to stop a process that exits on ``SIGTERM`` and stays in that group. + +To stop anything else, such as a job in an external system, implement ``StoppableTask``. This interface +adds ``onKill``, the Java counterpart of an operator's ``on_kill``: + +.. code-block:: java + + public class ReportTask implements StoppableTask { + private final ReportService service = new ReportService(); + private volatile String jobId; + + @Override + public void execute(Context context, Client client) throws Exception { + jobId = service.submit(); + if (context.isCancelled()) { + service.cancel(jobId); + throw new IllegalStateException("Airflow asked the task to stop"); + } + service.waitFor(jobId); // Throws if the job fails or is cancelled. + } + + @Override + public void onKill(Context context) throws Exception { + if (jobId != null) { + service.cancel(jobId); + } + } + } + +With annotations, mark the method that stops the task with ``@Airflow.OnKill`` and name the task. The +method can take a ``Context``, and the annotation processor generates a ``StoppableTask`` that calls it on +the same object as the task method: + +.. code-block:: java + + private final ReportService service = new ReportService(); + private volatile String jobId; + + @Airflow.Task(id = "build_report", retries = 2) + public void buildReport(Context context) throws Exception { + jobId = service.submit(); + if (context.isCancelled()) { + service.cancel(jobId); + throw new IllegalStateException("Airflow asked the task to stop"); + } + service.waitFor(jobId); // Throws if the job fails or is cancelled. + } + + @Airflow.OnKill(task = "build_report") + public void stopReport() throws Exception { + if (jobId != null) { + service.cancel(jobId); + } + } + +Once the SDK has started such a task, it catches the first ``SIGTERM`` or ``SIGINT`` and the JVM keeps +running. ``Context.isCancelled()`` then returns true and the SDK calls ``onKill``: + +* ``onKill`` runs once, on its own daemon thread, while the task keeps running. So ``onKill`` can stop a task Review Comment: Let's move these detailed manual to the docstring and link the API reference to keep the user facing doc more concise. ########## java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt: ########## @@ -210,7 +233,14 @@ class CoordinatorServer( aSocket(selector).tcp().connect(logs).use { socket -> logger.debug("Connected logs", mapOf("addr" to logs)) LogSender.configure(socket.openWriteChannel(autoFlush = true)) - deferral.await() + val hook = Thread({ sendQueuedLogs(socket::close, socket.socketContext, logFlushTimeout) }, "Airflow log flush") + Runtime.getRuntime().addShutdownHook(hook) + try { + deferral.await() + } finally { + // The use block closes the socket even when serving fails, so flush while it is still open. + LogSender.detachAndFlush() Review Comment: This flush has no idle timeout, so if the supervisor stops reading the log socket, `serveAsync` never returns even though the result is already reported. The hook added above is also never removed, so every `serveAsync` call leaves one behind. Calling `sendQueuedLogs` here bounds the wait, and after that the hook has nothing left to do: ```suggestion // The use block closes the socket even when serving fails, so flush while it is still open. sendQueuedLogs(socket::close, socket.socketContext, logFlushTimeout) try { Runtime.getRuntime().removeShutdownHook(hook) } catch (e: IllegalStateException) { // The JVM is already shutting down, and the hook is running or about to run. } ``` ########## java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Cancellation.kt: ########## @@ -0,0 +1,169 @@ +/* + * 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.apache.airflow.sdk.Context +import org.apache.airflow.sdk.StoppableTask +import sun.misc.Signal +import sun.misc.SignalHandler +import java.time.Duration +import java.util.concurrent.CompletableFuture +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. For a [StoppableTask], also catches the + * signal that carries that request. + */ +internal class Cancellation : AutoCloseable { + private val claimed = AtomicBoolean() + private val requested = CountDownLatch(1) + private val signalHandled = CompletableFuture<Unit>() + + @Volatile + private var held = false + + @Volatile + private var trap: AutoCloseable? = null + + @Volatile + private var hook: (() -> Unit)? = null + + fun check(): Boolean { + throwIfTaskCannotSeeRequest() + return requested.count == 0L + } + + fun await(timeout: Duration): Boolean { + throwIfTaskCannotSeeRequest() + return requested.await(TimeUnit.NANOSECONDS.convert(timeout), TimeUnit.NANOSECONDS) + } + + private fun throwIfTaskCannotSeeRequest() { + check(held) { + "Only a StoppableTask can see a stop request. Implement StoppableTask, or give an annotated task an @Airflow.OnKill method" + } + } + + /** Catches the first stop signal for [task] until [close], or throws [CannotCatchStopSignalsException]. */ + fun hold( + task: StoppableTask, + context: Context, + ) { + val classLoader = Thread.currentThread().contextClassLoader + hook = { + // On Java 11 and 17, a signal handler thread has no context class loader, and libraries can need one. + Thread.currentThread().contextClassLoader = classLoader + task.onKill(context) + } + // Set first: a signal can run onKill, which may check for a stop request, before trapStopSignals returns. + held = true + trap = trapStopSignals(::takeSignal) + ?: throw CannotCatchStopSignalsException( + "The SDK cannot catch SIGTERM or SIGINT in this JVM, so the task cannot see a stop request. The JVM might run with -Xrs, " + + "or without the jdk.unsupported module", + ) + } + + /** + * Stops later signals from calling onKill and waits for a call that already started. + * The wait has no timeout because the supervisor's SIGKILL ends a call that hangs. + */ + fun finish() { + hook = null + // join() keeps waiting if the thread is interrupted, and a task can leave the thread's interrupt flag set. + if (claimed.get()) signalHandled.join() + } Review Comment: If `onKill` hangs after `execute` has returned, `finish()` blocks here and the result never goes out. Could the wait before the report be bounded? The SDK still shouldn't cut `onKill` short: in the server-terminated path the supervisor allows up to `killed_task_cleanup_time` (60 s by default) before SIGKILL, so after the report it can keep waiting without a limit: ```suggestion /** * Stops later signals from calling onKill and waits for a call that already started, at most * [onKillTimeout], so a hung call cannot block the result report. [awaitOnKill] waits for the rest. */ fun finish() { hook = null if (!claimed.get()) return // join() keeps waiting if the thread is interrupted, and a task can leave the thread's interrupt flag set. signalHandled.copy().completeOnTimeout(Unit, onKillTimeout.toNanos(), TimeUnit.NANOSECONDS).join() if (!signalHandled.isDone) logger.warning("Stopped waiting for the task's onKill, so the SDK sends the result now") } /** Waits for a started onKill without a limit. Signal handler threads are daemons, so the JVM would cut it at exit. */ fun awaitOnKill() { if (claimed.get()) signalHandled.join() } ``` ```kotlin internal class Cancellation( private val onKillTimeout: Duration = Duration.ofSeconds(30), ) : AutoCloseable { // Server.kt, runTaskAndReport try { coordinator.communicate<Unit>(result) } finally { cancellation.awaitOnKill() } ``` ########## java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Cancellation.kt: ########## @@ -0,0 +1,169 @@ +/* + * 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.apache.airflow.sdk.Context +import org.apache.airflow.sdk.StoppableTask +import sun.misc.Signal +import sun.misc.SignalHandler +import java.time.Duration +import java.util.concurrent.CompletableFuture +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. For a [StoppableTask], also catches the + * signal that carries that request. + */ +internal class Cancellation : AutoCloseable { + private val claimed = AtomicBoolean() + private val requested = CountDownLatch(1) + private val signalHandled = CompletableFuture<Unit>() + + @Volatile + private var held = false + + @Volatile + private var trap: AutoCloseable? = null + + @Volatile + private var hook: (() -> Unit)? = null + + fun check(): Boolean { + throwIfTaskCannotSeeRequest() + return requested.count == 0L + } + + fun await(timeout: Duration): Boolean { + throwIfTaskCannotSeeRequest() + return requested.await(TimeUnit.NANOSECONDS.convert(timeout), TimeUnit.NANOSECONDS) + } + + private fun throwIfTaskCannotSeeRequest() { + check(held) { + "Only a StoppableTask can see a stop request. Implement StoppableTask, or give an annotated task an @Airflow.OnKill method" + } + } + + /** Catches the first stop signal for [task] until [close], or throws [CannotCatchStopSignalsException]. */ + fun hold( + task: StoppableTask, + context: Context, + ) { + val classLoader = Thread.currentThread().contextClassLoader + hook = { + // On Java 11 and 17, a signal handler thread has no context class loader, and libraries can need one. + Thread.currentThread().contextClassLoader = classLoader + task.onKill(context) + } + // Set first: a signal can run onKill, which may check for a stop request, before trapStopSignals returns. + held = true + trap = trapStopSignals(::takeSignal) + ?: throw CannotCatchStopSignalsException( + "The SDK cannot catch SIGTERM or SIGINT in this JVM, so the task cannot see a stop request. The JVM might run with -Xrs, " + + "or without the jdk.unsupported module", + ) + } + + /** + * Stops later signals from calling onKill and waits for a call that already started. + * The wait has no timeout because the supervisor's SIGKILL ends a call that hangs. + */ + fun finish() { + hook = null + // join() keeps waiting if the thread is interrupted, and a task can leave the thread's interrupt flag set. + if (claimed.get()) signalHandled.join() + } + + override fun close() { + trap?.close() + } + + private fun takeSignal(signal: String): Boolean { + // A second signal means the sender does not want to wait. + if (!claimed.compareAndSet(false, true)) { + logger.info("Letting the JVM handle the signal", mapOf("signal" to signal)) + return false + } + return try { + logger.warning("Asking the task to stop", mapOf("signal" to signal)) Review Comment: After `finish()` sets `hook = null`, the first signal is still claimed here and logs "Asking the task to stop", though there's no task left to stop. Holding the signal for the report makes sense, but the log could say so: ```suggestion val hook = hook if (hook == null) { // The task has returned, and the trap stays open until the result is sent. logger.warning("Holding the signal until the SDK has sent the task's result", mapOf("signal" to signal)) } else { logger.warning("Asking the task to stop", mapOf("signal" to signal)) } ``` ########## java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/AirflowProcessor.kt: ########## @@ -854,28 +927,67 @@ class AirflowProcessor : AbstractProcessor() { } } + val owner = ClassName.get(decl.owner) + val hook = decl.onKill + if (hook != null) { + // The instance is made here rather than in a field initializer. If the constructor throws, the task + // fails and its retries apply, as they do for a task without a hook. + executeSpec + .addStatement($$"this.instance = new $T()", owner) + // If the stop request came before this point, the hook had no instance to call, so the hook could + // not stop the task method. + .beginControlFlow("if (context.isCancelled())") + .addStatement($$"throw new $T($S)", IllegalStateException::class.java, "Airflow asked the task to stop before it started") + .endControlFlow() Review Comment: The stop check runs only after argument binding, so a stop request that arrives before or during a slow bind (XCom pulls, `TaskInput`) is only noticed once binding finishes. Could the generated code also check before binding? ```suggestion executeSpec.addStatement($$"this.instance = new $T()", owner) // If the stop request came before this point, the hook had no instance to call, so the hook could // not stop the task method. executeSpec.addCancellationCheck() ``` ```kotlin // Before argument binding, with `owner` and `hook` moved up. if (hook != null) { // Binding the arguments can be slow, for example when it pulls an XCom, so check before it starts too. executeSpec.addCancellationCheck() } private fun MethodSpec.Builder.addCancellationCheck() { beginControlFlow("if (context.isCancelled())") addStatement($$"throw new $T($S)", IllegalStateException::class.java, "Airflow asked the task to stop before it started") endControlFlow() } ``` ########## java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt: ########## @@ -216,48 +229,182 @@ internal class Logger( arguments: Map<String, Any>, ) { if (!Log.isEnabledForLevel(level, name)) return - LogSender.send(LogMessage(event, arguments, name, level)) + LogSender.send(LogMessage(event, arguments, name, level), waitsForRoom) } } +// This object's lock keeps send, configure and detachAndFlush from losing or reordering lines. internal object LogSender { + internal const val PENDING_CHARS = 8 * 1024 * 1024 + private var writer: ByteWriteChannel? = null + + // Kept after detachAndFlush clears writer, so a later call can cancel the channel an earlier one is stuck on. + private var configured: ByteWriteChannel? = null internal val messages: ConcurrentLinkedDeque<LogMessage> = ConcurrentLinkedDeque() + // Only the sender thread writes, so a caller's runBlocking cannot run the task under this object's lock. + private val pending = LinkedBlockingQueue<() -> Unit>() + + // One permit per queued character. Fair, so short lines cannot keep taking room a long line waits for. + private val room = Semaphore(PENDING_CHARS, true) + + // Not System.err, which a logging bridge can route back here and whose lock a waiting caller can hold. + @Volatile + internal var errors = PrintStream(FileOutputStream(FileDescriptor.err), true) + + private var failedChannel: ByteWriteChannel? = null + + // Printing a report when a line fails could stall the lines already queued, so wait for an empty queue or a flush. + private val reports = ArrayList<String>() + + @Volatile + private var finishedAt = System.nanoTime() + + private val sender = + thread(name = "Airflow log sender", isDaemon = true) { + while (true) { + try { + val task = + pending.poll() ?: run { + printReports() + pending.take() + } + task.invoke() + } catch (e: Throwable) { + // If the sender thread ended, flush would wait forever, and so would every caller that waits for room. + } + } + } + + @Synchronized fun configure(channel: ByteWriteChannel) { writer = channel + configured = channel if (!channel.isClosedForWrite) { while (messages.isNotEmpty()) { - sendTo(channel, messages.removeFirst()) + // Don't wait for room: these messages are already in memory, and waiting would block the serving thread. + queueLine(channel, buildLine(messages.removeFirst()), chars = 0) } } } - fun send(message: LogMessage) { - val channel = writer - if (channel == null || channel.isClosedForWrite) { - messages.addLast(message) - } else { - sendTo(channel, message) + fun send( + message: LogMessage, + waitForRoom: Boolean = true, + ) { + synchronized(this) { + if (!isQueuing()) { + messages.addLast(message) + return + } + } + // Build the line now, so that it shows the arguments as they were when the caller logged them. + val line = buildLine(message) + // Only the sender thread gives back room, so it never waits. A line over the limit asks for all the room. + val chars = if (waitForRoom && Thread.currentThread() !== sender) minOf(line.length, PENDING_CHARS) else 0 + // Wait before locking, or a logging sender thread can deadlock. A fair semaphore makes acquire(0) wait too. + if (chars > 0) room.acquireUninterruptibly(chars) + synchronized(this) { + val channel = writer + if (channel == null || channel.isClosedForWrite) { + // Queuing stopped while this caller waited for room. + messages.addLast(message) + room.release(chars) + } else { + queueLine(channel, line, chars) + } } } Review Comment: After `detachAndFlush` sets `writer = null`, every later `send` appends to `messages`, which nothing writes or bounds, so lines from other shutdown hooks are lost. Could a detached `LogSender` print them to `errors` instead? The supervisor logs stderr at ERROR level, so only WARNING and above should go there: ```suggestion val detachedNow = synchronized(this) { when { isQueuing() -> false detached -> true else -> { messages.addLast(message) return } } } // Nothing writes messages after detachAndFlush, so a late line goes to errors instead. if (detachedNow) return printAfterDetach(message) // Build the line now, so that it shows the arguments as they were when the caller logged them. val line = buildLine(message) // Only the sender thread gives back room, so it never waits. A line over the limit asks for all the room. val chars = if (waitForRoom && Thread.currentThread() !== sender) minOf(line.length, PENDING_CHARS) else 0 // Wait before locking, or a logging sender thread can deadlock. A fair semaphore makes acquire(0) wait too. if (chars > 0) room.acquireUninterruptibly(chars) val detachedWhileWaiting = synchronized(this) { val channel = writer if (channel == null || channel.isClosedForWrite) { // Queuing stopped while this caller waited for room. room.release(chars) if (!detached) messages.addLast(message) detached } else { queueLine(channel, line, chars) false } } if (detachedWhileWaiting) printAfterDetach(message, line) } // The supervisor logs stderr at ERROR level, so drop lines below WARNING. private fun printAfterDetach( message: LogMessage, line: String? = null, ) { if (message.level.value >= Level.WARNING.value) errors.print(line ?: buildLine(message)) } ``` ```kotlin // Set by detachAndFlush, cleared by configure and reset. Before the first configure, configure writes messages. private var detached = false // configure: detached = false // detachAndFlush, in its synchronized block: writer = null; detached = true // reset: mark @Synchronized and set detached = false ``` ########## java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt: ########## @@ -189,10 +192,63 @@ data class Context( /** Registration of the executing task; resolves wired data inputs. */ internal var taskDef: TaskDef? = null + // Only from() sets this field. A context made in any other way, including with copy(), keeps it null. + private var cancellation: Cancellation? = null + + /** + * 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. + * + * Only a [StoppableTask] can call this method, because the SDK catches the stop + * signal only for that type of task. For any other task, the JVM handles the + * signal as it would without the SDK: by default the JVM runs its shutdown hooks + * and exits. Such a task could never see a stop request, so this method throws + * instead of returning false. + * + * The context that the SDK passes to the task keeps seeing the stop request + * after the task returns. A context made in any other way never sees a stop + * request, and this method returns false for it. That includes a context that + * a unit test makes to call a task directly, and a context made with `copy()`. + * + * @throws IllegalStateException if the running task is not a [StoppableTask]. + */ + fun isCancelled(): Boolean = cancellation?.check() ?: false Review Comment: A non-stoppable task that calls `isCancelled()` or `awaitCancellation()` gets an `IllegalStateException` that is reported with `shouldRetry`, so every retry fails the same way. A context from `copy()` also loses `cancellation`, so it returns `false` even for a `StoppableTask`. Could the exception fail without a retry, and could a context fall back to the running task's `Cancellation`? The KDoc about `copy()` would need updating too. ```suggestion fun isCancelled(): Boolean = (cancellation ?: Cancellation.active)?.check() ?: false ``` ```kotlin // Context.awaitCancellation val cancellation = cancellation ?: Cancellation.active // Cancellation.kt, inside the class init { current.set(this) } companion object { private val current = AtomicReference<Cancellation?>() val active: Cancellation? get() = current.get() } override fun close() { trap?.close() current.compareAndSet(this, null) } private fun throwIfTaskCannotSeeRequest() { if (!held) throw TaskCannotSeeStopRequestException("Only a StoppableTask can see a stop request. ...") } // Extends IllegalStateException, so the documented contract of isCancelled stays the same. internal class TaskCannotSeeStopRequestException(message: String) : IllegalStateException(message) // Task.kt, TaskRunner.runTask } catch (e: TaskCannotSeeStopRequestException) { logger.error("Task asked for a stop request that only a StoppableTask can see", mapOf("ti" to request.ti, "error" to e)) // A retry runs the same code, so it would fail the same way. TaskResult.failure(shouldRetry = false) } ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
