cfxcore/pos/mempool/core_mempool/
ttl_cache.rs

1// Copyright (c) The Diem Core Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4// Copyright 2021 Conflux Foundation. All rights reserved.
5// Conflux is free software and distributed under GNU General Public License.
6// See http://www.gnu.org/licenses/
7
8use std::{
9    collections::{BTreeMap, HashMap},
10    time::{Duration, SystemTime},
11};
12
13struct ValueInfo<V> {
14    value: V,
15    ttl: SystemTime,
16}
17
18pub struct TtlCache<K, V> {
19    capacity: usize,
20    default_timeout: Duration,
21    data: HashMap<K, ValueInfo<V>>,
22    ttl_index: BTreeMap<SystemTime, K>,
23}
24
25impl<K, V> TtlCache<K, V>
26where K: std::cmp::Eq + std::hash::Hash + std::clone::Clone
27{
28    pub fn new(capacity: usize, default_timeout: Duration) -> Self {
29        Self {
30            capacity,
31            default_timeout,
32            data: HashMap::new(),
33            ttl_index: BTreeMap::new(),
34        }
35    }
36
37    pub fn get(&self, key: &K) -> Option<&V> {
38        self.data.get(key).map(|v| &v.value)
39    }
40
41    pub fn insert(&mut self, key: K, value: V) {
42        // Remove old entry if it exists.
43        match self.data.get(&key) {
44            Some(info) => {
45                self.ttl_index.remove(&info.ttl);
46            }
47            None => {
48                // Remove oldest entry if cache is still full.
49                if self.data.len() == self.capacity {
50                    let first_entry = self.ttl_index.keys().next().cloned();
51                    if let Some(tst) = first_entry {
52                        if let Some(key) = self.ttl_index.remove(&tst) {
53                            self.data.remove(&key);
54                        }
55                    }
56                }
57            }
58        }
59
60        // Insert the new transaction.
61        if let Some(expiration_time) =
62            SystemTime::now().checked_add(self.default_timeout)
63        {
64            self.ttl_index.insert(expiration_time, key.clone());
65            let value_info = ValueInfo {
66                value,
67                ttl: expiration_time,
68            };
69            self.data.insert(key, value_info);
70        }
71    }
72
73    pub fn remove(&mut self, key: &K) -> Option<V> {
74        match self.data.remove(&key) {
75            Some(info) => {
76                self.ttl_index.remove(&info.ttl);
77                Some(info.value)
78            }
79            None => None,
80        }
81    }
82
83    pub fn gc(&mut self, gc_time: SystemTime) {
84        // Remove the expired entries.
85        let mut active = self.ttl_index.split_off(&gc_time);
86        for key in self.ttl_index.values() {
87            self.data.remove(key);
88        }
89        self.ttl_index.clear();
90        self.ttl_index.append(&mut active);
91    }
92}
cfxcore/pos/mempool/core_mempool/ttl_cache.rs

cfxcore/pos/mempool/core_mempool/
ttl_cache.rs
