Repository: risingwavelabs/risingwave
File: src/frontend/src/datafusion/execute/memory_ctx.rs (Rust)

Code before:
// See the License for the specific language governing permissions and
// limitations under the License.

use datafusion::execution::memory_pool::{
    MemoryLimit, MemoryPool, MemoryReservation, human_readable_size,
};
use datafusion_common::{DataFusionError, Result as DFResult, resources_datafusion_err};
use prometheus::core::Atomic;
use risingwave_common::memory::MemoryContext;
use risingwave_common::metrics::TrAdderAtomic;

pub struct RwMemoryPool {
    ctx: MemoryContext,
}

impl RwMemoryPool {
    pub fn new(parent: MemoryContext) -> Self {
        let counter = TrAdderAtomic::new(0);
        let ctx = MemoryContext::new(Some(parent), counter);
        Self { ctx }
    }
}

impl std::fmt::Debug for RwMemoryPool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "DFMemoryPool")
    }
}

impl MemoryPool for RwMemoryPool {
    fn grow(&self, reservation: &MemoryReservation, additional: usize) {
        let success = self.ctx.add(additional as i64);
        if !success {
            tracing::warn!(
                info = %insufficient_capacity_err(
                    reservation,
                    additional,
                    self.ctx.mem_limit() as usize - self.reserved(),
                )
            )
        }
    }

    fn shrink(&self, _reservation: &MemoryReservation, shrink: usize) {
        self.ctx.add(-(shrink as i64));
    }

    fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> DFResult<()> {
        if self.ctx.add(additional as i64) {
            Ok(())
        } else {

Review comment (bug):
The Debug implementation outputs "DFMemoryPool" but the struct is named "RwMemoryPool". For consistency and clarity in debug logs, the Debug output should match the struct name.
```suggestion
        write!(f, "RwMemoryPool")
```

Code after:
// See the License for the specific language governing permissions and
// limitations under the License.

use datafusion::execution::memory_pool::{
    MemoryLimit, MemoryPool, MemoryReservation, human_readable_size,
};
use datafusion_common::{DataFusionError, Result as DFResult, resources_datafusion_err};
use prometheus::core::Atomic;
use risingwave_common::memory::MemoryContext;
use risingwave_common::metrics::TrAdderAtomic;
use thiserror_ext::AsReport;

pub struct RwMemoryPool {
    ctx: MemoryContext,
}

impl RwMemoryPool {
    pub fn new(parent: MemoryContext) -> Self {
        let counter = TrAdderAtomic::new(0);
        let ctx = MemoryContext::new(Some(parent), counter);
        Self { ctx }
    }

    fn available(&self) -> usize {
        let limit = self.ctx.mem_limit().try_into().unwrap_or(usize::MAX);
        let used = self.reserved();
        limit.saturating_sub(used)
    }
}

impl std::fmt::Debug for RwMemoryPool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "RwMemoryPool")
    }
}

impl MemoryPool for RwMemoryPool {
    fn grow(&self, reservation: &MemoryReservation, additional: usize) {
        let success = self.ctx.add(additional as i64);
        if !success {
            tracing::warn!(
                error = %insufficient_capacity_err(
                    reservation,
                    additional,
                    self.available()
                ).as_report()
            )
        }
    }

    fn shrink(&self, _reservation: &MemoryReservation, shrink: usize) {