Describe the bug
Follow-up to #6224, which covers native acquires through CometTaskMemoryManager. JVM consumers of the same task have the same exposure and are not covered there.
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. A caller parked in the loop that wakes after the removal gets NoSuchElementException: key not found: <taskAttemptId>. TaskMemoryManager.releaseExecutionMemory takes no monitor, so a release can run while an acquire of the same task is parked.
CometUnifiedShuffleMemoryAllocator.allocate calls allocatePage, and Spark's own operators in the task (sorters, aggregates) do the same. When one of them is parked and a native CometTaskMemoryManager.releaseMemory empties the task's balance, the parked allocatePage fails with that exception instead of completing or returning a page of zero size. Any off-heap consumer's freeMemory can do the same to another parked consumer of the task.
Steps to reproduce
Component-level, in the style of CometTaskMemoryManagerSuite: a 100-byte off-heap UnifiedMemoryManager, one TaskMemoryManager. A CometTaskMemoryManager acquires 10 bytes for the task and another task takes 90. On a second thread, CometUnifiedShuffleMemoryAllocator.allocate(20) parks below the task's minimum share. The main thread calls CometTaskMemoryManager.releaseMemory(10). The parked allocate throws NoSuchElementException: key not found: <taskAttemptId>.
Expected behavior
A parked page allocation completes once memory is free, or fails with Spark's usual out-of-memory path. It does not fail because another consumer of the same task released memory.
Additional context
The root cause is in Spark: reading the entry with getOrElseUpdate(taskAttemptId, 0L) inside the wait loop would remove the window for every consumer. Until that lands, Comet can only guard its own callers. CometTaskMemoryManager retries the acquire (#6224); the allocator could do the same in allocate, and Spark's own operators cannot be guarded from Comet.
Describe the bug
Follow-up to #6224, which covers native acquires through
CometTaskMemoryManager. JVM consumers of the same task have the same exposure and are not covered there.Spark's
ExecutionMemoryPool.acquireMemoryregisters the task'smemoryForTaskentry once, before its wait loop, and reads it withmemoryForTask(taskAttemptId)on every pass.releaseMemoryremoves the entry when the task's balance reaches zero. A caller parked in the loop that wakes after the removal getsNoSuchElementException: key not found: <taskAttemptId>.TaskMemoryManager.releaseExecutionMemorytakes no monitor, so a release can run while an acquire of the same task is parked.CometUnifiedShuffleMemoryAllocator.allocatecallsallocatePage, and Spark's own operators in the task (sorters, aggregates) do the same. When one of them is parked and a nativeCometTaskMemoryManager.releaseMemoryempties the task's balance, the parkedallocatePagefails with that exception instead of completing or returning a page of zero size. Any off-heap consumer'sfreeMemorycan do the same to another parked consumer of the task.Steps to reproduce
Component-level, in the style of
CometTaskMemoryManagerSuite: a 100-byte off-heapUnifiedMemoryManager, oneTaskMemoryManager. ACometTaskMemoryManageracquires 10 bytes for the task and another task takes 90. On a second thread,CometUnifiedShuffleMemoryAllocator.allocate(20)parks below the task's minimum share. The main thread callsCometTaskMemoryManager.releaseMemory(10). The parked allocate throwsNoSuchElementException: key not found: <taskAttemptId>.Expected behavior
A parked page allocation completes once memory is free, or fails with Spark's usual out-of-memory path. It does not fail because another consumer of the same task released memory.
Additional context
The root cause is in Spark: reading the entry with
getOrElseUpdate(taskAttemptId, 0L)inside the wait loop would remove the window for every consumer. Until that lands, Comet can only guard its own callers.CometTaskMemoryManagerretries the acquire (#6224); the allocator could do the same inallocate, and Spark's own operators cannot be guarded from Comet.