84 lines
1.9 KiB
Rust
84 lines
1.9 KiB
Rust
use std::{collections::HashMap, pin::Pin};
|
|
|
|
use anyhow::Ok;
|
|
use async_trait::async_trait;
|
|
use tokio_stream::{Stream, StreamExt};
|
|
use tokio_util::bytes::Bytes;
|
|
|
|
use crate::config::MemoryStorageConfig;
|
|
|
|
use super::Storage;
|
|
|
|
pub struct MemoryStorage {
|
|
config: MemoryStorageConfig,
|
|
items: HashMap<String, Vec<u8>>,
|
|
}
|
|
|
|
impl MemoryStorage {
|
|
pub fn new(config: MemoryStorageConfig) -> Self {
|
|
Self {
|
|
config,
|
|
items: HashMap::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Storage for MemoryStorage {
|
|
async fn init(&mut self) -> anyhow::Result<()> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn eligible(&self, _src: String) -> bool {
|
|
true
|
|
}
|
|
|
|
async fn tick(&mut self) -> anyhow::Result<()> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn delete(&mut self, key: String) -> anyhow::Result<()> {
|
|
self.items
|
|
.remove(&key)
|
|
.ok_or_else(|| anyhow::anyhow!(""))
|
|
.map(|_| ())
|
|
}
|
|
|
|
async fn retrieve(
|
|
&self,
|
|
key: String,
|
|
) -> 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 stream = futures::stream::iter([Ok(bytes)]);
|
|
|
|
Some(Box::pin(stream))
|
|
}
|
|
None => None,
|
|
}
|
|
}
|
|
|
|
async fn save(
|
|
&mut self,
|
|
key: String,
|
|
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());
|
|
}
|
|
|
|
println!("save {}", key);
|
|
|
|
while let Some(chunk) = stream.next().await {
|
|
self.items
|
|
.get_mut(&key)
|
|
.unwrap()
|
|
.append(&mut chunk?.to_vec());
|
|
}
|
|
|
|
println!("{} saved", key);
|
|
|
|
Ok(())
|
|
}
|
|
}
|