postage/sync/
ref_count.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3#[derive(Debug)]
4pub struct RefCount {
5    count: AtomicUsize,
6}
7
8pub enum TryDecrement {
9    Alive(usize),
10    Dead,
11}
12
13impl TryDecrement {
14    #[allow(dead_code)]
15    #[track_caller]
16    pub fn expect_dead(&self, message: &str) {
17        if let Self::Alive(_) = self {
18            panic!("TryDecrement unwrapped on an Alive value: {}", message);
19        }
20    }
21}
22
23impl RefCount {
24    pub fn new(count: usize) -> Self {
25        Self {
26            count: AtomicUsize::new(count),
27        }
28    }
29
30    pub fn is_alive(&self) -> bool {
31        self.count.load(Ordering::Acquire) > 0
32    }
33
34    pub fn increment(&self) {
35        self.count.fetch_add(1, Ordering::AcqRel);
36    }
37
38    pub fn decrement(&self) -> TryDecrement {
39        loop {
40            let state = self.count.load(Ordering::Acquire);
41
42            if state == 0 {
43                return TryDecrement::Dead;
44            }
45
46            if let Ok(_prev) =
47                self.count
48                    .compare_exchange(state, state - 1, Ordering::AcqRel, Ordering::Relaxed)
49            {
50                if state == 1 {
51                    return TryDecrement::Dead;
52                } else {
53                    return TryDecrement::Alive(state - 1);
54                }
55            }
56        }
57    }
58}
postage/sync/ref_count.rs

postage/sync/
ref_count.rs
