andygrove opened a new pull request, #5384: URL: https://github.com/apache/datafusion-comet/pull/5384
## Which issue does this PR close? Part of #5383 (which is in turn positions 17 and 18 of #5212). This is the measurement step; it does not change how much memory is acquired or when. ## Rationale for this change Every DataFusion `try_grow` and `shrink` against `CometUnifiedMemoryPool` or `CometFairMemoryPool` is a JNI round-trip into `CometTaskMemoryManager`, which then takes Spark's executor-wide `synchronized` memory manager lock. There is no batching and no hysteresis: memory is acquired in exactly the amount requested and released the moment a reservation shrinks, however small. Before adding chunked acquisition and slack retention I want a before-and-after number. There is currently no way to tell how many round-trips a query makes, so there is no way to show that the added complexity buys anything, or to spot a pool whose traffic is dominated by tiny grow/shrink churn. ## What changes are included in this PR? Both pools carried their own copy of the JNI acquire/release helpers. Those move into a single `SparkMemoryClient` (`native/core/src/execution/memory_pools/spark_client.rs`) which records, per task attempt: - acquire calls, bytes requested, bytes granted, and how many grants came back short - release calls and bytes - cumulative wall-clock time spent inside the round-trip The counters are relaxed atomics and are always collected. The summary is logged when the pool is dropped — for the task-shared unified pools, when the last native plan for the task is released: ``` Task 42 CometFairMemoryPool memory pool stats: acquire(calls=1832, requested=... bytes, granted=... bytes, short=0), release(calls=1790, bytes=...), backend_time=61.4ms ``` That is one line per task attempt, so it logs at debug level. To collect it without turning on debug logging for every module, a `log4rs.yaml` (via `COMET_CONF_DIR` or the `comet.log.file.path` system property) can raise the level for `comet::execution::memory_pools` alone; the doc comment on `log_stats` spells this out. The JNI call itself is reached through a small `SparkMemoryBackend` trait. Two things follow from that: - The pools can be constructed without a JVM, so `CometUnifiedMemoryPool` gets its first unit tests. - The `unsafe impl Send`/`Sync` moves off both pools and onto `JniMemoryBackend`, the only type that actually holds the JNI global reference. The pools now satisfy the auto traits on their own. No behavioural change — the same calls are made in the same order with the same arguments. ## How are these changes tested? Three new unit tests for the unified pool, over an in-process backend that mimics Spark's "grant as much as is available" behaviour: - `grow_and_shrink_track_spark_grants` — grow/shrink accounting stays in step with what the backend has handed out - `short_grant_is_returned_to_spark_and_reported_as_an_error` — a partial grant is released rather than retained, and the error names the requested size - `every_grow_and_shrink_reaches_spark` — 10 grow/shrink cycles produce 10 acquires and 10 releases, pinning the current no-batching behaviour that #5383 aims to reduce Plus three tests for the counters themselves in `spark_client.rs`. `cargo test -p datafusion-comet --lib` passes (163 tests), and `cargo clippy --all-targets` and `cargo fmt --check` are clean. `CometJoinSuite` passes (29 tests) as an end-to-end check of the JNI path, which every Comet suite exercises since `CometTestBase` enables off-heap mode. I deliberately left `CometFairMemoryPool` without unit tests here: any assertion about its limit would encode the `pool_size / num_consumers` bug from position 1 of #5212. Those tests belong with that fix. -- 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]
