diff --git a/native/core/src/execution/memory_pools/fair_pool.rs b/native/core/src/execution/memory_pools/fair_pool.rs index 347c3d8ef69..50bcd2b864e 100644 --- a/native/core/src/execution/memory_pools/fair_pool.rs +++ b/native/core/src/execution/memory_pools/fair_pool.rs @@ -15,14 +15,9 @@ // specific language governing permissions and limitations // under the License. -use std::{ - fmt::{Debug, Display, Formatter, Result as FmtResult}, - sync::Arc, -}; - -use jni::objects::{Global, JObject}; +use std::fmt::{Debug, Display, Formatter, Result as FmtResult}; -use crate::{errors::CometResult, jvm_bridge::JVMClasses}; +use crate::execution::memory_pools::spark_client::SparkMemoryClient; use datafusion::common::resources_err; use datafusion::execution::memory_pool::MemoryConsumer; use datafusion::{ @@ -34,7 +29,7 @@ use parking_lot::Mutex; /// A DataFusion fair `MemoryPool` implementation for Comet. Internally this is /// implemented via delegating calls to [`crate::jvm_bridge::CometTaskMemoryManager`]. pub struct CometFairMemoryPool { - task_memory_manager_handle: Arc>>, + client: SparkMemoryClient, pool_size: usize, state: Mutex, } @@ -56,30 +51,18 @@ impl Debug for CometFairMemoryPool { } impl CometFairMemoryPool { - pub fn new( - task_memory_manager_handle: Arc>>, - pool_size: usize, - ) -> CometFairMemoryPool { + pub fn new(client: SparkMemoryClient, pool_size: usize) -> CometFairMemoryPool { Self { - task_memory_manager_handle, + client, pool_size, state: Mutex::new(CometFairPoolState { used: 0, num: 0 }), } } +} - fn acquire(&self, additional: usize) -> CometResult { - let handle = self.task_memory_manager_handle.as_obj(); - JVMClasses::with_env(|env| unsafe { - jni_call!(env, - comet_task_memory_manager(handle).acquire_memory(additional as i64) -> i64) - }) - } - - fn release(&self, size: usize) -> CometResult<()> { - let handle = self.task_memory_manager_handle.as_obj(); - JVMClasses::with_env(|env| unsafe { - jni_call!(env, comet_task_memory_manager(handle).release_memory(size as i64) -> ()) - }) +impl Drop for CometFairMemoryPool { + fn drop(&mut self) { + self.client.log_stats(self.name()); } } @@ -94,9 +77,6 @@ impl Display for CometFairMemoryPool { } } -unsafe impl Send for CometFairMemoryPool {} -unsafe impl Sync for CometFairMemoryPool {} - impl MemoryPool for CometFairMemoryPool { fn name(&self) -> &str { "CometFairMemoryPool" @@ -134,7 +114,8 @@ impl MemoryPool for CometFairMemoryPool { state.used ) } - self.release(subtractive) + self.client + .release(subtractive) .unwrap_or_else(|_| panic!("Failed to release {subtractive} bytes")); state.used = state.used.checked_sub(subtractive).unwrap(); } @@ -162,12 +143,12 @@ impl MemoryPool for CometFairMemoryPool { ); } - let acquired = self.acquire(additional)?; + let acquired = self.client.acquire(additional)?; // If the number of bytes we acquired is less than the requested, return an error, // and hopefully will trigger spilling from the caller side. if acquired < additional as i64 { // Release the acquired bytes before throwing error - self.release(acquired as usize)?; + self.client.release(acquired as usize)?; return resources_err!( "Failed to acquire {} bytes, only got {} bytes. Reserved: {} bytes", diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index 389e3489907..913f2bb697c 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -18,6 +18,7 @@ mod config; mod fair_pool; pub mod logging_pool; +mod spark_client; mod task_shared; mod unified_pool; @@ -27,6 +28,7 @@ use datafusion::execution::memory_pool::{ use fair_pool::CometFairMemoryPool; use jni::objects::{Global, JObject}; use once_cell::sync::OnceCell; +use spark_client::SparkMemoryClient; use std::num::NonZeroUsize; use std::sync::Arc; use unified_pool::CometUnifiedMemoryPool; @@ -46,10 +48,10 @@ pub(crate) fn create_memory_pool( let per_task_memory_pool = memory_pool_map.entry(task_attempt_id).or_insert_with(|| { let pool: Arc = Arc::new(TrackConsumersPool::new( - CometUnifiedMemoryPool::new( + CometUnifiedMemoryPool::new(SparkMemoryClient::new( Arc::clone(&comet_task_memory_manager), task_attempt_id, - ), + )), NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(), )); PerTaskMemoryPool::new(pool) @@ -63,7 +65,10 @@ pub(crate) fn create_memory_pool( memory_pool_map.entry(task_attempt_id).or_insert_with(|| { let pool: Arc = Arc::new(TrackConsumersPool::new( CometFairMemoryPool::new( - Arc::clone(&comet_task_memory_manager), + SparkMemoryClient::new( + Arc::clone(&comet_task_memory_manager), + task_attempt_id, + ), memory_pool_config.pool_size, ), NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(), diff --git a/native/core/src/execution/memory_pools/spark_client.rs b/native/core/src/execution/memory_pools/spark_client.rs new file mode 100644 index 00000000000..0d1a0d95162 --- /dev/null +++ b/native/core/src/execution/memory_pools/spark_client.rs @@ -0,0 +1,347 @@ +// 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::{ + fmt::{Debug, Display, Formatter, Result as FmtResult}, + sync::{ + atomic::{AtomicU64, Ordering::Relaxed}, + Arc, + }, + time::{Duration, Instant}, +}; + +use jni::objects::{Global, JObject}; + +use crate::{errors::CometResult, jvm_bridge::JVMClasses}; + +/// The source of memory for the unified memory pools. In production this is +/// Spark's off-heap executor memory pool, reached over JNI; tests substitute an +/// in-process implementation so that the pools can be exercised without a JVM. +pub(crate) trait SparkMemoryBackend: Send + Sync { + /// Request `size` bytes, returning the number of bytes actually granted, + /// which may be less than requested. + fn acquire(&self, size: usize) -> CometResult; + + /// Return `size` bytes. + fn release(&self, size: usize) -> CometResult<()>; +} + +/// A [`SparkMemoryBackend`] that delegates to Spark's unified memory manager by +/// calling [`crate::jvm_bridge::CometTaskMemoryManager`] over JNI. +struct JniMemoryBackend { + task_memory_manager_handle: Arc>>, +} + +// The JNI global reference is safe to use from any thread that attaches to the +// JVM, which `JVMClasses::with_env` guarantees. +unsafe impl Send for JniMemoryBackend {} +unsafe impl Sync for JniMemoryBackend {} + +impl SparkMemoryBackend for JniMemoryBackend { + fn acquire(&self, size: usize) -> CometResult { + let handle = self.task_memory_manager_handle.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.task_memory_manager_handle.as_obj(); + JVMClasses::with_env(|env| unsafe { + jni_call!(env, comet_task_memory_manager(handle).release_memory(size as i64) -> ()) + }) + } +} + +/// Counters describing the traffic a memory pool has sent to its backend. +/// +/// Each `try_grow` or `shrink` that reaches the backend costs a JNI round-trip +/// plus contention on Spark's executor-wide memory manager lock, so these +/// counters are the input to any decision about batching acquisitions or +/// releasing with hysteresis (see issue #5383). They are plain relaxed atomics +/// and are always collected. +#[derive(Default)] +pub(crate) struct MemoryPoolStats { + acquire_calls: AtomicU64, + acquire_requested_bytes: AtomicU64, + acquire_granted_bytes: AtomicU64, + /// Acquisitions that were granted less than they asked for. + short_grants: AtomicU64, + release_calls: AtomicU64, + release_bytes: AtomicU64, + /// Wall-clock time spent inside the backend, i.e. in the JNI call and the + /// Spark-side accounting it performs. + backend_nanos: AtomicU64, +} + +impl MemoryPoolStats { + fn record_acquire(&self, requested: usize, granted: i64, elapsed: Duration) { + self.acquire_calls.fetch_add(1, Relaxed); + self.acquire_requested_bytes + .fetch_add(requested as u64, Relaxed); + if granted > 0 { + self.acquire_granted_bytes + .fetch_add(granted as u64, Relaxed); + } + if granted < requested as i64 { + self.short_grants.fetch_add(1, Relaxed); + } + self.record_elapsed(elapsed); + } + + fn record_release(&self, size: usize, elapsed: Duration) { + self.release_calls.fetch_add(1, Relaxed); + self.release_bytes.fetch_add(size as u64, Relaxed); + self.record_elapsed(elapsed); + } + + fn record_elapsed(&self, elapsed: Duration) { + self.backend_nanos + .fetch_add(elapsed.as_nanos() as u64, Relaxed); + } + + pub(crate) fn acquire_calls(&self) -> u64 { + self.acquire_calls.load(Relaxed) + } + + pub(crate) fn release_calls(&self) -> u64 { + self.release_calls.load(Relaxed) + } + + /// True if no call ever reached the backend, in which case there is nothing + /// worth logging. + fn is_empty(&self) -> bool { + self.acquire_calls() == 0 && self.release_calls() == 0 + } +} + +impl Debug for MemoryPoolStats { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + f.debug_struct("MemoryPoolStats") + .field("acquire_calls", &self.acquire_calls()) + .field( + "acquire_requested_bytes", + &self.acquire_requested_bytes.load(Relaxed), + ) + .field( + "acquire_granted_bytes", + &self.acquire_granted_bytes.load(Relaxed), + ) + .field("short_grants", &self.short_grants.load(Relaxed)) + .field("release_calls", &self.release_calls()) + .field("release_bytes", &self.release_bytes.load(Relaxed)) + .field("backend_nanos", &self.backend_nanos.load(Relaxed)) + .finish() + } +} + +impl Display for MemoryPoolStats { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + write!( + f, + "acquire(calls={}, requested={} bytes, granted={} bytes, short={}), \ + release(calls={}, bytes={}), backend_time={:?}", + self.acquire_calls(), + self.acquire_requested_bytes.load(Relaxed), + self.acquire_granted_bytes.load(Relaxed), + self.short_grants.load(Relaxed), + self.release_calls(), + self.release_bytes.load(Relaxed), + Duration::from_nanos(self.backend_nanos.load(Relaxed)), + ) + } +} + +/// Handle used by the unified memory pools to acquire and release memory from +/// Spark, recording [`MemoryPoolStats`] for every call that it makes. +pub(crate) struct SparkMemoryClient { + backend: Arc, + task_attempt_id: i64, + stats: MemoryPoolStats, +} + +impl SparkMemoryClient { + /// Create a client that acquires memory from Spark over JNI. + pub(crate) fn new( + task_memory_manager_handle: Arc>>, + task_attempt_id: i64, + ) -> Self { + Self::with_backend( + Arc::new(JniMemoryBackend { + task_memory_manager_handle, + }), + task_attempt_id, + ) + } + + pub(crate) fn with_backend(backend: Arc, task_attempt_id: i64) -> Self { + Self { + backend, + task_attempt_id, + stats: MemoryPoolStats::default(), + } + } + + pub(crate) fn task_attempt_id(&self) -> i64 { + self.task_attempt_id + } + + #[cfg(test)] + pub(crate) fn stats(&self) -> &MemoryPoolStats { + &self.stats + } + + /// Request `size` bytes from Spark, returning the number of bytes granted. + pub(crate) fn acquire(&self, size: usize) -> CometResult { + let start = Instant::now(); + let result = self.backend.acquire(size); + let elapsed = start.elapsed(); + match &result { + Ok(granted) => self.stats.record_acquire(size, *granted, elapsed), + // A failed call still costs a round-trip, so account for its time. + Err(_) => self.stats.record_acquire(size, 0, elapsed), + } + result + } + + /// Return `size` bytes to Spark. + pub(crate) fn release(&self, size: usize) -> CometResult<()> { + let start = Instant::now(); + let result = self.backend.release(size); + self.stats.record_release(size, start.elapsed()); + result + } + + /// Emit the accumulated statistics for this task. Called when a pool is + /// dropped, which for the task-shared unified pools is when the last native + /// plan for the task is released. + /// + /// This logs at debug level because it produces one line per task attempt. + /// To collect it without turning on debug logging for every module, supply a + /// `log4rs.yaml` (via `COMET_CONF_DIR` or the `comet.log.file.path` system + /// property) that raises the level for this module alone: + /// + /// ```yaml + /// loggers: + /// comet::execution::memory_pools: + /// level: debug + /// ``` + pub(crate) fn log_stats(&self, pool_name: &str) { + if !self.stats.is_empty() { + log::debug!( + "Task {} {pool_name} memory pool stats: {}", + self.task_attempt_id, + self.stats + ); + } + } +} + +impl Debug for SparkMemoryClient { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + f.debug_struct("SparkMemoryClient") + .field("task_attempt_id", &self.task_attempt_id) + .field("stats", &self.stats) + .finish() + } +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + + /// An in-process [`SparkMemoryBackend`] standing in for Spark's memory + /// manager, so that the pools can be tested without a JVM. + pub(crate) struct TestMemoryBackend { + /// Total memory the backend is willing to hand out. + capacity: usize, + granted: std::sync::Mutex, + } + + impl TestMemoryBackend { + pub(crate) fn new(capacity: usize) -> Self { + Self { + capacity, + granted: std::sync::Mutex::new(0), + } + } + + pub(crate) fn granted(&self) -> usize { + *self.granted.lock().unwrap() + } + } + + impl SparkMemoryBackend for TestMemoryBackend { + fn acquire(&self, size: usize) -> CometResult { + // Mirror Spark's behaviour of granting as much as is available. + let mut granted = self.granted.lock().unwrap(); + let available = self.capacity - *granted; + let grant = size.min(available); + *granted += grant; + Ok(grant as i64) + } + + fn release(&self, size: usize) -> CometResult<()> { + let mut granted = self.granted.lock().unwrap(); + *granted -= size; + Ok(()) + } + } + + fn client(capacity: usize) -> SparkMemoryClient { + SparkMemoryClient::with_backend(Arc::new(TestMemoryBackend::new(capacity)), 1) + } + + #[test] + fn stats_count_acquires_and_releases() { + let client = client(1024); + + assert_eq!(client.acquire(100).unwrap(), 100); + assert_eq!(client.acquire(200).unwrap(), 200); + client.release(100).unwrap(); + + let stats = client.stats(); + assert_eq!(stats.acquire_calls(), 2); + assert_eq!(stats.acquire_requested_bytes.load(Relaxed), 300); + assert_eq!(stats.acquire_granted_bytes.load(Relaxed), 300); + assert_eq!(stats.short_grants.load(Relaxed), 0); + assert_eq!(stats.release_calls(), 1); + assert_eq!(stats.release_bytes.load(Relaxed), 100); + } + + #[test] + fn stats_count_short_grants() { + let client = client(100); + + assert_eq!(client.acquire(150).unwrap(), 100); + + let stats = client.stats(); + assert_eq!(stats.acquire_calls(), 1); + assert_eq!(stats.acquire_requested_bytes.load(Relaxed), 150); + assert_eq!(stats.acquire_granted_bytes.load(Relaxed), 100); + assert_eq!(stats.short_grants.load(Relaxed), 1); + } + + #[test] + fn stats_are_empty_before_any_call() { + let client = client(1024); + assert!(client.stats().is_empty()); + client.acquire(1).unwrap(); + assert!(!client.stats().is_empty()); + } +} diff --git a/native/core/src/execution/memory_pools/unified_pool.rs b/native/core/src/execution/memory_pools/unified_pool.rs index e72c734d04d..562c38b07d7 100644 --- a/native/core/src/execution/memory_pools/unified_pool.rs +++ b/native/core/src/execution/memory_pools/unified_pool.rs @@ -17,27 +17,22 @@ use std::{ fmt::{Debug, Display, Formatter, Result as FmtResult}, - sync::{ - atomic::{AtomicUsize, Ordering::Relaxed}, - Arc, - }, + sync::atomic::{AtomicUsize, Ordering::Relaxed}, }; -use crate::{errors::CometResult, jvm_bridge::JVMClasses}; +use crate::execution::memory_pools::spark_client::SparkMemoryClient; use datafusion::{ common::{resources_datafusion_err, DataFusionError}, execution::memory_pool::{MemoryPool, MemoryReservation}, }; -use jni::objects::{Global, JObject}; use log::warn; /// A DataFusion `MemoryPool` implementation for Comet that delegates to /// Spark's off-heap executor memory pool via JNI by calling /// [`crate::jvm_bridge::CometTaskMemoryManager`]. pub struct CometUnifiedMemoryPool { - task_memory_manager_handle: Arc>>, + client: SparkMemoryClient, used: AtomicUsize, - task_attempt_id: i64, } impl Debug for CometUnifiedMemoryPool { @@ -49,42 +44,26 @@ impl Debug for CometUnifiedMemoryPool { } impl CometUnifiedMemoryPool { - pub fn new( - task_memory_manager_handle: Arc>>, - task_attempt_id: i64, - ) -> CometUnifiedMemoryPool { + pub fn new(client: SparkMemoryClient) -> CometUnifiedMemoryPool { Self { - task_memory_manager_handle, - task_attempt_id, + client, used: AtomicUsize::new(0), } } - /// Request memory from Spark's off-heap memory pool via JNI - fn acquire_from_spark(&self, additional: usize) -> CometResult { - let handle = self.task_memory_manager_handle.as_obj(); - JVMClasses::with_env(|env| unsafe { - jni_call!(env, - comet_task_memory_manager(handle).acquire_memory(additional as i64) -> i64) - }) - } - - /// Release memory to Spark's off-heap memory pool via JNI - fn release_to_spark(&self, size: usize) -> CometResult<()> { - let handle = self.task_memory_manager_handle.as_obj(); - JVMClasses::with_env(|env| unsafe { - jni_call!(env, comet_task_memory_manager(handle).release_memory(size as i64) -> ()) - }) + fn task_attempt_id(&self) -> i64 { + self.client.task_attempt_id() } } impl Drop for CometUnifiedMemoryPool { fn drop(&mut self) { + self.client.log_stats(self.name()); let used = self.used.load(Relaxed); if used != 0 { warn!( "Task {} dropped CometUnifiedMemoryPool with {used} bytes still reserved", - self.task_attempt_id + self.task_attempt_id() ); } } @@ -100,9 +79,6 @@ impl Display for CometUnifiedMemoryPool { } } -unsafe impl Send for CometUnifiedMemoryPool {} -unsafe impl Sync for CometUnifiedMemoryPool {} - impl MemoryPool for CometUnifiedMemoryPool { fn name(&self) -> &str { "CometUnifiedMemoryPool" @@ -113,10 +89,10 @@ impl MemoryPool for CometUnifiedMemoryPool { } fn shrink(&self, _: &MemoryReservation, size: usize) { - if let Err(e) = self.release_to_spark(size) { + if let Err(e) = self.client.release(size) { panic!( "Task {} failed to return {size} bytes to Spark: {e:?}", - self.task_attempt_id + self.task_attempt_id() ); } if let Err(prev) = self @@ -125,23 +101,23 @@ impl MemoryPool for CometUnifiedMemoryPool { { panic!( "Task {} overflow when releasing {size} of {prev} bytes", - self.task_attempt_id + self.task_attempt_id() ); } } fn try_grow(&self, _: &MemoryReservation, additional: usize) -> Result<(), DataFusionError> { if additional > 0 { - let acquired = self.acquire_from_spark(additional)?; + let acquired = self.client.acquire(additional)?; // If the number of bytes we acquired is less than the requested, return an error, // and hopefully will trigger spilling from the caller side. if acquired < additional as i64 { // Release the acquired bytes before throwing error - self.release_to_spark(acquired as usize)?; + self.client.release(acquired as usize)?; return Err(resources_datafusion_err!( "Task {} failed to acquire {} bytes, only got {}. Reserved: {}", - self.task_attempt_id, + self.task_attempt_id(), additional, acquired, self.reserved() @@ -153,7 +129,7 @@ impl MemoryPool for CometUnifiedMemoryPool { { return Err(resources_datafusion_err!( "Task {} failed to acquire {} bytes due to overflow. Reserved: {}", - self.task_attempt_id, + self.task_attempt_id(), additional, prev )); @@ -166,3 +142,81 @@ impl MemoryPool for CometUnifiedMemoryPool { self.used.load(Relaxed) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution::memory_pools::spark_client::tests::TestMemoryBackend; + use datafusion::execution::memory_pool::MemoryConsumer; + use std::sync::Arc; + + /// Returns the pool both as a trait object, for registering consumers, and + /// as the concrete type, for inspecting the statistics it recorded. + fn pool( + capacity: usize, + ) -> ( + Arc, + Arc, + Arc, + ) { + let backend = Arc::new(TestMemoryBackend::new(capacity)); + let client = SparkMemoryClient::with_backend(Arc::clone(&backend) as _, 1); + let pool = Arc::new(CometUnifiedMemoryPool::new(client)); + (Arc::clone(&pool) as _, pool, backend) + } + + #[test] + fn grow_and_shrink_track_spark_grants() { + let (pool, unified, backend) = pool(1024); + let reservation = MemoryConsumer::new("test").register(&pool); + + reservation.try_grow(100).unwrap(); + assert_eq!(pool.reserved(), 100); + assert_eq!(backend.granted(), 100); + + reservation.try_grow(200).unwrap(); + assert_eq!(pool.reserved(), 300); + assert_eq!(backend.granted(), 300); + + reservation.shrink(300); + assert_eq!(pool.reserved(), 0); + assert_eq!(backend.granted(), 0); + + let stats = unified.client.stats(); + assert_eq!(stats.acquire_calls(), 2); + assert_eq!(stats.release_calls(), 1); + } + + #[test] + fn short_grant_is_returned_to_spark_and_reported_as_an_error() { + let (pool, _unified, backend) = pool(100); + let reservation = MemoryConsumer::new("test").register(&pool); + + let err = reservation.try_grow(150).unwrap_err(); + assert!( + err.to_string().contains("failed to acquire 150 bytes"), + "unexpected error: {err}" + ); + // The partial grant must not be retained, otherwise the pool would hold + // memory that no reservation accounts for. + assert_eq!(backend.granted(), 0); + assert_eq!(pool.reserved(), 0); + } + + #[test] + fn every_grow_and_shrink_reaches_spark() { + // Records the current behaviour that issue #5383 is about: one round-trip + // per grow and per shrink, with no batching or hysteresis. + let (pool, unified, _) = pool(1024); + let reservation = MemoryConsumer::new("test").register(&pool); + + for _ in 0..10 { + reservation.try_grow(10).unwrap(); + reservation.shrink(10); + } + + let stats = unified.client.stats(); + assert_eq!(stats.acquire_calls(), 10); + assert_eq!(stats.release_calls(), 10); + } +}