From 71c10cd7d65554cc2f0fce824fe9f123310c77d7 Mon Sep 17 00:00:00 2001
From: sinavir <sinavir@sinavir.fr>
Date: Sun, 21 Jul 2024 11:51:46 +0200
Subject: [PATCH] feat(nar-bridge): Use store composition

---
 tvix/nar-bridge/src/bin/nar-bridge.rs | 12 +++++--
 tvix/store/src/bin/tvix-store.rs      | 39 ++++------------------
 tvix/store/src/utils.rs               | 48 +++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 35 deletions(-)

diff --git a/tvix/nar-bridge/src/bin/nar-bridge.rs b/tvix/nar-bridge/src/bin/nar-bridge.rs
index 6823be9e9..da967c3ee 100644
--- a/tvix/nar-bridge/src/bin/nar-bridge.rs
+++ b/tvix/nar-bridge/src/bin/nar-bridge.rs
@@ -17,6 +17,12 @@ struct Cli {
     #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
     path_info_service_addr: String,
 
+    /// URL to a PathInfoService that's considered "remote".
+    /// If set, the other one is considered "local", and a "cache" for the
+    /// "remote" one.
+    #[arg(long, env)]
+    remote_path_info_service_addr: Option<String>,
+
     /// The priority to announce at the `nix-cache-info` endpoint.
     /// A lower number means it's *more preferred.
     #[arg(long, env, default_value_t = 39)]
@@ -48,12 +54,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         builder.build()?
     };
 
-    // initialize stores
-    let (blob_service, directory_service, path_info_service, _nar_calculation_service) =
-        tvix_store::utils::construct_services(
+    let (blob_service, directory_service, path_info_service, _) =
+        tvix_store::utils::initialize_stores(
             cli.blob_service_addr,
             cli.directory_service_addr,
             cli.path_info_service_addr,
+            cli.remote_path_info_service_addr,
         )
         .await?;
 
diff --git a/tvix/store/src/bin/tvix-store.rs b/tvix/store/src/bin/tvix-store.rs
index 99323d2a5..1a252cd4f 100644
--- a/tvix/store/src/bin/tvix-store.rs
+++ b/tvix/store/src/bin/tvix-store.rs
@@ -208,39 +208,14 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error + Send + Sync
             path_info_service_addr,
             remote_path_info_service_addr,
         } => {
-            // initialize stores
-            let mut configs = tvix_store::utils::addrs_to_configs(
-                blob_service_addr,
-                directory_service_addr,
-                path_info_service_addr,
-            )?;
-
-            // if remote_path_info_service_addr has been specified,
-            // update path_info_service to point to a cache combining the two.
-            if let Some(addr) = remote_path_info_service_addr {
-                use tvix_store::composition::{with_registry, DeserializeWithRegistry, REG};
-                use tvix_store::pathinfoservice::CachePathInfoServiceConfig;
-
-                let remote_url = url::Url::parse(&addr)?;
-                let remote_config = with_registry(&REG, || remote_url.try_into())?;
-
-                let local = configs.pathinfoservices.insert(
-                    "default".into(),
-                    DeserializeWithRegistry(Box::new(CachePathInfoServiceConfig {
-                        near: "local".into(),
-                        far: "remote".into(),
-                    })),
-                );
-                configs
-                    .pathinfoservices
-                    .insert("local".into(), local.unwrap());
-                configs
-                    .pathinfoservices
-                    .insert("remote".into(), remote_config);
-            }
-
             let (blob_service, directory_service, path_info_service, nar_calculation_service) =
-                tvix_store::utils::construct_services_from_configs(configs).await?;
+                tvix_store::utils::initialize_stores(
+                    blob_service_addr,
+                    directory_service_addr,
+                    path_info_service_addr,
+                    remote_path_info_service_addr,
+                )
+                .await?;
 
             let mut server = Server::builder().layer(
                 ServiceBuilder::new()
diff --git a/tvix/store/src/utils.rs b/tvix/store/src/utils.rs
index a09786386..ab26c75e8 100644
--- a/tvix/store/src/utils.rs
+++ b/tvix/store/src/utils.rs
@@ -29,6 +29,54 @@ pub struct CompositionConfigs {
     >,
 }
 
+pub async fn initialize_stores(
+    blob_service_addr: impl AsRef<str>,
+    directory_service_addr: impl AsRef<str>,
+    path_info_service_addr: impl AsRef<str>,
+    remote_path_info_service_addr: Option<impl AsRef<str>>,
+) -> Result<
+    (
+        Arc<dyn BlobService>,
+        Arc<dyn DirectoryService>,
+        Arc<dyn PathInfoService>,
+        Box<dyn NarCalculationService>,
+    ),
+    Box<dyn std::error::Error + Send + Sync>,
+> {
+    // initialize stores
+    let mut configs = crate::utils::addrs_to_configs(
+        blob_service_addr,
+        directory_service_addr,
+        path_info_service_addr,
+    )?;
+
+    // if remote_path_info_service_addr has been specified,
+    // update path_info_service to point to a cache combining the two.
+    if let Some(addr) = remote_path_info_service_addr {
+        use crate::composition::{with_registry, DeserializeWithRegistry, REG};
+        use crate::pathinfoservice::CachePathInfoServiceConfig;
+
+        let remote_url = url::Url::parse(addr.as_ref())?;
+        let remote_config = with_registry(&REG, || remote_url.try_into())?;
+
+        let local = configs.pathinfoservices.insert(
+            "default".into(),
+            DeserializeWithRegistry(Box::new(CachePathInfoServiceConfig {
+                near: "local".into(),
+                far: "remote".into(),
+            })),
+        );
+        configs
+            .pathinfoservices
+            .insert("local".into(), local.unwrap());
+        configs
+            .pathinfoservices
+            .insert("remote".into(), remote_config);
+    }
+
+    Ok(crate::utils::construct_services_from_configs(configs).await?)
+}
+
 pub fn addrs_to_configs(
     blob_service_addr: impl AsRef<str>,
     directory_service_addr: impl AsRef<str>,