sc_transaction_pool/graph/
tracked_map.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
5
6// This program is free software: you can redistribute it and/or modify
7// it under the terms of the GNU General Public License as published by
8// the Free Software Foundation, either version 3 of the License, or
9// (at your option) any later version.
10
11// This program is distributed in the hope that it will be useful,
12// but WITHOUT ANY WARRANTY; without even the implied warranty of
13// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14// GNU General Public License for more details.
15
16// You should have received a copy of the GNU General Public License
17// along with this program. If not, see <https://www.gnu.org/licenses/>.
18
19use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
20use std::{
21	collections::HashMap,
22	sync::{
23		atomic::{AtomicIsize, Ordering as AtomicOrdering},
24		Arc,
25	},
26};
27
28/// Something that can report its size.
29pub trait Size {
30	fn size(&self) -> usize;
31}
32
33/// Map with size tracking.
34///
35/// Size reported might be slightly off and only approximately true.
36#[derive(Debug)]
37pub struct TrackedMap<K, V> {
38	index: Arc<RwLock<HashMap<K, V>>>,
39	bytes: AtomicIsize,
40	length: AtomicIsize,
41}
42
43impl<K, V> Default for TrackedMap<K, V> {
44	fn default() -> Self {
45		Self { index: Arc::new(HashMap::default().into()), bytes: 0.into(), length: 0.into() }
46	}
47}
48
49impl<K, V> TrackedMap<K, V> {
50	/// Current tracked length of the content.
51	pub fn len(&self) -> usize {
52		std::cmp::max(self.length.load(AtomicOrdering::Relaxed), 0) as usize
53	}
54
55	/// Current sum of content length.
56	pub fn bytes(&self) -> usize {
57		std::cmp::max(self.bytes.load(AtomicOrdering::Relaxed), 0) as usize
58	}
59
60	/// Lock map for read.
61	pub fn read(&self) -> TrackedMapReadAccess<K, V> {
62		TrackedMapReadAccess { inner_guard: self.index.read() }
63	}
64
65	/// Lock map for write.
66	pub fn write(&self) -> TrackedMapWriteAccess<K, V> {
67		TrackedMapWriteAccess {
68			inner_guard: self.index.write(),
69			bytes: &self.bytes,
70			length: &self.length,
71		}
72	}
73}
74
75impl<K: Clone, V: Clone> TrackedMap<K, V> {
76	/// Clone the inner map.
77	pub fn clone_map(&self) -> HashMap<K, V> {
78		self.index.read().clone()
79	}
80}
81
82pub struct TrackedMapReadAccess<'a, K, V> {
83	inner_guard: RwLockReadGuard<'a, HashMap<K, V>>,
84}
85
86impl<'a, K, V> TrackedMapReadAccess<'a, K, V>
87where
88	K: Eq + std::hash::Hash,
89{
90	/// Returns true if map contains key.
91	pub fn contains_key(&self, key: &K) -> bool {
92		self.inner_guard.contains_key(key)
93	}
94
95	/// Returns reference to the contained value by key, if exists.
96	pub fn get(&self, key: &K) -> Option<&V> {
97		self.inner_guard.get(key)
98	}
99
100	/// Returns iterator over all values.
101	pub fn values(&self) -> std::collections::hash_map::Values<K, V> {
102		self.inner_guard.values()
103	}
104}
105
106pub struct TrackedMapWriteAccess<'a, K, V> {
107	bytes: &'a AtomicIsize,
108	length: &'a AtomicIsize,
109	inner_guard: RwLockWriteGuard<'a, HashMap<K, V>>,
110}
111
112impl<'a, K, V> TrackedMapWriteAccess<'a, K, V>
113where
114	K: Eq + std::hash::Hash,
115	V: Size,
116{
117	/// Insert value and return previous (if any).
118	pub fn insert(&mut self, key: K, val: V) -> Option<V> {
119		let new_bytes = val.size();
120		self.bytes.fetch_add(new_bytes as isize, AtomicOrdering::Relaxed);
121		self.length.fetch_add(1, AtomicOrdering::Relaxed);
122		self.inner_guard.insert(key, val).inspect(|old_val| {
123			self.bytes.fetch_sub(old_val.size() as isize, AtomicOrdering::Relaxed);
124			self.length.fetch_sub(1, AtomicOrdering::Relaxed);
125		})
126	}
127
128	/// Remove value by key.
129	pub fn remove(&mut self, key: &K) -> Option<V> {
130		let val = self.inner_guard.remove(key);
131		if let Some(size) = val.as_ref().map(Size::size) {
132			self.bytes.fetch_sub(size as isize, AtomicOrdering::Relaxed);
133			self.length.fetch_sub(1, AtomicOrdering::Relaxed);
134		}
135		val
136	}
137
138	/// Returns mutable reference to the contained value by key, if exists.
139	pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
140		self.inner_guard.get_mut(key)
141	}
142}
143
144#[cfg(test)]
145mod tests {
146
147	use super::*;
148
149	impl Size for i32 {
150		fn size(&self) -> usize {
151			*self as usize / 10
152		}
153	}
154
155	#[test]
156	fn basic() {
157		let map = TrackedMap::default();
158		map.write().insert(5, 10);
159		map.write().insert(6, 20);
160
161		assert_eq!(map.bytes(), 3);
162		assert_eq!(map.len(), 2);
163
164		map.write().insert(6, 30);
165
166		assert_eq!(map.bytes(), 4);
167		assert_eq!(map.len(), 2);
168
169		map.write().remove(&6);
170		assert_eq!(map.bytes(), 1);
171		assert_eq!(map.len(), 1);
172	}
173}
sc_transaction_pool/graph/tracked_map.rs

sc_transaction_pool/graph/
tracked_map.rs
