add tick function for memory storage

This commit is contained in:
qpismont 2024-07-14 22:22:10 +02:00
parent 319a407ae3
commit d1bee14213

View file

@ -1,7 +1,8 @@
use std::{collections::HashMap, pin::Pin};
use anyhow::Ok;
use async_trait::async_trait;
use std::time::{Duration, SystemTime};
use std::{collections::HashMap, pin::Pin};
use tokio::time::Instant;
use tokio_stream::{Stream, StreamExt};
use tokio_util::bytes::Bytes;
@ -9,9 +10,14 @@ use crate::config::MemoryStorageConfig;
use super::Storage;
struct MemoryStorageItem {
data: Vec<u8>,
added_at: Instant,
}
pub struct MemoryStorage {
config: MemoryStorageConfig,
items: HashMap<String, Vec<u8>>,
items: HashMap<String, MemoryStorageItem>,
}
impl MemoryStorage {
@ -34,6 +40,10 @@ impl Storage for MemoryStorage {
}
async fn tick(&mut self) -> anyhow::Result<()> {
let ttl = self.config.ttl;
self.items
.retain(|_, elt| elt.added_at.elapsed().as_secs() < ttl as u64);
Ok(())
}
@ -50,7 +60,7 @@ impl Storage for MemoryStorage {
) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>> {
match self.items.get(&key) {
Some(item) => {
let bytes = tokio_util::bytes::Bytes::from(item.clone());
let bytes = tokio_util::bytes::Bytes::from(item.data.clone());
let stream = futures::stream::iter([Ok(bytes)]);
Some(Box::pin(stream))
@ -65,7 +75,13 @@ impl Storage for MemoryStorage {
mut stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
) -> anyhow::Result<()> {
if !self.items.contains_key(&key) {
self.items.insert(key.clone(), Vec::new());
self.items.insert(
key.clone(),
MemoryStorageItem {
data: Vec::new(),
added_at: Instant::now(),
},
);
}
println!("save {}", key);
@ -74,6 +90,7 @@ impl Storage for MemoryStorage {
self.items
.get_mut(&key)
.unwrap()
.data
.append(&mut chunk?.to_vec());
}