|
| 1 | +// Licensed to the Apache Software Foundation (ASF) under one |
| 2 | +// or more contributor license agreements. See the NOTICE file |
| 3 | +// distributed with this work for additional information |
| 4 | +// regarding copyright ownership. The ASF licenses this file |
| 5 | +// to you under the Apache License, Version 2.0 (the |
| 6 | +// "License"); you may not use this file except in compliance |
| 7 | +// with the License. You may obtain a copy of the License at |
| 8 | +// |
| 9 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +// |
| 11 | +// Unless required by applicable law or agreed to in writing, |
| 12 | +// software distributed under the License is distributed on an |
| 13 | +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +// KIND, either express or implied. See the License for the |
| 15 | +// specific language governing permissions and limitations |
| 16 | +// under the License. |
| 17 | + |
| 18 | +use std::{ |
| 19 | + fmt::{Debug, Formatter, Result as FmtResult}, |
| 20 | + sync::Arc, |
| 21 | +}; |
| 22 | + |
| 23 | +use jni::objects::GlobalRef; |
| 24 | + |
| 25 | +use crate::{ |
| 26 | + errors::CometResult, |
| 27 | + jvm_bridge::{jni_call, JVMClasses}, |
| 28 | +}; |
| 29 | +use datafusion::{ |
| 30 | + common::DataFusionError, |
| 31 | + execution::memory_pool::{MemoryPool, MemoryReservation}, |
| 32 | +}; |
| 33 | +use datafusion_common::resources_err; |
| 34 | +use datafusion_execution::memory_pool::MemoryConsumer; |
| 35 | +use parking_lot::Mutex; |
| 36 | + |
| 37 | +/// A DataFusion fair `MemoryPool` implementation for Comet. Internally this is |
| 38 | +/// implemented via delegating calls to [`crate::jvm_bridge::CometTaskMemoryManager`]. |
| 39 | +pub struct CometFairMemoryPool { |
| 40 | + task_memory_manager_handle: Arc<GlobalRef>, |
| 41 | + pool_size: usize, |
| 42 | + state: Mutex<CometFairPoolState>, |
| 43 | +} |
| 44 | + |
| 45 | +struct CometFairPoolState { |
| 46 | + used: usize, |
| 47 | + num: usize, |
| 48 | +} |
| 49 | + |
| 50 | +impl Debug for CometFairMemoryPool { |
| 51 | + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { |
| 52 | + let state = self.state.lock(); |
| 53 | + f.debug_struct("CometFairMemoryPool") |
| 54 | + .field("pool_size", &self.pool_size) |
| 55 | + .field("used", &state.used) |
| 56 | + .field("num", &state.num) |
| 57 | + .finish() |
| 58 | + } |
| 59 | +} |
| 60 | + |
| 61 | +impl CometFairMemoryPool { |
| 62 | + pub fn new( |
| 63 | + task_memory_manager_handle: Arc<GlobalRef>, |
| 64 | + pool_size: usize, |
| 65 | + ) -> CometFairMemoryPool { |
| 66 | + Self { |
| 67 | + task_memory_manager_handle, |
| 68 | + pool_size, |
| 69 | + state: Mutex::new(CometFairPoolState { used: 0, num: 0 }), |
| 70 | + } |
| 71 | + } |
| 72 | + |
| 73 | + fn acquire(&self, additional: usize) -> CometResult<i64> { |
| 74 | + let mut env = JVMClasses::get_env()?; |
| 75 | + let handle = self.task_memory_manager_handle.as_obj(); |
| 76 | + unsafe { |
| 77 | + jni_call!(&mut env, |
| 78 | + comet_task_memory_manager(handle).acquire_memory(additional as i64) -> i64) |
| 79 | + } |
| 80 | + } |
| 81 | + |
| 82 | + fn release(&self, size: usize) -> CometResult<()> { |
| 83 | + let mut env = JVMClasses::get_env()?; |
| 84 | + let handle = self.task_memory_manager_handle.as_obj(); |
| 85 | + unsafe { |
| 86 | + jni_call!(&mut env, comet_task_memory_manager(handle).release_memory(size as i64) -> ()) |
| 87 | + } |
| 88 | + } |
| 89 | +} |
| 90 | + |
| 91 | +unsafe impl Send for CometFairMemoryPool {} |
| 92 | +unsafe impl Sync for CometFairMemoryPool {} |
| 93 | + |
| 94 | +impl MemoryPool for CometFairMemoryPool { |
| 95 | + fn register(&self, _: &MemoryConsumer) { |
| 96 | + let mut state = self.state.lock(); |
| 97 | + state.num = state |
| 98 | + .num |
| 99 | + .checked_add(1) |
| 100 | + .expect("unexpected amount of register happened"); |
| 101 | + } |
| 102 | + |
| 103 | + fn unregister(&self, _: &MemoryConsumer) { |
| 104 | + let mut state = self.state.lock(); |
| 105 | + state.num = state |
| 106 | + .num |
| 107 | + .checked_sub(1) |
| 108 | + .expect("unexpected amount of unregister happened"); |
| 109 | + } |
| 110 | + |
| 111 | + fn grow(&self, reservation: &MemoryReservation, additional: usize) { |
| 112 | + self.try_grow(reservation, additional).unwrap(); |
| 113 | + } |
| 114 | + |
| 115 | + fn shrink(&self, reservation: &MemoryReservation, subtractive: usize) { |
| 116 | + if subtractive > 0 { |
| 117 | + let mut state = self.state.lock(); |
| 118 | + let size = reservation.size(); |
| 119 | + if size < subtractive { |
| 120 | + panic!("Failed to release {subtractive} bytes where only {size} bytes reserved") |
| 121 | + } |
| 122 | + self.release(subtractive) |
| 123 | + .unwrap_or_else(|_| panic!("Failed to release {} bytes", subtractive)); |
| 124 | + state.used = state.used.checked_sub(subtractive).unwrap(); |
| 125 | + } |
| 126 | + } |
| 127 | + |
| 128 | + fn try_grow( |
| 129 | + &self, |
| 130 | + reservation: &MemoryReservation, |
| 131 | + additional: usize, |
| 132 | + ) -> Result<(), DataFusionError> { |
| 133 | + if additional > 0 { |
| 134 | + let mut state = self.state.lock(); |
| 135 | + let num = state.num; |
| 136 | + let limit = self.pool_size.checked_div(num).unwrap(); |
| 137 | + let size = reservation.size(); |
| 138 | + if limit < size + additional { |
| 139 | + return resources_err!( |
| 140 | + "Failed to acquire {additional} bytes where {size} bytes already reserved and the fair limit is {limit} bytes, {num} registered" |
| 141 | + ); |
| 142 | + } |
| 143 | + |
| 144 | + let acquired = self.acquire(additional)?; |
| 145 | + // If the number of bytes we acquired is less than the requested, return an error, |
| 146 | + // and hopefully will trigger spilling from the caller side. |
| 147 | + if acquired < additional as i64 { |
| 148 | + // Release the acquired bytes before throwing error |
| 149 | + self.release(acquired as usize)?; |
| 150 | + |
| 151 | + return resources_err!( |
| 152 | + "Failed to acquire {} bytes, only got {} bytes. Reserved: {} bytes", |
| 153 | + additional, |
| 154 | + acquired, |
| 155 | + state.used |
| 156 | + ); |
| 157 | + } |
| 158 | + state.used = state.used.checked_add(additional).unwrap(); |
| 159 | + } |
| 160 | + Ok(()) |
| 161 | + } |
| 162 | + |
| 163 | + fn reserved(&self) -> usize { |
| 164 | + self.state.lock().used |
| 165 | + } |
| 166 | +} |
0 commit comments