Repository: tikv/tikv
File: components/tikv_util/src/memory.rs (Rust)

Code before:
        self.capacity.load(Ordering::Relaxed)
    }

    pub fn set_capacity(&self, mut capacity: usize) {
        // Value bigger than isize::MAX just means unlimited,
        // so replace it with isize::MAX to avoid overflow.
        if capacity > isize::MAX as usize {
            capacity = isize::MAX as usize;
        }
        self.capacity.store(capacity, Ordering::Relaxed);
    }

    pub fn alloc_force(&self, bytes: usize) {
        self.in_use.fetch_add(bytes as isize, Ordering::Relaxed);
    }

    pub fn alloc(&self, bytes: usize) -> Result<(), MemoryQuotaExceeded> {
        let capacity = self.capacity.load(Ordering::Relaxed) as isize;
        let in_use_bytes = self.in_use.load(Ordering::Relaxed);
        if bytes > capacity.saturating_sub(in_use_bytes) as usize {
            return Err(MemoryQuotaExceeded);
        }
        let bytes = bytes as isize;
        let new_in_use_bytes = self.in_use.fetch_add(bytes, Ordering::Relaxed);
        if bytes > capacity - new_in_use_bytes {
            self.in_use.fetch_sub(bytes, Ordering::Relaxed);
            return Err(MemoryQuotaExceeded);
        }
        Ok(())
    }

    pub fn free(&self, bytes: usize) {
        // Alloc memory higher than isize::MAX must fail because it exceeds the maxinum
        // capacity, but we still try to handle the overflow here to avoid
        // caller's misuse.
        let bytes = std::cmp::min(bytes, isize::MAX as usize) as isize;
        let new_in_use_bytes = self.in_use.fetch_sub(bytes, Ordering::Relaxed) - bytes;
        if new_in_use_bytes < 0 {
            // handle overflow.
            self.in_use
                .fetch_add(std::cmp::min(-new_in_use_bytes, bytes), Ordering::Relaxed);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_memory_quota() {

Review comment (bug):
Another issue

https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=69bef8ec9096033e89d90b3c6b86d16a

Code after:
        // Saturating at the numeric bounds instead of overflowing.
        // we handle negative overflow here to make `free` implementation simpler.
        std::cmp::max(value, 0) as usize
    }

    /// Returns a floating number between [0, 1] presents the current memory
    /// status.
    pub fn used_ratio(&self) -> f64 {
        self.in_use() as f64 / self.capacity() as f64
    }

    pub fn capacity(&self) -> usize {
        self.capacity.load(Ordering::Relaxed)
    }

    pub fn set_capacity(&self, capacity: usize) {
        let capacity = Self::adjust_capacity(capacity);
        self.capacity.store(capacity, Ordering::Relaxed);
    }

    pub fn alloc_force(&self, bytes: usize) {
        if bytes > MAX_MEMORY_ALLOC_SIZE {
            warn!("bytes size exceeds the max memory alloc size, force alloc is ignored"; "bytes" => bytes);
            return;
        }
        self.in_use.fetch_add(bytes as isize, Ordering::Relaxed);
    }

    pub fn alloc(&self, bytes: usize) -> Result<(), MemoryQuotaExceeded> {
        let capacity = self.capacity.load(Ordering::Relaxed) as isize;
        let in_use_bytes = self.in_use.load(Ordering::Relaxed);
        if bytes > capacity.saturating_sub(in_use_bytes) as usize || bytes > MAX_MEMORY_ALLOC_SIZE {
            return Err(MemoryQuotaExceeded);
        }
        let bytes = bytes as isize;
        let new_in_use_bytes = self.in_use.fetch_add(bytes, Ordering::Relaxed);
        if bytes > capacity - new_in_use_bytes {
            self.in_use.fetch_sub(bytes, Ordering::Relaxed);
            return Err(MemoryQuotaExceeded);
        }
        Ok(())
    }

    pub fn free(&self, bytes: usize) {
        // bytes higher than `MAX_MEMORY_ALLOC_SIZE` means the memory is acquired with
        // `alloc_force`, so also ignore the `free` call.
        if bytes > MAX_MEMORY_ALLOC_SIZE {
            warn!("bytes size exceeds the max memory alloc size, free is ignored"; "bytes" => bytes);
            return;
        }
        let bytes = bytes as isize;