1use cfx_types::{AddressWithSpace, U256};
6use heap_map::HeapMap;
7use malloc_size_of_derive::MallocSizeOf as DeriveMallocSizeOf;
8use std::cmp::{Ord, Ordering, PartialEq, PartialOrd, Reverse};
9
10#[derive(Default, Eq, PartialEq, Copy, Clone, Debug, DeriveMallocSizeOf)]
16pub struct GarbageCollectorValue {
17 pub count: usize,
20 pub has_ready_tx: bool,
23 pub first_tx_gas_price: U256,
27 pub timestamp: u64,
31}
32
33impl Ord for GarbageCollectorValue {
34 fn cmp(&self, other: &Self) -> Ordering {
35 (
36 self.count,
37 Reverse(self.has_ready_tx),
38 Reverse(self.first_tx_gas_price),
39 Reverse(self.timestamp),
40 )
41 .cmp(&(
42 other.count,
43 Reverse(other.has_ready_tx),
44 Reverse(other.first_tx_gas_price),
45 Reverse(other.timestamp),
46 ))
47 }
48}
49
50impl PartialOrd for GarbageCollectorValue {
51 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
52 Some(self.cmp(other))
53 }
54}
55
56#[derive(Default, DeriveMallocSizeOf)]
59pub struct GarbageCollector {
60 heap_map: HeapMap<AddressWithSpace, GarbageCollectorValue>,
61 gc_size: usize,
62}
63
64impl GarbageCollector {
65 pub fn insert(
68 &mut self, sender: &AddressWithSpace, count: usize, timestamp: u64,
69 has_ready_tx: bool, first_tx_gas_price: U256,
70 ) {
71 let value = GarbageCollectorValue {
72 count,
73 has_ready_tx,
74 first_tx_gas_price,
75 timestamp,
76 };
77 if let Some(origin) = self.heap_map.get(sender) {
78 self.gc_size -= origin.count;
79 }
80 self.gc_size += count;
81 self.heap_map.insert(sender, value);
82 }
83
84 pub fn pop(&mut self) -> Option<(AddressWithSpace, GarbageCollectorValue)> {
89 let item = self.heap_map.pop();
90 if let Some((_, v)) = &item {
91 self.gc_size -= v.count;
92 }
93 item
94 }
95
96 pub fn clear(&mut self) {
97 self.heap_map.clear();
98 self.gc_size = 0;
99 }
100
101 pub fn get_timestamp(&self, sender: &AddressWithSpace) -> Option<u64> {
102 self.heap_map.get(sender).map(|v| v.timestamp)
103 }
104
105 pub fn is_empty(&self) -> bool { self.heap_map.is_empty() }
106
107 #[cfg(test)]
108 pub fn len(&self) -> usize { self.heap_map.len() }
109
110 #[inline]
111 pub fn gc_size(&self) -> usize { self.gc_size }
112
113 pub fn top(&self) -> Option<(&AddressWithSpace, &GarbageCollectorValue)> {
114 self.heap_map.top()
115 }
116}
117
118#[cfg(test)]
119mod garbage_collector_test {
120 use super::{GarbageCollector, GarbageCollectorValue};
121 use cfx_types::{Address, AddressSpaceUtil, AddressWithSpace, U256};
122 use rand::{RngCore, SeedableRng};
123 use rand_xorshift::XorShiftRng;
124 use std::collections::HashMap;
125
126 #[test]
127 fn test_basic_operation() {
128 let mut gc = GarbageCollector::default();
129 assert!(gc.is_empty());
130 assert_eq!(gc.gc_size(), 0);
131 assert!(gc.top().is_none());
132 assert!(gc.pop().is_none());
133
134 let mut addr = Vec::new();
135 for _ in 0..10 {
136 addr.push(Address::random().with_native_space());
137 }
138 gc.insert(&addr[0], 10, 10, false, 0.into());
139 assert_eq!(gc.len(), 1);
140 assert_eq!(gc.gc_size(), 10);
141 assert_eq!(*gc.top().unwrap().0, addr[0]);
142 assert_eq!(gc.top().unwrap().1.count, 10);
143 assert_eq!(gc.top().unwrap().1.timestamp, 10);
144
145 gc.insert(&addr[1], 10, 5, false, 0.into());
146 assert_eq!(gc.len(), 2);
147 assert_eq!(gc.gc_size(), 20);
148 assert_eq!(*gc.top().unwrap().0, addr[1]);
149 assert_eq!(gc.top().unwrap().1.count, 10);
150 assert_eq!(gc.top().unwrap().1.timestamp, 5);
151
152 gc.insert(&addr[2], 11, 5, false, 0.into());
153 assert_eq!(gc.len(), 3);
154 assert_eq!(gc.gc_size(), 31);
155 assert_eq!(*gc.top().unwrap().0, addr[2]);
156 assert_eq!(gc.top().unwrap().1.count, 11);
157 assert_eq!(gc.top().unwrap().1.timestamp, 5);
158
159 gc.insert(&addr[0], 15, 0, false, 0.into());
160 assert_eq!(gc.len(), 3);
161 assert_eq!(gc.gc_size(), 36);
162 assert_eq!(*gc.top().unwrap().0, addr[0]);
163 assert_eq!(gc.top().unwrap().1.count, 15);
164 assert_eq!(gc.top().unwrap().1.timestamp, 0);
165
166 assert_eq!(gc.get_timestamp(&addr[0]), Some(0));
167 assert_eq!(gc.get_timestamp(&addr[1]), Some(5));
168 assert_eq!(gc.get_timestamp(&addr[2]), Some(5));
169 assert_eq!(gc.get_timestamp(&addr[3]), None);
170
171 let top = gc.pop().unwrap();
172 assert_eq!(top.0, addr[0]);
173 assert_eq!(top.1.count, 15);
174 assert_eq!(top.1.timestamp, 0);
175 assert_eq!(gc.len(), 2);
176 assert_eq!(gc.gc_size(), 21);
177 let top = gc.pop().unwrap();
178 assert_eq!(top.0, addr[2]);
179 assert_eq!(top.1.count, 11);
180 assert_eq!(top.1.timestamp, 5);
181 assert_eq!(gc.len(), 1);
182 assert_eq!(gc.gc_size(), 10);
183 let top = gc.pop().unwrap();
184 assert_eq!(top.0, addr[1]);
185 assert_eq!(top.1.count, 10);
186 assert_eq!(top.1.timestamp, 5);
187 assert_eq!(gc.len(), 0);
188 assert_eq!(gc.gc_size(), 0);
189 assert!(gc.pop().is_none());
190 }
191
192 #[test]
193 fn test_ready_accounts() {
194 let mut gc = GarbageCollector::default();
195 assert!(gc.is_empty());
196 assert_eq!(gc.gc_size(), 0);
197 assert!(gc.top().is_none());
198 assert!(gc.pop().is_none());
199
200 let mut addr = Vec::new();
201 for _ in 0..10 {
202 addr.push(Address::random().with_native_space());
203 }
204 gc.insert(&addr[0], 0, 10, false, 0.into());
205 assert_eq!(gc.len(), 1);
206 assert_eq!(gc.gc_size(), 0);
207 assert_eq!(*gc.top().unwrap().0, addr[0]);
208 assert_eq!(gc.top().unwrap().1.count, 0);
209 assert_eq!(gc.top().unwrap().1.timestamp, 10);
210
211 gc.insert(&addr[1], 0, 5, false, 0.into());
212 assert_eq!(gc.len(), 2);
213 assert_eq!(gc.gc_size(), 0);
214 assert_eq!(*gc.top().unwrap().0, addr[1]);
215 assert_eq!(gc.top().unwrap().1.count, 0);
216 assert_eq!(gc.top().unwrap().1.timestamp, 5);
217
218 gc.insert(&addr[1], 0, 5, false, 1.into());
219 assert_eq!(gc.len(), 2);
220 assert_eq!(gc.gc_size(), 0);
221 assert_eq!(*gc.top().unwrap().0, addr[0]);
222 assert_eq!(gc.top().unwrap().1.count, 0);
223 assert_eq!(gc.top().unwrap().1.timestamp, 10);
224
225 gc.insert(&addr[0], 0, 10, true, 1.into());
226 assert_eq!(gc.len(), 2);
227 assert_eq!(gc.gc_size(), 0);
228 assert_eq!(*gc.top().unwrap().0, addr[1]);
229 assert_eq!(gc.top().unwrap().1.count, 0);
230 assert_eq!(gc.top().unwrap().1.timestamp, 5);
231
232 gc.insert(&addr[1], 0, 5, true, 2.into());
233 assert_eq!(gc.len(), 2);
234 assert_eq!(gc.gc_size(), 0);
235 assert_eq!(*gc.top().unwrap().0, addr[0]);
236 assert_eq!(gc.top().unwrap().1.count, 0);
237 assert_eq!(gc.top().unwrap().1.timestamp, 10);
238
239 gc.insert(&addr[0], 0, 10, true, 3.into());
240 assert_eq!(gc.len(), 2);
241 assert_eq!(gc.gc_size(), 0);
242 assert_eq!(*gc.top().unwrap().0, addr[1]);
243 assert_eq!(gc.top().unwrap().1.count, 0);
244 assert_eq!(gc.top().unwrap().1.timestamp, 5);
245
246 gc.insert(&addr[0], 0, 10, false, 1.into());
247 assert_eq!(gc.len(), 2);
248 assert_eq!(gc.gc_size(), 0);
249 assert_eq!(*gc.top().unwrap().0, addr[0]);
250 assert_eq!(gc.top().unwrap().1.count, 0);
251 assert_eq!(gc.top().unwrap().1.timestamp, 10);
252
253 gc.insert(&addr[2], 1, 5, false, 0.into());
254 assert_eq!(gc.len(), 3);
255 assert_eq!(gc.gc_size(), 1);
256 assert_eq!(*gc.top().unwrap().0, addr[2]);
257 assert_eq!(gc.top().unwrap().1.count, 1);
258 assert_eq!(gc.top().unwrap().1.timestamp, 5);
259
260 let top = gc.pop().unwrap();
261 assert_eq!(top.0, addr[2]);
262 assert_eq!(top.1.count, 1);
263 assert_eq!(top.1.timestamp, 5);
264 assert_eq!(top.1.first_tx_gas_price, U256::from(0));
265 assert_eq!(gc.len(), 2);
266 assert_eq!(gc.gc_size(), 0);
267 let top = gc.pop().unwrap();
268 assert_eq!(top.0, addr[0]);
269 assert_eq!(top.1.count, 0);
270 assert_eq!(top.1.timestamp, 10);
271 assert_eq!(top.1.first_tx_gas_price, U256::from(1));
272 assert_eq!(top.1.has_ready_tx, false);
273 assert_eq!(gc.len(), 1);
274 assert_eq!(gc.gc_size(), 0);
275 let top = gc.pop().unwrap();
276 assert_eq!(top.0, addr[1]);
277 assert_eq!(top.1.count, 0);
278 assert_eq!(top.1.timestamp, 5);
279 assert_eq!(top.1.first_tx_gas_price, U256::from(2));
280 assert_eq!(top.1.has_ready_tx, true);
281 assert_eq!(gc.len(), 0);
282 assert_eq!(gc.gc_size(), 0);
283 assert!(gc.pop().is_none());
284 }
285
286 fn get_max(
287 mapping: &HashMap<AddressWithSpace, GarbageCollectorValue>,
288 ) -> Option<GarbageCollectorValue> {
289 mapping
290 .iter()
291 .max_by(|x, y| x.1.cmp(&y.1))
292 .and_then(|x| Some(*x.1))
293 }
294
295 #[test]
296 fn test_correctness() {
297 let mut rng = XorShiftRng::from_os_rng();
298 let mut addr = Vec::new();
299 for _ in 0..10000 {
300 addr.push(Address::random().with_native_space());
301 }
302
303 let mut gc = GarbageCollector::default();
304 let mut mapping = HashMap::new();
305 let mut sum = 0;
306
307 for _ in 0..100000 {
308 let opt: usize = rng.next_u64() as usize % 4;
309 if opt <= 2 {
310 let idx: usize = rng.next_u64() as usize % 10000;
311 let count: usize = rng.next_u64() as usize % 10;
312 let timestamp: u64 = rng.next_u64() % 1000;
313 let has_ready_tx: bool = rng.next_u64() % 2 == 0;
314 let first_tx_gas_price: u64 = rng.next_u64();
315 let node = GarbageCollectorValue {
316 count,
317 has_ready_tx,
318 first_tx_gas_price: first_tx_gas_price.into(),
319 timestamp,
320 };
321 gc.insert(
322 &addr[idx],
323 count,
324 timestamp,
325 has_ready_tx,
326 first_tx_gas_price.into(),
327 );
328 let old = mapping.insert(addr[idx], node);
329 sum += count;
330 if old.is_some() {
331 sum -= old.unwrap().count;
332 }
333 } else {
334 if gc.is_empty() {
335 assert!(gc.pop().is_none());
336 assert_eq!(gc.gc_size(), 0);
337 } else {
338 let max = get_max(&mapping).unwrap();
339 let gc_max = gc.pop().unwrap();
340 assert_eq!(gc_max.1.count, max.count);
341 assert_eq!(gc_max.1.timestamp, max.timestamp);
342 mapping.remove(&gc_max.0);
343 sum -= gc_max.1.count;
344 assert_eq!(gc.len(), mapping.len());
345 }
346 }
347 assert_eq!(gc.len(), mapping.len());
348 assert_eq!(gc.gc_size(), sum);
349 if gc.is_empty() {
350 assert!(mapping.is_empty());
351 } else {
352 assert_eq!(
353 gc.top().unwrap().1.count,
354 get_max(&mapping).unwrap().count
355 );
356 assert_eq!(
357 gc.top().unwrap().1.timestamp,
358 get_max(&mapping).unwrap().timestamp
359 );
360 assert_eq!(
361 gc.top().unwrap().1.has_ready_tx,
362 get_max(&mapping).unwrap().has_ready_tx
363 );
364 assert_eq!(
365 gc.top().unwrap().1.first_tx_gas_price,
366 get_max(&mapping).unwrap().first_tx_gas_price
367 );
368 }
369 }
370 for (addr, value) in mapping.iter() {
371 assert_eq!(value.timestamp, gc.get_timestamp(addr).unwrap());
372 }
373 }
374}