add tick function for memory storage
This commit is contained in:
parent
319a407ae3
commit
d1bee14213
1 changed files with 22 additions and 5 deletions
|
@ -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());
|
||||
}
|
||||
|
||||
|
|
Loading…
Reference in a new issue