1
use std::sync::atomic::Ordering as O;
2

            
3
pub struct Bucket {
4
    level:	std::sync::atomic::AtomicU32,
5
}
6

            
7
impl Bucket {
8
1
    pub fn new(level: u32) -> Self {
9
1
	Self {
10
1
	    level:	level.into(),
11
1
	}
12
1
    }
13

            
14
    #[cfg(test)]
15
10
    pub fn level(&self) -> u32 {
16
10
	self.level.load(O::Relaxed)
17
10
    }
18

            
19
6
    pub fn acquire(&self) -> BucketGuard {
20
6
	let do_release = self.level
21
6
	    .fetch_update(O::Relaxed, O::Relaxed, |v| match v {
22
1
		0	=> None,
23
5
		v	=> Some(v - 1)
24
6
	    }).is_ok();
25
6

            
26
6
	BucketGuard {
27
6
	    bucket:	self,
28
6
	    do_release:	do_release,
29
6
	}
30
6
    }
31

            
32
5
    fn release(&self) {
33
5
	self.level.fetch_add(1, O::Relaxed);
34
5
    }
35
}
36

            
37
pub struct BucketGuard<'a> {
38
    bucket:	&'a Bucket,
39
    do_release:	bool,
40
}
41

            
42
impl Drop for BucketGuard<'_> {
43
6
    fn drop(&mut self) {
44
6
	if self.do_release {
45
5
	    self.bucket.release()
46
1
	}
47
6
    }
48
}
49

            
50
impl BucketGuard<'_> {
51
6
    pub fn is_ok(&self) -> bool {
52
6
	self.do_release
53
6
    }
54

            
55
1
    pub fn release(self) {
56
1
    }
57
}
58

            
59
#[cfg(test)]
60
mod test {
61
    use super::*;
62

            
63
1
    #[test]
64
1
    fn test_bucket() {
65
1
	let bucket = Bucket::new(4);
66
1

            
67
1
	assert_eq!(bucket.level(), 4);
68

            
69
	{
70
1
	    let g0 = bucket.acquire();
71
1
	    assert_eq!(bucket.level(), 3);
72
1
	    assert!(g0.is_ok());
73

            
74
1
	    let g1 = bucket.acquire();
75
1
	    assert_eq!(bucket.level(), 2);
76
1
	    assert!(g1.is_ok());
77

            
78
	    {
79
1
		let g2 = bucket.acquire();
80
1
		assert_eq!(bucket.level(), 1);
81
1
		assert!(g2.is_ok());
82

            
83
1
		let g3 = bucket.acquire();
84
1
		assert_eq!(bucket.level(), 0);
85
1
		assert!(g3.is_ok());
86

            
87
		{
88
1
		    let g4 = bucket.acquire();
89
1
		    assert_eq!(bucket.level(), 0);
90
1
		    assert!(!g4.is_ok());
91
		}
92

            
93
1
		assert_eq!(bucket.level(), 0);
94

            
95
1
		g3.release();
96
1

            
97
1
		assert_eq!(bucket.level(), 1);
98

            
99
1
		let g5 = bucket.acquire();
100
1
		assert_eq!(bucket.level(), 0);
101
1
		assert!(g5.is_ok());
102
	    }
103

            
104
1
	    assert_eq!(bucket.level(), 2);
105
	}
106
1
    }
107
}