use std::sync::Arc;
use futures::stream::BoxStream;
use futures::StreamExt;
use futures::TryFutureExt;
use futures::TryStreamExt;
use tonic::async_trait;
use tracing::{instrument, trace};
use super::{Directory, DirectoryGraph, DirectoryService, RootToLeavesValidator, SimplePutter};
use crate::composition::{CompositionContext, ServiceBuilder};
use crate::directoryservice::DirectoryPutter;
use crate::B3Digest;
use crate::Error;
#[derive(Clone)]
pub struct Cache<DS1, DS2> {
instance_name: String,
near: DS1,
far: DS2,
}
impl<DS1, DS2> Cache<DS1, DS2> {
pub fn new(instance_name: String, near: DS1, far: DS2) -> Self {
Self {
instance_name,
near,
far,
}
}
}
#[async_trait]
impl<DS1, DS2> DirectoryService for Cache<DS1, DS2>
where
DS1: DirectoryService + Clone + 'static,
DS2: DirectoryService + Clone + 'static,
{
#[instrument(skip(self, digest), fields(directory.digest = %digest, instance_name = %self.instance_name))]
async fn get(&self, digest: &B3Digest) -> Result<Option<Directory>, Error> {
match self.near.get(digest).await? {
Some(directory) => {
trace!("serving from cache");
Ok(Some(directory))
}
None => {
trace!("not found in near, asking remote…");
let mut copy = DirectoryGraph::with_order(
RootToLeavesValidator::new_with_root_digest(digest.clone()),
);
let mut stream = self.far.get_recursive(digest);
let root = stream.try_next().await?;
if let Some(root) = root.clone() {
copy.add(root)
.map_err(|e| Error::StorageError(e.to_string()))?;
}
while let Some(dir) = stream.try_next().await? {
copy.add(dir)
.map_err(|e| Error::StorageError(e.to_string()))?;
}
let copy = copy
.validate()
.map_err(|e| Error::StorageError(e.to_string()))?;
let mut put = self.near.put_multiple_start();
for dir in copy.drain_leaves_to_root() {
put.put(dir).await?;
}
put.close().await?;
Ok(root)
}
}
}
#[instrument(skip_all, fields(instance_name = %self.instance_name))]
async fn put(&self, _directory: Directory) -> Result<B3Digest, Error> {
Err(Error::StorageError("unimplemented".to_string()))
}
#[instrument(skip_all, fields(directory.digest = %root_directory_digest, instance_name = %self.instance_name))]
fn get_recursive(
&self,
root_directory_digest: &B3Digest,
) -> BoxStream<'static, Result<Directory, Error>> {
let near = self.near.clone();
let far = self.far.clone();
let digest = root_directory_digest.clone();
Box::pin(
(async move {
let mut stream = near.get_recursive(&digest);
match stream.try_next().await? {
Some(first) => {
trace!("serving from cache");
Ok(futures::stream::once(async { Ok(first) })
.chain(stream)
.left_stream())
}
None => {
trace!("not found in near, asking remote…");
let mut copy_for_near = DirectoryGraph::with_order(
RootToLeavesValidator::new_with_root_digest(digest.clone()),
);
let mut copy_for_client = vec![];
let mut stream = far.get_recursive(&digest);
while let Some(dir) = stream.try_next().await? {
copy_for_near
.add(dir.clone())
.map_err(|e| Error::StorageError(e.to_string()))?;
copy_for_client.push(dir);
}
let copy_for_near = copy_for_near
.validate()
.map_err(|e| Error::StorageError(e.to_string()))?;
let mut put = near.put_multiple_start();
for dir in copy_for_near.drain_leaves_to_root() {
put.put(dir).await?;
}
put.close().await?;
Ok(futures::stream::iter(copy_for_client.into_iter().map(Ok))
.right_stream())
}
}
})
.try_flatten_stream(),
)
}
#[instrument(skip_all)]
fn put_multiple_start(&self) -> Box<(dyn DirectoryPutter + 'static)> {
Box::new(SimplePutter::new((*self).clone()))
}
}
#[derive(serde::Deserialize, Debug)]
#[serde(deny_unknown_fields)]
pub struct CacheConfig {
near: String,
far: String,
}
impl TryFrom<url::Url> for CacheConfig {
type Error = Box<dyn std::error::Error + Send + Sync>;
fn try_from(url: url::Url) -> Result<Self, Self::Error> {
if url.has_host() || !url.path().is_empty() {
return Err(Error::StorageError("invalid url".to_string()).into());
}
Ok(serde_qs::from_str(url.query().unwrap_or_default())?)
}
}
#[async_trait]
impl ServiceBuilder for CacheConfig {
type Output = dyn DirectoryService;
async fn build<'a>(
&'a self,
instance_name: &str,
context: &CompositionContext,
) -> Result<Arc<dyn DirectoryService>, Box<dyn std::error::Error + Send + Sync + 'static>> {
let (near, far) = futures::join!(
context.resolve::<Self::Output>(self.near.clone()),
context.resolve::<Self::Output>(self.far.clone())
);
Ok(Arc::new(Cache {
instance_name: instance_name.to_string(),
near: near?,
far: far?,
}))
}
}