1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

//! # raii-counter
//! Rust type for a RAII Counter (counts number of held instances,
//! decrements count on `Drop`), implemented with `Arc<AtomicUsize>`.
//!
//! Useful for tracking the number of holders exist for a handle,
//! tracking the number of transactions that are in-flight, etc.
//!
//! ## Demo
//!
//! ```rust
//! extern crate raii_counter;
//! use raii_counter::Counter;
//!
//! let counter = Counter::new();
//! assert_eq!(counter.count(), 1);
//!
//! let weak = counter.downgrade();
//! assert_eq!(weak.count(), 0);
//!
//! {
//!     let _counter1 = weak.spawn_upgrade();
//!     assert_eq!(weak.count(), 1);
//!     let _counter2 = weak.spawn_upgrade();
//!     assert_eq!(weak.count(), 2);
//! }
//!
//! assert_eq!(weak.count(), 0);
//! ```

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

/// Essentially an AtomicUsize that is clonable and whose count is based
/// on the number of copies. The count is automatically updated on Drop.
pub struct Counter {
    counter: Arc<AtomicUsize>,
    size: usize,
}

/// A 'weak' Counter that does not affect the count.
#[derive(Clone)]
pub struct WeakCounter(Arc<AtomicUsize>);

impl Counter {
    pub fn new() -> Counter {
        Counter::new_with_size(1)
    }

    pub fn new_with_size(size: usize) -> Counter {
        Counter {
            counter: Arc::new(AtomicUsize::new(1)),
            size,
        }
    }

    /// Consume self (causing the count to decrease by 1)
    /// and return a weak reference to the count through a WeakCounter
    pub fn downgrade(self) -> WeakCounter {
        WeakCounter(self.counter.clone())
    }

    /// This method is inherently racey. Assume the count will have changed once
    /// the value is observed.
    #[inline]
    pub fn count(&self) -> usize {
        self.counter.load(Ordering::Acquire)
    }
}

impl Clone for Counter {
    fn clone(&self) -> Self {
        self.counter.fetch_add(self.size, Ordering::AcqRel);
        Counter { counter: self.counter.clone(), size: self.size }
    }
}

impl Drop for Counter {
    fn drop(&mut self) {
        self.counter.fetch_sub(self.size, Ordering::AcqRel);
    }
}

impl WeakCounter {
    pub fn new() -> WeakCounter {
        WeakCounter(Arc::new(AtomicUsize::new(0)))
    }

    /// This method is inherently racey. Assume the count will have changed once
    /// the value is observed.
    #[inline]
    pub fn count(&self) -> usize {
        self.0.load(Ordering::Acquire)
    }

    /// Consumes self, becomes a Counter
    pub fn upgrade(self) -> Counter {
        self.spawn_upgrade()
    }

    /// Instead of clone + upgrade, this will only clone once
    /// Defaults to a Counter of size 1
    pub fn spawn_upgrade(&self) -> Counter {
        self.spawn_upgrade_with_size(1)
    }

    /// Instead of clone + upgrade, this will only clone once
    pub fn spawn_upgrade_with_size(&self, size: usize) -> Counter {
        self.0.fetch_add(size, Ordering::AcqRel);
        Counter { counter: self.0.clone(), size }
    }
}

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

    #[test]
    fn it_works() {
        let counter = Counter::new();
        assert_eq!(counter.count(), 1);

        let weak = counter.downgrade();
        assert_eq!(weak.count(), 0);

        {
            let _counter1 = weak.spawn_upgrade();
            assert_eq!(weak.count(), 1);
            let _counter2 = weak.spawn_upgrade();
            assert_eq!(weak.count(), 2);
        }

        assert_eq!(weak.count(), 0);
    }

    #[test]
    fn different_sizes_work() {
        let weak = WeakCounter::new();
        assert_eq!(weak.count(), 0);

        let counter = weak.spawn_upgrade_with_size(5);
        assert_eq!(weak.count(), 5);

        {
            let _counter1 = counter.clone();
            assert_eq!(weak.count(), 10);
            let _counter2 = weak.spawn_upgrade();
            assert_eq!(weak.count(), 11);
        }

        assert_eq!(weak.count(), 5);
    }
}