dwsmith1983 opened a new pull request, #6310: URL: https://github.com/apache/datafusion-comet/pull/6310
## Which issue does this PR close? Closes #6224. ## Rationale for this change Spark's `ExecutionMemoryPool.acquireMemory` registers the task's `memoryForTask` entry once, before its wait loop, and reads it with `memoryForTask(taskAttemptId)` on every pass. `releaseMemory` removes the entry when the task's balance reaches zero and wakes the waiters. A caller that wakes after the removal throws `NoSuchElementException: key not found: <taskAttemptId>`. The code is the same in Spark 3.4 through 4.1 and on master. `TaskMemoryManager.acquireExecutionMemory` holds the task memory manager's monitor while it is parked, but `releaseExecutionMemory` takes no monitor. So a release can empty the task's balance while an acquire of the same task is parked: under `greedy_unified`, one native thread's release fails another native thread's parked acquire; under `fair_unified`, a JVM consumer of the task freeing its last bytes with `freeMemory` does the same. The exception reaches the native side as a plain error rather than `ResourcesExhausted`, so the operator cannot spill and the task fails. ## What changes are included in this PR? `CometTaskMemoryManager.acquireMemory` calls Spark through a helper that catches this one exception, identified by its message (`key not found:` plus this task's attempt id) and a frame in `ExecutionMemoryPool`, and calls `acquireExecutionMemory` again. The retry re-registers the entry, re-checks the task's share and parks again if it has to, so the caller ends up with what Spark would have granted had the entry stayed. After three attempts it returns 0, which the native side treats as a refusal and spills on. Any other exception is rethrown. A retry logs at info level, and giving up logs a warning. The retry is safe because Spark only removes the entry at a zero balance, which cannot hold while the failed call has a partial grant. So the failed call has nothing to reconcile, and `used` and Spark's counters have not moved. The helper holds no lock of its own across attempts. Retrying from the acquire side reaches the same end state as fixing the loop in Spark (reading the entry with `getOrElseUpdate(taskAttemptId, 0L)`), without coordinating the release side, which does not know whether an acquire is parked. The Spark fix would also cover JVM consumers that call `allocatePage` directly, such as `CometUnifiedShuffleMemoryAllocator` and Spark's own operators; they are not covered here and are tracked in #6304. The contributor guide's memory management page describes the new failure and why the retry is safe. ## How are these changes tested? Four tests in `CometTaskMemoryManagerSuite` on a 100-byte off-heap `UnifiedMemoryManager`: - two threads share one `CometTaskMemoryManager`; the task holds 10 bytes and another task holds 90; an acquire of 20 parks below the minimum share; releasing the 10 bytes empties the balance. On main the parked acquire throws `key not found: 0`. With this change it completes with 20 once the other task frees its memory, and `getUsed` and the task's consumption both read 20. - the same, with the 10 bytes held by a sibling `MemoryConsumer` of the task and released with `freeMemory`. - a `NoSuchElementException` with another task's id, or without an `ExecutionMemoryPool` frame, is rethrown unchanged after one call. - when every attempt throws the matching exception, the helper returns 0 after the cap with accounting unchanged. `CometTaskMemoryManagerSuite`, `CometUnboundedShuffleMemoryAllocatorSuite` and `CometExecIteratorLifecycleSuite` pass on Spark 3.4, 3.5, 4.0 and 4.1, the spill tests in `CometExecSuite` pass on 3.5, and the native memory pool tests pass. With the retry disabled, the two race tests and the cap test fail again. -- 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]
