94 lines
2.4 KiB
Rust
94 lines
2.4 KiB
Rust
pub mod disk;
|
|
pub mod memory;
|
|
pub mod mixed;
|
|
|
|
use std::pin::Pin;
|
|
|
|
use async_trait::async_trait;
|
|
use tokio_stream::Stream;
|
|
use tokio_util::bytes::Bytes;
|
|
|
|
use crate::{
|
|
config::{StorageConfig, StorageStrategyConfig},
|
|
crypt,
|
|
};
|
|
|
|
use self::{disk::DiskStorage, memory::MemoryStorage, mixed::MixedStorage};
|
|
|
|
#[async_trait]
|
|
pub trait Storage {
|
|
async fn init(&mut self) -> anyhow::Result<()>;
|
|
async fn eligible(&self, src: String) -> bool;
|
|
async fn delete(&mut self, key: String) -> anyhow::Result<()>;
|
|
async fn retrieve(
|
|
&self,
|
|
key: String,
|
|
) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>>;
|
|
async fn save(
|
|
&mut self,
|
|
key: String,
|
|
stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
|
|
) -> anyhow::Result<()>;
|
|
}
|
|
|
|
pub struct StoragePool {
|
|
storages: Vec<Box<dyn Storage + Sync + Send>>,
|
|
}
|
|
|
|
impl StoragePool {
|
|
pub fn from_config(items: Vec<StorageConfig>) -> Self {
|
|
Self {
|
|
storages: items.into_iter().map(create_storage).collect(),
|
|
}
|
|
}
|
|
|
|
pub async fn init(&mut self) -> anyhow::Result<()> {
|
|
for item in &mut self.storages {
|
|
item.init().await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn retrieve(
|
|
&self,
|
|
src: String,
|
|
) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>> {
|
|
let mut stream = None;
|
|
let key = crypt::compute_key(src);
|
|
|
|
for item in &self.storages {
|
|
if let Some(storage_stream) = item.retrieve(key.clone()).await {
|
|
stream = Some(storage_stream);
|
|
break;
|
|
}
|
|
}
|
|
|
|
stream
|
|
}
|
|
|
|
pub async fn save(
|
|
&mut self,
|
|
src: String,
|
|
stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
|
|
) -> anyhow::Result<()> {
|
|
let key = crypt::compute_key(src);
|
|
|
|
for item in self.storages.iter_mut() {
|
|
if item.eligible(key.clone()).await {
|
|
item.save(key, stream).await?;
|
|
break;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
fn create_storage(item: StorageConfig) -> Box<dyn Storage + Sync + Send> {
|
|
match item.strategy {
|
|
StorageStrategyConfig::Memory(config) => Box::new(MemoryStorage::new(config)),
|
|
StorageStrategyConfig::Disk(config) => Box::new(DiskStorage::new(config)),
|
|
StorageStrategyConfig::Mixed(config) => Box::new(MixedStorage::new(config)),
|
|
}
|
|
}
|