diff --git a/tvix/castore/src/directoryservice/combinators.rs b/tvix/castore/src/directoryservice/combinators.rs index 428314223..5b695b0ff 100644 --- a/tvix/castore/src/directoryservice/combinators.rs +++ b/tvix/castore/src/directoryservice/combinators.rs @@ -178,3 +178,75 @@ impl ServiceBuilder for CacheConfig { })) } } + +#[derive(Clone)] +pub struct Router { + writes: DS1, + reads: DS2, +} + +#[async_trait] +impl DirectoryService for Router +where + DS1: DirectoryService + Clone + 'static, + DS2: DirectoryService + Clone + 'static, +{ + #[instrument(skip(self, digest), fields(directory.digest = %digest))] + async fn get(&self, digest: &B3Digest) -> Result, Error> { + self.reads.get(digest).await + } + + #[instrument(skip_all)] + async fn put(&self, directory: Directory) -> Result { + self.writes.put(directory).await + } + + #[instrument(skip_all, fields(directory.digest = %root_directory_digest))] + fn get_recursive( + &self, + root_directory_digest: &B3Digest, + ) -> BoxStream<'static, Result> { + self.reads.get_recursive(root_directory_digest) + } + + #[instrument(skip_all)] + fn put_multiple_start(&self) -> Box<(dyn DirectoryPutter + 'static)> { + self.writes.put_multiple_start() + } +} + +#[derive(serde::Deserialize, Debug)] +#[serde(deny_unknown_fields)] +pub struct RouterConfig { + writes: String, + reads: String, +} + +impl TryFrom for RouterConfig { + type Error = Box; + fn try_from(_url: url::Url) -> Result { + Err(Error::StorageError( + "Instantiating a CombinedDirectoryService from a url is not supported".into(), + ) + .into()) + } +} + +#[async_trait] +impl ServiceBuilder for RouterConfig { + type Output = dyn DirectoryService; + async fn build<'a>( + &'a self, + _instance_name: &str, + context: &CompositionContext, + ) -> Result, Box> { + let (writes, reads) = futures::join!( + context.resolve(self.writes.clone()), + context.resolve(self.reads.clone()) + ); + Ok(Arc::new(Router { + writes: writes?, + reads: reads?, + })) + } +} diff --git a/tvix/castore/src/directoryservice/mod.rs b/tvix/castore/src/directoryservice/mod.rs index 25162e4de..28b97e852 100644 --- a/tvix/castore/src/directoryservice/mod.rs +++ b/tvix/castore/src/directoryservice/mod.rs @@ -18,7 +18,7 @@ pub mod tests; mod traverse; mod utils; -pub use self::combinators::{Cache, CacheConfig}; +pub use self::combinators::{Cache, CacheConfig, Router, RouterConfig}; pub use self::directory_graph::{DirectoryGraph, ValidatedDirectoryGraph}; pub use self::from_addr::from_addr; pub use self::grpc::{GRPCDirectoryService, GRPCDirectoryServiceConfig}; @@ -137,6 +137,7 @@ pub(crate) fn register_directory_services(reg: &mut Registry) { reg.register::>, super::directoryservice::ObjectStoreDirectoryServiceConfig>("objectstore"); reg.register::>, super::directoryservice::MemoryDirectoryServiceConfig>("memory"); reg.register::>, super::directoryservice::CacheConfig>("cache"); + reg.register::>, super::directoryservice::RouterConfig>("router"); reg.register::>, super::directoryservice::GRPCDirectoryServiceConfig>("grpc"); reg.register::>, super::directoryservice::SledDirectoryServiceConfig>("sled"); reg.register::>, super::directoryservice::RedbDirectoryServiceConfig>("redb");