LuciferYang opened a new pull request, #12425: URL: https://github.com/apache/gluten/pull/12425
### What changes were proposed in this pull request? `DynamicOffHeapSizingMemoryTarget.borrow(size)` used to unconditionally book the requested `size` into `USED_OFF_HEAP_BYTES` and its per-instance recorder, call `target.borrow(size)` while discarding the wrapped target's return value, and return `size` to the caller. The wrapped target — typically a `TreeMemoryConsumer.Node` bottoming out in Spark's `MemoryConsumer.acquireMemory` — is allowed to grant less than `size` (Spark's execution memory pool routinely returns partial reservations under pressure), so the counters drifted above the true reservation. `repay` had the mirrored bug: `TreeMemoryConsumer.Node.repay` bounds the freed amount via `Math.min(usedBytes(), size)` and returns what it actually released, but this class decremented by the full requested `size`. There was also a permanent leak on the throw path — the pre-fix code incremented the counters before calling into the wrapped target, so any exception from Spark's memory manager (or a `Preconditions.checkState` inside `ensureFreeCapacity`, or a task-teardown race) permanently added the requested bytes to the static counter for the lifetime of the executor JVM. This PR rewrites `borrow` to capture `granted = target.borrow(size)` and record only `granted`, and `repay` to capture `freed = target.repay(size)` and record only `freed`. Ledger updates now happen strictly after the wrapped call returns normally, so exceptions leave both counters untouched. A `repay(0)` short-circuit is added for symmetry with the existing `borrow(0)` guard. Two package-private test hooks (`usedOffHeapBytesForTesting` / `resetUsedOffHeapBytesForTesting`, both `@VisibleForTesting`) expose the static counter to the new JUnit 4 tests. ### Why are the changes needed? Closes #12424. The static `USED_OFF_HEAP_BYTES` counter feeds this class's own `exceedsMaxMemoryUsage()` gate, so overshoot there causes subsequent in-budget borrows to be rejected across the whole executor. Under-frees on `repay` leave phantom bytes that never come back, causing the same rejection over time. The two symptoms compound: heavy off-heap workloads under memory pressure see spurious OOMs even when the actual Spark pool has room, and the discrepancy grows with each partial reservation. The throw-path leak makes the drift permanent rather than transient. ### Does this PR introduce any user-facing change? No functional change to the API surface. Callers now see the actual grant / free amount instead of a false `size` echo — which is what the `MemoryTarget` contract already promised. `ThrowOnOomMemoryTarget`'s existing partial-grant retry logic starts working correctly (it was previously unreachable through this wrapper). ### How was this patch tested? `mvn -pl gluten-core -Pspark-3.5 spotless:check` passes. The new `DynamicOffHeapSizingMemoryTargetTest` (9 tests) covers full-grant, partial-grant, zero-grant, sequential accumulation, partial-repay drain, `borrow(0)` / `repay(0)` short-circuits, and throw-safety on both `borrow` and `repay`. Each test was verified against the pre-fix implementation: reverting the borrow/repay bodies to the original shape red-fails exactly the three ledger-drift assertions the fix protects, plus `@After`'s static-counter-must-be-zero invariant. Every borrow in the suite pairs with a repay so the shared static counter returns to zero across test methods. -- 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]
