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

Code before:
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 {
            Err(insufficient_capacity_err(
                reservation,
                additional,
                self.ctx.mem_limit() as usize - self.reserved(),
            ))
        }
    }

    fn reserved(&self) -> usize {
        let bytes_used = self.ctx.get_bytes_used();
        if bytes_used <= 0 {
            0
        } else if let Ok(v) = usize::try_from(bytes_used) {
            v
        } else {
            usize::MAX
        }
    }

    fn memory_limit(&self) -> MemoryLimit {
        MemoryLimit::Finite(self.ctx.mem_limit() as usize)
    }
}

Review comment (bug):
The `grow` method only logs a warning when memory allocation fails, but DataFusion's `MemoryPool::grow` trait method is documented as a method that should succeed. If the memory limit is exceeded, the available memory calculation `self.ctx.mem_limit() as usize - self.reserved()` might underflow if `self.reserved()` is greater than `mem_limit()`, which could happen with concurrent allocations. Consider using `saturating_sub` to prevent potential underflow in the available memory calculation.

Code after:
    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) {
        self.ctx.add(-(shrink as i64));
    }

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

    fn reserved(&self) -> usize {
        let bytes_used = self.ctx.get_bytes_used();
        if bytes_used <= 0 {
            0
        } else {
            bytes_used.try_into().unwrap_or(usize::MAX)
        }
    }