changelog shortlog graph tags branches changeset files revisions annotate raw help

Mercurial > core / rust/lib/alch/src/tr.rs

changeset 698: 96958d3eb5b0
parent: 1227f932b628
author: Richard Westhaver <ellis@rwest.io>
date: Fri, 04 Oct 2024 22:04:59 -0400
permissions: -rw-r--r--
description: fixes 
     1 use std::{
     2   alloc::{GlobalAlloc, Layout, System},
     3   sync::atomic::{AtomicU64, Ordering},
     4 };
     5 
     6 /// a simple global allocator for tracing memory usage
     7 /// from https://github.com/datafuselabs/databend/issues/1148#issuecomment-907829698
     8 pub struct TrAlloc(System, AtomicU64, AtomicU64);
     9 
    10 unsafe impl GlobalAlloc for TrAlloc {
    11   unsafe fn alloc(&self, l: Layout) -> *mut u8 {
    12     self.1.fetch_add(l.size() as u64, Ordering::SeqCst);
    13     self.2.fetch_add(l.size() as u64, Ordering::SeqCst);
    14     self.0.alloc(l)
    15   }
    16   unsafe fn dealloc(&self, ptr: *mut u8, l: Layout) {
    17     self.0.dealloc(ptr, l);
    18     self.1.fetch_sub(l.size() as u64, Ordering::SeqCst);
    19   }
    20 }
    21 
    22 impl TrAlloc {
    23   pub const fn new(s: System) -> Self {
    24     TrAlloc(s, AtomicU64::new(0), AtomicU64::new(0))
    25   }
    26 
    27   pub fn reset(&self) {
    28     self.1.store(0, Ordering::SeqCst);
    29     self.2.store(0, Ordering::SeqCst);
    30   }
    31 
    32   pub fn get(&self) -> u64 {
    33     self.1.load(Ordering::SeqCst)
    34   }
    35 
    36   pub fn get_sum(&self) -> u64 {
    37     self.2.load(Ordering::SeqCst)
    38   }
    39 }