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 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());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue