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 anyhow::Ok;
use async_trait::async_trait; 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_stream::{Stream, StreamExt};
use tokio_util::bytes::Bytes; use tokio_util::bytes::Bytes;
@ -9,9 +10,14 @@ use crate::config::MemoryStorageConfig;
use super::Storage; use super::Storage;
struct MemoryStorageItem {
data: Vec<u8>,
added_at: Instant,
}
pub struct MemoryStorage { pub struct MemoryStorage {
config: MemoryStorageConfig, config: MemoryStorageConfig,
items: HashMap<String, Vec<u8>>, items: HashMap<String, MemoryStorageItem>,
} }
impl MemoryStorage { impl MemoryStorage {
@ -34,6 +40,10 @@ impl Storage for MemoryStorage {
} }
async fn tick(&mut self) -> anyhow::Result<()> { 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(()) Ok(())
} }
@ -50,7 +60,7 @@ impl Storage for MemoryStorage {
) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>> { ) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>> {
match self.items.get(&key) { match self.items.get(&key) {
Some(item) => { 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)]); let stream = futures::stream::iter([Ok(bytes)]);
Some(Box::pin(stream)) Some(Box::pin(stream))
@ -65,7 +75,13 @@ impl Storage for MemoryStorage {
mut stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>, mut stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
) -> anyhow::Result<()> { ) -> anyhow::Result<()> {
if !self.items.contains_key(&key) { 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); println!("save {}", key);
@ -74,6 +90,7 @@ impl Storage for MemoryStorage {
self.items self.items
.get_mut(&key) .get_mut(&key)
.unwrap() .unwrap()
.data
.append(&mut chunk?.to_vec()); .append(&mut chunk?.to_vec());
} }