1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
use std::{
    collections::HashMap,
    sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
    time::{Duration, Instant},
};

use sessions_core::{anyhow, async_trait, Data, Result, Storage};

#[derive(Clone, Debug)]
struct State(Instant, Data);

impl State {
    fn new(i: Instant, d: Data) -> Self {
        Self(i, d)
    }
}

#[derive(Clone, Debug)]
pub struct MemoryStorage {
    inner: Arc<RwLock<HashMap<String, State>>>,
}

impl MemoryStorage {
    pub fn new() -> Self {
        Self { inner: Arc::default() }
    }

    fn read(&self) -> Result<RwLockReadGuard<'_, HashMap<String, State>>> {
        self.inner.read().map_err(|e| anyhow!(e.to_string()))
    }

    fn write(&self) -> Result<RwLockWriteGuard<'_, HashMap<String, State>>> {
        self.inner.write().map_err(|e| anyhow!(e.to_string()))
    }
}

#[async_trait]
impl Storage for MemoryStorage {
    async fn get(&self, key: &str) -> Result<Option<Data>> {
        let state = self.read()?.get(key).cloned();
        if let Some(State(time, data)) = state {
            if time >= Instant::now() {
                return Ok(Some(data));
            } else {
                self.remove(key).await?;
            }
        }

        Ok(None)
    }

    async fn set(&self, key: &str, val: Data, exp: Duration) -> Result<()> {
        self.write()?
            .insert(key.to_string(), State::new(Instant::now() + exp, val));
        Ok(())
    }

    async fn remove(&self, key: &str) -> Result<()> {
        self.write()?.remove(key);
        Ok(())
    }

    async fn reset(&self) -> Result<()> {
        Ok(self.write()?.clear())
    }
}