1+ // This file is part of Substrate.
2+
3+ // Copyright (C) 2018-2020 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+
19+ use std:: {
20+ collections:: HashMap ,
21+ sync:: { Arc , atomic:: { AtomicIsize , Ordering as AtomicOrdering } } ,
22+ } ;
23+ use parking_lot:: { RwLock , RwLockWriteGuard , RwLockReadGuard } ;
24+
25+ /// Something that can report it's size.
26+ pub trait TrackedSize {
27+ fn tracked_size ( & self ) -> usize ;
28+ }
29+
30+ /// Map with size tracking.
31+ ///
32+ /// Size reported might be slightly off and only approximately true.
33+ #[ derive( Debug , parity_util_mem:: MallocSizeOf ) ]
34+ pub struct TrackedMap < K , V > {
35+ index : Arc < RwLock < HashMap < K , V > > > ,
36+ bytes : AtomicIsize ,
37+ length : AtomicIsize ,
38+ }
39+
40+ impl < K , V > Default for TrackedMap < K , V > {
41+ fn default ( ) -> Self {
42+ Self {
43+ index : Arc :: new ( HashMap :: default ( ) . into ( ) ) ,
44+ bytes : 0 . into ( ) ,
45+ length : 0 . into ( ) ,
46+ }
47+ }
48+ }
49+
50+ impl < K , V > TrackedMap < K , V > {
51+ /// Current tracked length of the content.
52+ pub fn len ( & self ) -> usize {
53+ std:: cmp:: max ( self . length . load ( AtomicOrdering :: Relaxed ) , 0 ) as usize
54+ }
55+
56+ /// Current sum of content length.
57+ pub fn bytes ( & self ) -> usize {
58+ std:: cmp:: max ( self . bytes . load ( AtomicOrdering :: Relaxed ) , 0 ) as usize
59+ }
60+
61+ /// Read-only clone of the interior.
62+ pub fn clone ( & self ) -> ReadOnlyTrackedMap < K , V > {
63+ ReadOnlyTrackedMap ( self . index . clone ( ) )
64+ }
65+
66+ /// Lock map for read.
67+ pub fn read < ' a > ( & ' a self ) -> TrackedMapReadAccess < ' a , K , V > {
68+ TrackedMapReadAccess {
69+ inner_guard : self . index . read ( ) ,
70+ }
71+ }
72+
73+ /// Lock map for write.
74+ pub fn write < ' a > ( & ' a self ) -> TrackedMapWriteAccess < ' a , K , V > {
75+ TrackedMapWriteAccess {
76+ inner_guard : self . index . write ( ) ,
77+ bytes : & self . bytes ,
78+ length : & self . length ,
79+ }
80+ }
81+ }
82+
83+ /// Read-only access to map.
84+ ///
85+ /// The only thing can be done is .read().
86+ pub struct ReadOnlyTrackedMap < K , V > ( Arc < RwLock < HashMap < K , V > > > ) ;
87+
88+ impl < K , V > ReadOnlyTrackedMap < K , V >
89+ where
90+ K : Eq + std:: hash:: Hash
91+ {
92+ /// Lock map for read.
93+ pub fn read < ' a > ( & ' a self ) -> TrackedMapReadAccess < ' a , K , V > {
94+ TrackedMapReadAccess {
95+ inner_guard : self . 0 . read ( ) ,
96+ }
97+ }
98+ }
99+
100+ pub struct TrackedMapReadAccess < ' a , K , V > {
101+ inner_guard : RwLockReadGuard < ' a , HashMap < K , V > > ,
102+ }
103+
104+ impl < ' a , K , V > TrackedMapReadAccess < ' a , K , V >
105+ where
106+ K : Eq + std:: hash:: Hash
107+ {
108+ /// Returns true if map contains key.
109+ pub fn contains_key ( & self , key : & K ) -> bool {
110+ self . inner_guard . contains_key ( key)
111+ }
112+
113+ /// Returns reference to the contained value by key, if exists.
114+ pub fn get ( & self , key : & K ) -> Option < & V > {
115+ self . inner_guard . get ( key)
116+ }
117+
118+ /// Returns iterator over all values.
119+ pub fn values ( & self ) -> std:: collections:: hash_map:: Values < K , V > {
120+ self . inner_guard . values ( )
121+ }
122+ }
123+
124+ pub struct TrackedMapWriteAccess < ' a , K , V > {
125+ bytes : & ' a AtomicIsize ,
126+ length : & ' a AtomicIsize ,
127+ inner_guard : RwLockWriteGuard < ' a , HashMap < K , V > > ,
128+ }
129+
130+ impl < ' a , K , V > TrackedMapWriteAccess < ' a , K , V >
131+ where
132+ K : Eq + std:: hash:: Hash , V : TrackedSize
133+ {
134+ /// Insert value and return previous (if any).
135+ pub fn insert ( & mut self , key : K , val : V ) -> Option < V > {
136+ let new_bytes = val. tracked_size ( ) ;
137+ self . bytes . fetch_add ( new_bytes as isize , AtomicOrdering :: Relaxed ) ;
138+ self . length . fetch_add ( 1 , AtomicOrdering :: Relaxed ) ;
139+ self . inner_guard . insert ( key, val) . and_then ( |old_val| {
140+ self . bytes . fetch_sub ( old_val. tracked_size ( ) as isize , AtomicOrdering :: Relaxed ) ;
141+ self . length . fetch_sub ( 1 , AtomicOrdering :: Relaxed ) ;
142+ Some ( old_val)
143+ } )
144+ }
145+
146+ /// Remove value by key.
147+ pub fn remove ( & mut self , key : & K ) -> Option < V > {
148+ let val = self . inner_guard . remove ( key) ;
149+ if let Some ( size) = val. as_ref ( ) . map ( TrackedSize :: tracked_size) {
150+ self . bytes . fetch_sub ( size as isize , AtomicOrdering :: Relaxed ) ;
151+ self . length . fetch_sub ( 1 , AtomicOrdering :: Relaxed ) ;
152+ }
153+ val
154+ }
155+
156+ /// Returns mutable reference to the contained value by key, if exists.
157+ pub fn get_mut ( & mut self , key : & K ) -> Option < & mut V > {
158+ self . inner_guard . get_mut ( key)
159+ }
160+ }
161+
162+ #[ cfg( test) ]
163+ mod tests {
164+
165+ use super :: * ;
166+
167+ impl TrackedSize for i32 {
168+ fn tracked_size ( & self ) -> usize { * self as usize / 10 }
169+ }
170+
171+ #[ test]
172+ fn basic ( ) {
173+ let map = TrackedMap :: default ( ) ;
174+ map. write ( ) . insert ( 5 , 10 ) ;
175+ map. write ( ) . insert ( 6 , 20 ) ;
176+
177+ assert_eq ! ( map. bytes( ) , 3 ) ;
178+ assert_eq ! ( map. len( ) , 2 ) ;
179+
180+ map. write ( ) . insert ( 6 , 30 ) ;
181+
182+ assert_eq ! ( map. bytes( ) , 4 ) ;
183+ assert_eq ! ( map. len( ) , 2 ) ;
184+
185+ map. write ( ) . remove ( & 6 ) ;
186+ assert_eq ! ( map. bytes( ) , 1 ) ;
187+ assert_eq ! ( map. len( ) , 1 ) ;
188+ }
189+ }
0 commit comments