andygrove commented on code in PR #6128: URL: https://github.com/apache/datafusion-comet/pull/6128#discussion_r4085117247
########## native/core/src/execution/memory_pools/spark_memory.rs: ########## @@ -0,0 +1,331 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::{ + atomic::{AtomicUsize, Ordering::Relaxed}, + Arc, +}; + +use jni::objects::{Global, JObject}; +use log::warn; + +use crate::{errors::CometResult, jvm_bridge::JVMClasses}; + +/// Spark's side of a Comet pool: the calls that acquire and release off-heap execution memory. +pub(super) trait SparkMemoryManager { + /// Asks Spark for `size` bytes and returns how many it granted. + fn acquire(&self, size: usize) -> CometResult<i64>; + fn release(&self, size: usize) -> CometResult<()>; +} + +/// Calls [`crate::jvm_bridge::CometTaskMemoryManager`] over JNI. +pub(super) struct JniMemoryManager(Arc<Global<JObject<'static>>>); + +impl SparkMemoryManager for JniMemoryManager { + fn acquire(&self, size: usize) -> CometResult<i64> { + let handle = self.0.as_obj(); + JVMClasses::with_env(|env| unsafe { + jni_call!(env, + comet_task_memory_manager(handle).acquire_memory(size as i64) -> i64) + }) + } + + fn release(&self, size: usize) -> CometResult<()> { + let handle = self.0.as_obj(); + JVMClasses::with_env(|env| unsafe { + jni_call!(env, comet_task_memory_manager(handle).release_memory(size as i64) -> ()) + }) + } +} + +/// Memory a Comet pool holds from Spark, including any it has recorded without Spark's grant. +/// +/// `MemoryPool::grow` must always succeed: DataFusion calls it for memory that already exists, +/// such as a spilled batch read back from disk. When Spark grants less than [`Self::acquire`] +/// asked for, the shortfall is carried as overcommit rather than failing. [`Self::release`] repays +/// it before returning anything to Spark, so Spark is never handed back more than it granted. +/// While any is outstanding, [`Self::try_acquire`] also asks Spark for it, so a pool refuses +/// `try_grow` until Spark can cover both the request and the debt, and the operator spills. +/// +/// # Invariant +/// +/// Every byte a pool records through this type is backed either by Spark's grant or by +/// overcommit, so Spark's grant plus `overcommit` equals the bytes recorded and not yet released. +/// Spark is handed back more than it granted only if a release takes less from `overcommit` than +/// it could. `CometUnifiedMemoryPool` calls in here from several threads without a lock, and +/// updates its own `used` separately, so this rests on three things: +/// +/// - `overcommit` only grows by bytes that are being recorded in the same call. +/// - Each repayment takes its share of `overcommit` in a single atomic update, so two concurrent +/// calls can never repay the same debt. +/// - A caller never releases more than it recorded. DataFusion guarantees this, because a shrink +/// cannot exceed the reservation it comes from. +pub(super) struct SparkMemory { + manager: Box<dyn SparkMemoryManager>, + overcommit: AtomicUsize, + task_attempt_id: i64, +} + +/// Why [`SparkMemory::try_acquire`] refused a request. +#[derive(Debug, PartialEq, Eq)] +pub(super) struct Refusal { + /// Outstanding overcommit that was asked for on top of the request. + pub(super) overcommit: usize, + /// What Spark offered before it was handed back. + pub(super) granted: usize, +} + +impl SparkMemory { + pub(super) fn new(handle: Arc<Global<JObject<'static>>>, task_attempt_id: i64) -> Self { + Self::with_manager(Box::new(JniMemoryManager(handle)), task_attempt_id) + } + + pub(super) fn with_manager(manager: Box<dyn SparkMemoryManager>, task_attempt_id: i64) -> Self { + Self { + manager, + overcommit: AtomicUsize::new(0), + task_attempt_id, + } + } + + pub(super) fn task_attempt_id(&self) -> i64 { + self.task_attempt_id + } + + /// Acquires `size` bytes plus any outstanding overcommit, or nothing. A full grant repays the + /// overcommit; a partial one is handed back and reported as a [`Refusal`]. + pub(super) fn try_acquire(&self, size: usize) -> CometResult<Result<(), Refusal>> { + let debt = self.overcommit.load(Relaxed); + let request = size.saturating_add(debt); + let granted = granted(request, self.manager.acquire(request)?); + if granted < request { + if granted > 0 { + self.manager.release(granted)?; + } + return Ok(Err(Refusal { + overcommit: debt, + granted, + })); + } + if debt > 0 { + // A concurrent release may have repaid part of the debt since it was read, in which + // case Spark granted more than is still owed and the excess goes back. + let owed = self.repay(debt); + if owed < debt { + self.manager.release(debt - owed)?; + } + } Review Comment: Added both. `try_acquire_hands_back_debt_repaid_by_a_concurrent_release` gives `FakeSpark` a hook that runs inside its next `acquire`, and uses it to release 30 of the 60 bytes owed on the same `SparkMemory` while `try_acquire` is waiting on Spark. It checks that only those 30 go back and that Spark ends up holding exactly what is still recorded. `concurrent_consumers_hand_spark_back_exactly_what_it_granted` runs 8 consumers doing random `grow`, `try_grow` and `shrink` on `CometUnifiedMemoryPool` against a `FakeSpark` whose limit keeps moving, then checks that `reserved()`, `overcommit()` and `fake.held()` are all 0. With a temporary counter in that branch it went through it about 400 times per run. Both tests fail with the excess release removed or with the whole debt handed back, and the stress test also fails with a `repay` that loads and stores separately. It passed 1,700 runs locally, 1,200 of them with eight copies running at once. `SparkMemoryManager` is now `Send + Sync` so the fake can be shared across threads, and the fake uses `parking_lot` so a failed assertion in one thread doesn't poison the lock for all the others. -- 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]
