viirya opened a new pull request, #59663:
URL: https://github.com/apache/spark/pull/59663
### What changes were proposed in this pull request?
When another thread is invoking task listeners at the moment the task
completes (e.g. a kill thread running failure listeners because an interrupt
listener threw, or a non-task thread that called `markTaskFailed` through
`Utils.tryWithSafeFinallyAndFailureCallbacks`),
`TaskContextImpl.markTaskCompleted` returns without running the completion
listeners, and the other thread only drains its own stack, so the completion
listeners never run. Failure listeners queued while the completing thread runs
completion listeners, and completion listeners added after completion while
another thread holds the invocation, are dropped the same way.
This PR keeps the completion listeners on the completing thread:
- `markTaskCompleted` claims the listener invocation in the same
synchronized block that sets `completed`. If another thread holds it,
`markTaskCompleted` registers itself as the next invoker and waits; the holder
hands the invocation over directly when it is done, so no third thread can take
it in between.
- When a listener throws, the failure listeners run nested in the current
invocation instead of clearing and reclaiming `listenerInvocationThread`, which
left a window in which another thread could take over the listeners.
- Failure listeners left to the invoking thread run before its next
listener, so failure cleanup (e.g. a state store `abort()`) runs before the
completion listener that releases the store, also after a handover.
- Before releasing the invocation, a thread runs the completion listeners
other threads left to it, unless a thread waiting in `markTaskCompleted` takes
them over. The release and the handover are in a `finally`.
- The handover wait keeps interrupts (they are restored once it is done) and
logs a warning periodically; the first one carries the holder's stack trace,
taken outside the TaskContext lock. The interval is settable for tests.
- The StreamingShuffleSuite test that documented the skipped completion
listener now checks that it runs on the completing thread, and the comments in
`ErrorNotifier` and `StreamingShuffleReader` are updated accordingly.
This was split out of #59332 during its review, since it changes listener
behavior for every task.
### Why are the changes needed?
Completion listeners do the task's cleanup (closing readers, releasing state
stores, freeing task-managed resources). When they are skipped, those resources
leak; when they run late on another thread, they run after the executor has
already released the task's memory. Failure listeners that are dropped skip
failure cleanup such as aborting a state store update.
Note on the wait: it has no time bound, so if the holder's listeners block,
the completing task waits too. On the kill path this is not new:
`TaskRunner.kill` runs `Task.kill` (and so the failure listeners) holding the
`TaskRunner` lock, which the task thread needs anyway in
`setTaskFinishedAndClearInterruptStatus` before the task can finish. The new
wait only applies when a non-kill thread holds the invocation;
`spark.task.reaper.enabled` with a positive `spark.task.reaper.killTimeout`
bounds it in that case, and the periodic warning names the holder and shows its
stack.
Out of scope: post-status-update listeners can be dropped the same way if
another thread holds the invocation when `invokePostStatusUpdateListeners`
runs; this PR does not change that.
### Does this PR introduce _any_ user-facing change?
No API change. In the cases above, completion and failure listeners that
were skipped now run, and `markTaskCompleted` waits for the other thread to
finish its listeners first.
### How was this patch tested?
New tests in `TaskContextSuite`:
- a listener added from another thread as completion starts does not take
over the listeners;
- a failing listener does not let another thread take over the listeners;
- completion waits for another thread invoking listeners (a kill whose
interrupt listener throws) and runs the listeners on the completing thread;
- a listener on the invoking thread that completes the task runs the
completion listeners nested;
- a failure listener left to the completing thread runs before its next
completion listener, without and with a handover;
- completion listeners left to a failure invoker run before it releases, a
failure of such a listener reaches the invoker's caller, and the invocation is
still released;
- an interrupt does not stop the handover wait and is kept;
- the handover wait warns with the holder's name and, on the first warning
only, its stack trace.
Updated test in `StreamingShuffleSuite`: "background markTaskFailed races
with task thread markTaskCompleted".
Each behavior was checked by reverting it alone and confirming that a test
fails. All core and sql suites that call `markTaskFailed`, `markTaskCompleted`
or `markInterrupted` pass locally (`TaskContextSuite`, `BasePythonRunnerSuite`,
`StreamingShuffle{,Reader,Writer}Suite`, `PartiallySerializedBlockSuite`,
`ShuffleBlockFetcherIteratorSuite`, `SupportsArchiveFormatSuite`,
`JdbcTaskInterruptSuite`, `JsonV1Suite`, `JsonV2Suite`,
`ParquetStorageFilterSuite`, `RocksDB*Suite`, `ValueStateSuite`), as well as
`ExecutorSuite` and `JobCancellationSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
