aboutsummaryrefslogtreecommitdiff
path: root/src/util/sled_counter.rs
blob: bc54cea0d0db0631bae3a842e6d936559c1118b9 (plain) (blame)
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
use std::sync::{
	atomic::{AtomicUsize, Ordering},
	Arc,
};

use sled::{CompareAndSwapError, IVec, Iter, Result, Tree};

#[derive(Clone)]
pub struct SledCountedTree(Arc<SledCountedTreeInternal>);

struct SledCountedTreeInternal {
	tree: Tree,
	len: AtomicUsize,
}

impl SledCountedTree {
	pub fn new(tree: Tree) -> Self {
		let len = tree.len();
		Self(Arc::new(SledCountedTreeInternal {
			tree,
			len: AtomicUsize::new(len),
		}))
	}

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

	pub fn is_empty(&self) -> bool {
		self.0.tree.is_empty()
	}

	pub fn get<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<IVec>> {
		self.0.tree.get(key)
	}

	pub fn iter(&self) -> Iter {
		self.0.tree.iter()
	}

	// ---- writing functions ----

	pub fn insert<K, V>(&self, key: K, value: V) -> Result<Option<IVec>>
	where
		K: AsRef<[u8]>,
		V: Into<IVec>,
	{
		let res = self.0.tree.insert(key, value);
		if res == Ok(None) {
			self.0.len.fetch_add(1, Ordering::Relaxed);
		}
		res
	}

	pub fn remove<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<IVec>> {
		let res = self.0.tree.remove(key);
		if matches!(res, Ok(Some(_))) {
			self.0.len.fetch_sub(1, Ordering::Relaxed);
		}
		res
	}

	pub fn pop_min(&self) -> Result<Option<(IVec, IVec)>> {
		let res = self.0.tree.pop_min();
		if let Ok(Some(_)) = &res {
			self.0.len.fetch_sub(1, Ordering::Relaxed);
		};
		res
	}

	pub fn compare_and_swap<K, OV, NV>(
		&self,
		key: K,
		old: Option<OV>,
		new: Option<NV>,
	) -> Result<std::result::Result<(), CompareAndSwapError>>
	where
		K: AsRef<[u8]>,
		OV: AsRef<[u8]>,
		NV: Into<IVec>,
	{
		let old_some = old.is_some();
		let new_some = new.is_some();

		let res = self.0.tree.compare_and_swap(key, old, new);

		if res == Ok(Ok(())) {
			match (old_some, new_some) {
				(false, true) => {
					self.0.len.fetch_add(1, Ordering::Relaxed);
				}
				(true, false) => {
					self.0.len.fetch_sub(1, Ordering::Relaxed);
				}
				_ => (),
			}
		}
		res
	}
}