longport/quote/
cache.rs

1use std::{
2    collections::HashMap,
3    hash::Hash,
4    time::{Duration, Instant},
5};
6
7use futures_util::Future;
8use tokio::sync::Mutex;
9
10struct Item<V> {
11    deadline: Instant,
12    value: V,
13}
14
15struct Inner<K, V> {
16    timeout: Duration,
17    values: HashMap<K, Item<V>>,
18}
19
20pub(crate) struct CacheWithKey<K, V> {
21    inner: Mutex<Inner<K, V>>,
22}
23
24impl<K, V> CacheWithKey<K, V>
25where
26    K: Eq + Hash + Clone,
27    V: Clone,
28{
29    pub(crate) fn new(timeout: Duration) -> Self {
30        CacheWithKey {
31            inner: Mutex::new(Inner {
32                timeout,
33                values: HashMap::new(),
34            }),
35        }
36    }
37
38    pub(crate) async fn get_or_update<F, Fut, Err>(&self, key: K, f: F) -> Result<V, Err>
39    where
40        F: FnOnce(K) -> Fut,
41        Fut: Future<Output = Result<V, Err>>,
42    {
43        let mut inner = self.inner.lock().await;
44        match inner.values.get(&key) {
45            Some(Item { deadline, value }) if deadline < &Instant::now() => Ok(value.clone()),
46            _ => {
47                let value = f(key.clone()).await?;
48                let deadline = Instant::now() + inner.timeout;
49                inner.values.insert(
50                    key,
51                    Item {
52                        deadline,
53                        value: value.clone(),
54                    },
55                );
56                Ok(value)
57            }
58        }
59    }
60}
61
62pub(crate) struct Cache<V> {
63    inner: CacheWithKey<(), V>,
64}
65
66impl<V> Cache<V>
67where
68    V: Clone,
69{
70    pub(crate) fn new(timeout: Duration) -> Self {
71        Cache {
72            inner: CacheWithKey::new(timeout),
73        }
74    }
75
76    pub(crate) async fn get_or_update<F, Fut, Err>(&self, f: F) -> Result<V, Err>
77    where
78        F: FnOnce() -> Fut,
79        Fut: Future<Output = Result<V, Err>>,
80    {
81        self.inner.get_or_update((), |_| f()).await
82    }
83}