refactor(tvix): move service addrs into shared clap struct r/8399

Change-Id: I7cab29ecfa1823c2103b4c47b7d784bc31459d55
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12008
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
Autosubmit: yuka <yuka@yuka.dev>

12 files changed, 132 insertions, 149 deletions

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 3d0ba837d8a3..d41c678144d4 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -4878,6 +4878,7 @@ dependencies = [
  "async-compression",
  "bstr",
  "bytes",
+ "clap",
  "criterion",
  "data-encoding",
  "futures",
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index b1821b09d21c..bb2a1dd62979 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -15982,6 +15982,11 @@ rec {
             packageId = "bytes";
           }
           {
+            name = "clap";
+            packageId = "clap";
+            usesDefaultFeatures = false;
+          }
+          {
             name = "data-encoding";
             packageId = "data-encoding";
           }
diff --git a/tvix/cli/src/args.rs b/tvix/cli/src/args.rs
index ebb3c0dc107e..a4d1bcae157e 100644
--- a/tvix/cli/src/args.rs
+++ b/tvix/cli/src/args.rs
@@ -2,6 +2,7 @@ use std::path::PathBuf;
 
 use clap::Parser;
 use tracing::Level;
+use tvix_store::utils::ServiceUrlsMemory;
 
 #[derive(Parser, Clone)]
 pub struct Args {
@@ -58,14 +59,8 @@ pub struct Args {
     #[clap(long)]
     pub strict: bool,
 
-    #[arg(long, env, default_value = "memory://")]
-    pub blob_service_addr: String,
-
-    #[arg(long, env, default_value = "memory://")]
-    pub directory_service_addr: String,
-
-    #[arg(long, env, default_value = "memory://")]
-    pub path_info_service_addr: String,
+    #[clap(flatten)]
+    pub service_addrs: ServiceUrlsMemory,
 
     #[arg(long, env, default_value = "dummy://")]
     pub build_service_addr: String,
diff --git a/tvix/cli/src/lib.rs b/tvix/cli/src/lib.rs
index 070e88e358f2..7686cc2c0fe8 100644
--- a/tvix/cli/src/lib.rs
+++ b/tvix/cli/src/lib.rs
@@ -27,19 +27,9 @@ pub use repl::Repl;
 pub fn init_io_handle(tokio_runtime: &tokio::runtime::Runtime, args: &Args) -> Rc<TvixStoreIO> {
     let (blob_service, directory_service, path_info_service, nar_calculation_service) =
         tokio_runtime
-            .block_on({
-                let blob_service_addr = args.blob_service_addr.clone();
-                let directory_service_addr = args.directory_service_addr.clone();
-                let path_info_service_addr = args.path_info_service_addr.clone();
-                async move {
-                    tvix_store::utils::construct_services(
-                        blob_service_addr,
-                        directory_service_addr,
-                        path_info_service_addr,
-                    )
-                    .await
-                }
-            })
+            .block_on(tvix_store::utils::construct_services(
+                args.service_addrs.clone(),
+            ))
             .expect("unable to setup {blob|directory|pathinfo}service before interpreter setup");
 
     let build_service = tokio_runtime
diff --git a/tvix/glue/Cargo.toml b/tvix/glue/Cargo.toml
index 6fe35c7a5a89..4c72ed676bc6 100644
--- a/tvix/glue/Cargo.toml
+++ b/tvix/glue/Cargo.toml
@@ -31,6 +31,7 @@ sha1 = "0.10.6"
 md-5 = "0.10.6"
 url = "2.4.0"
 walkdir = "2.4.0"
+clap = { version = "4.4.0", default-features = false }
 
 [dependencies.wu-manber]
 git = "https://github.com/tvlfyi/wu-manber.git"
diff --git a/tvix/glue/benches/eval.rs b/tvix/glue/benches/eval.rs
index ee3d554dcbab..3468fa4e431e 100644
--- a/tvix/glue/benches/eval.rs
+++ b/tvix/glue/benches/eval.rs
@@ -1,3 +1,4 @@
+use clap::Parser;
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use lazy_static::lazy_static;
 use std::{env, rc::Rc, sync::Arc, time::Duration};
@@ -11,7 +12,7 @@ use tvix_glue::{
     tvix_io::TvixIO,
     tvix_store_io::TvixStoreIO,
 };
-use tvix_store::utils::construct_services;
+use tvix_store::utils::{construct_services, ServiceUrlsMemory};
 
 #[cfg(not(target_env = "msvc"))]
 #[global_allocator]
@@ -27,7 +28,9 @@ fn interpret(code: &str) {
     // piece of code. b/262
     let (blob_service, directory_service, path_info_service, nar_calculation_service) =
         TOKIO_RUNTIME
-            .block_on(async { construct_services("memory://", "memory://", "memory://").await })
+            .block_on(async {
+                construct_services(ServiceUrlsMemory::parse_from(std::iter::empty::<&str>())).await
+            })
             .unwrap();
 
     // We assemble a complete store in memory.
diff --git a/tvix/glue/src/builtins/mod.rs b/tvix/glue/src/builtins/mod.rs
index e1d1c9c84f25..21685484867a 100644
--- a/tvix/glue/src/builtins/mod.rs
+++ b/tvix/glue/src/builtins/mod.rs
@@ -60,12 +60,13 @@ mod tests {
     use crate::tvix_store_io::TvixStoreIO;
 
     use super::{add_derivation_builtins, add_fetcher_builtins, add_import_builtins};
+    use clap::Parser;
     use nix_compat::store_path::hash_placeholder;
     use rstest::rstest;
     use tempfile::TempDir;
     use tvix_build::buildservice::DummyBuildService;
     use tvix_eval::{EvalIO, EvaluationResult};
-    use tvix_store::utils::construct_services;
+    use tvix_store::utils::{construct_services, ServiceUrlsMemory};
 
     /// evaluates a given nix expression and returns the result.
     /// Takes care of setting up the evaluator so it knows about the
@@ -74,7 +75,9 @@ mod tests {
         // We assemble a complete store in memory.
         let runtime = tokio::runtime::Runtime::new().expect("Failed to build a Tokio runtime");
         let (blob_service, directory_service, path_info_service, nar_calculation_service) = runtime
-            .block_on(async { construct_services("memory://", "memory://", "memory://").await })
+            .block_on(async {
+                construct_services(ServiceUrlsMemory::parse_from(std::iter::empty::<&str>())).await
+            })
             .expect("Failed to construct store services in memory");
 
         let io = Rc::new(TvixStoreIO::new(
diff --git a/tvix/glue/src/tests/mod.rs b/tvix/glue/src/tests/mod.rs
index e9f21c32912e..f68ea5425899 100644
--- a/tvix/glue/src/tests/mod.rs
+++ b/tvix/glue/src/tests/mod.rs
@@ -1,10 +1,11 @@
 use std::{rc::Rc, sync::Arc};
 
+use clap::Parser;
 use pretty_assertions::assert_eq;
 use std::path::PathBuf;
 use tvix_build::buildservice::DummyBuildService;
 use tvix_eval::{EvalIO, Value};
-use tvix_store::utils::construct_services;
+use tvix_store::utils::{construct_services, ServiceUrlsMemory};
 
 use rstest::rstest;
 
@@ -35,7 +36,9 @@ fn eval_test(code_path: PathBuf, expect_success: bool) {
     let tokio_runtime = tokio::runtime::Runtime::new().unwrap();
     let (blob_service, directory_service, path_info_service, nar_calculation_service) =
         tokio_runtime
-            .block_on(async { construct_services("memory://", "memory://", "memory://").await })
+            .block_on(async {
+                construct_services(ServiceUrlsMemory::parse_from(std::iter::empty::<&str>())).await
+            })
             .unwrap();
 
     let tvix_store_io = Rc::new(TvixStoreIO::new(
diff --git a/tvix/glue/src/tvix_store_io.rs b/tvix/glue/src/tvix_store_io.rs
index cfc037e8dada..c25194d50fda 100644
--- a/tvix/glue/src/tvix_store_io.rs
+++ b/tvix/glue/src/tvix_store_io.rs
@@ -627,10 +627,11 @@ mod tests {
     use std::{path::Path, rc::Rc, sync::Arc};
 
     use bstr::ByteSlice;
+    use clap::Parser;
     use tempfile::TempDir;
     use tvix_build::buildservice::DummyBuildService;
     use tvix_eval::{EvalIO, EvaluationResult};
-    use tvix_store::utils::construct_services;
+    use tvix_store::utils::{construct_services, ServiceUrlsMemory};
 
     use super::TvixStoreIO;
     use crate::builtins::{add_derivation_builtins, add_fetcher_builtins, add_import_builtins};
@@ -642,7 +643,10 @@ mod tests {
         let tokio_runtime = tokio::runtime::Runtime::new().unwrap();
         let (blob_service, directory_service, path_info_service, nar_calculation_service) =
             tokio_runtime
-                .block_on(async { construct_services("memory://", "memory://", "memory://").await })
+                .block_on(async {
+                    construct_services(ServiceUrlsMemory::parse_from(std::iter::empty::<&str>()))
+                        .await
+                })
                 .unwrap();
 
         let io = Rc::new(TvixStoreIO::new(
diff --git a/tvix/nar-bridge/src/bin/nar-bridge.rs b/tvix/nar-bridge/src/bin/nar-bridge.rs
index 6823be9e9582..cf312649d147 100644
--- a/tvix/nar-bridge/src/bin/nar-bridge.rs
+++ b/tvix/nar-bridge/src/bin/nar-bridge.rs
@@ -3,19 +3,14 @@ use nar_bridge::AppState;
 use tower::ServiceBuilder;
 use tower_http::trace::{DefaultMakeSpan, TraceLayer};
 use tracing::info;
+use tvix_store::utils::ServiceUrlsGrpc;
 
 /// Expose the Nix HTTP Binary Cache protocol for a tvix-store.
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
 struct Cli {
-    #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-    blob_service_addr: String,
-
-    #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-    directory_service_addr: String,
-
-    #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-    path_info_service_addr: String,
+    #[clap(flatten)]
+    service_addrs: ServiceUrlsGrpc,
 
     /// The priority to announce at the `nix-cache-info` endpoint.
     /// A lower number means it's *more preferred.
@@ -50,12 +45,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
 
     // initialize stores
     let (blob_service, directory_service, path_info_service, _nar_calculation_service) =
-        tvix_store::utils::construct_services(
-            cli.blob_service_addr,
-            cli.directory_service_addr,
-            cli.path_info_service_addr,
-        )
-        .await?;
+        tvix_store::utils::construct_services(cli.service_addrs).await?;
 
     let state = AppState::new(blob_service, directory_service, path_info_service);
 
diff --git a/tvix/store/src/bin/tvix-store.rs b/tvix/store/src/bin/tvix-store.rs
index 99323d2a50a1..fd85b8bd36f3 100644
--- a/tvix/store/src/bin/tvix-store.rs
+++ b/tvix/store/src/bin/tvix-store.rs
@@ -18,6 +18,7 @@ use tvix_castore::import::fs::ingest_path;
 use tvix_store::nar::NarCalculationService;
 use tvix_store::proto::NarInfo;
 use tvix_store::proto::PathInfo;
+use tvix_store::utils::{ServiceUrls, ServiceUrlsGrpc};
 
 use tvix_castore::proto::blob_service_server::BlobServiceServer;
 use tvix_castore::proto::directory_service_server::DirectoryServiceServer;
@@ -67,22 +68,8 @@ enum Commands {
         #[clap(flatten)]
         listen_args: tokio_listener::ListenerAddressLFlag,
 
-        #[arg(
-            long,
-            env,
-            default_value = "objectstore+file:///var/lib/tvix-store/blobs.object_store"
-        )]
-        blob_service_addr: String,
-
-        #[arg(
-            long,
-            env,
-            default_value = "sled:///var/lib/tvix-store/directories.sled"
-        )]
-        directory_service_addr: String,
-
-        #[arg(long, env, default_value = "sled:///var/lib/tvix-store/pathinfo.sled")]
-        path_info_service_addr: String,
+        #[clap(flatten)]
+        service_addrs: ServiceUrls,
 
         /// URL to a PathInfoService that's considered "remote".
         /// If set, the other one is considered "local", and a "cache" for the
@@ -95,26 +82,14 @@ enum Commands {
         #[clap(value_name = "PATH")]
         paths: Vec<PathBuf>,
 
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        blob_service_addr: String,
-
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        directory_service_addr: String,
-
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        path_info_service_addr: String,
+        #[clap(flatten)]
+        service_addrs: ServiceUrlsGrpc,
     },
 
     /// Copies a list of store paths on the system into tvix-store.
     Copy {
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        blob_service_addr: String,
-
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        directory_service_addr: String,
-
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        path_info_service_addr: String,
+        #[clap(flatten)]
+        service_addrs: ServiceUrlsGrpc,
 
         /// A path pointing to a JSON file produced by the Nix
         /// `__structuredAttrs` containing reference graph information provided
@@ -134,14 +109,8 @@ enum Commands {
         #[clap(value_name = "PATH")]
         dest: PathBuf,
 
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        blob_service_addr: String,
-
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        directory_service_addr: String,
-
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        path_info_service_addr: String,
+        #[clap(flatten)]
+        service_addrs: ServiceUrlsGrpc,
 
         /// Number of FUSE threads to spawn.
         #[arg(long, env, default_value_t = default_threads())]
@@ -170,14 +139,8 @@ enum Commands {
         #[clap(value_name = "PATH")]
         socket: PathBuf,
 
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        blob_service_addr: String,
-
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        directory_service_addr: String,
-
-        #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
-        path_info_service_addr: String,
+        #[clap(flatten)]
+        service_addrs: ServiceUrlsGrpc,
 
         /// Whether to list elements at the root of the mount point.
         /// This is useful if your PathInfoService doesn't provide an
@@ -203,17 +166,11 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error + Send + Sync
     match cli.command {
         Commands::Daemon {
             listen_args,
-            blob_service_addr,
-            directory_service_addr,
-            path_info_service_addr,
+            service_addrs,
             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,
-            )?;
+            let mut configs = tvix_store::utils::addrs_to_configs(service_addrs)?;
 
             // if remote_path_info_service_addr has been specified,
             // update path_info_service to point to a cache combining the two.
@@ -295,18 +252,11 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error + Send + Sync
         }
         Commands::Import {
             paths,
-            blob_service_addr,
-            directory_service_addr,
-            path_info_service_addr,
+            service_addrs,
         } => {
             // FUTUREWORK: allow flat for single files?
             let (blob_service, directory_service, path_info_service, nar_calculation_service) =
-                tvix_store::utils::construct_services(
-                    blob_service_addr,
-                    directory_service_addr,
-                    path_info_service_addr,
-                )
-                .await?;
+                tvix_store::utils::construct_services(service_addrs).await?;
 
             // Arc NarCalculationService, as we clone it .
             let nar_calculation_service: Arc<dyn NarCalculationService> =
@@ -345,18 +295,11 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error + Send + Sync
             try_join_all(tasks).await?;
         }
         Commands::Copy {
-            blob_service_addr,
-            directory_service_addr,
-            path_info_service_addr,
+            service_addrs,
             reference_graph_path,
         } => {
             let (blob_service, directory_service, path_info_service, _nar_calculation_service) =
-                tvix_store::utils::construct_services(
-                    blob_service_addr,
-                    directory_service_addr,
-                    path_info_service_addr,
-                )
-                .await?;
+                tvix_store::utils::construct_services(service_addrs).await?;
 
             // Parse the file at reference_graph_path.
             let reference_graph_json = tokio::fs::read(&reference_graph_path).await?;
@@ -457,21 +400,14 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error + Send + Sync
         #[cfg(feature = "fuse")]
         Commands::Mount {
             dest,
-            blob_service_addr,
-            directory_service_addr,
-            path_info_service_addr,
+            service_addrs,
             list_root,
             threads,
             allow_other,
             show_xattr,
         } => {
             let (blob_service, directory_service, path_info_service, _nar_calculation_service) =
-                tvix_store::utils::construct_services(
-                    blob_service_addr,
-                    directory_service_addr,
-                    path_info_service_addr,
-                )
-                .await?;
+                tvix_store::utils::construct_services(service_addrs).await?;
 
             let fuse_daemon = tokio::task::spawn_blocking(move || {
                 let fs = make_fs(
@@ -507,19 +443,12 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error + Send + Sync
         #[cfg(feature = "virtiofs")]
         Commands::VirtioFs {
             socket,
-            blob_service_addr,
-            directory_service_addr,
-            path_info_service_addr,
+            service_addrs,
             list_root,
             show_xattr,
         } => {
             let (blob_service, directory_service, path_info_service, _nar_calculation_service) =
-                tvix_store::utils::construct_services(
-                    blob_service_addr,
-                    directory_service_addr,
-                    path_info_service_addr,
-                )
-                .await?;
+                tvix_store::utils::construct_services(service_addrs).await?;
 
             tokio::task::spawn_blocking(move || {
                 let fs = make_fs(
diff --git a/tvix/store/src/utils.rs b/tvix/store/src/utils.rs
index a09786386eba..8f97651a03a5 100644
--- a/tvix/store/src/utils.rs
+++ b/tvix/store/src/utils.rs
@@ -29,16 +29,81 @@ pub struct CompositionConfigs {
     >,
 }
 
+#[derive(clap::Parser, Clone)]
+pub struct ServiceUrls {
+    #[arg(
+        long,
+        env,
+        default_value = "objectstore+file:///var/lib/tvix-store/blobs.object_store"
+    )]
+    blob_service_addr: String,
+
+    #[arg(
+        long,
+        env,
+        default_value = "sled:///var/lib/tvix-store/directories.sled"
+    )]
+    directory_service_addr: String,
+
+    #[arg(long, env, default_value = "sled:///var/lib/tvix-store/pathinfo.sled")]
+    path_info_service_addr: String,
+}
+
+/// like ServiceUrls, but with different clap defaults
+#[derive(clap::Parser, Clone)]
+pub struct ServiceUrlsGrpc {
+    #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
+    blob_service_addr: String,
+
+    #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
+    directory_service_addr: String,
+
+    #[arg(long, env, default_value = "grpc+http://[::1]:8000")]
+    path_info_service_addr: String,
+}
+
+/// like ServiceUrls, but with different clap defaults
+#[derive(clap::Parser, Clone)]
+pub struct ServiceUrlsMemory {
+    #[arg(long, env, default_value = "memory://")]
+    blob_service_addr: String,
+
+    #[arg(long, env, default_value = "memory://")]
+    directory_service_addr: String,
+
+    #[arg(long, env, default_value = "memory://")]
+    path_info_service_addr: String,
+}
+
+impl From<ServiceUrlsGrpc> for ServiceUrls {
+    fn from(urls: ServiceUrlsGrpc) -> ServiceUrls {
+        ServiceUrls {
+            blob_service_addr: urls.blob_service_addr,
+            directory_service_addr: urls.directory_service_addr,
+            path_info_service_addr: urls.path_info_service_addr,
+        }
+    }
+}
+
+impl From<ServiceUrlsMemory> for ServiceUrls {
+    fn from(urls: ServiceUrlsMemory) -> ServiceUrls {
+        ServiceUrls {
+            blob_service_addr: urls.blob_service_addr,
+            directory_service_addr: urls.directory_service_addr,
+            path_info_service_addr: urls.path_info_service_addr,
+        }
+    }
+}
+
 pub fn addrs_to_configs(
-    blob_service_addr: impl AsRef<str>,
-    directory_service_addr: impl AsRef<str>,
-    path_info_service_addr: impl AsRef<str>,
+    urls: impl Into<ServiceUrls>,
 ) -> Result<CompositionConfigs, Box<dyn std::error::Error + Send + Sync>> {
+    let urls: ServiceUrls = urls.into();
     let mut configs: CompositionConfigs = Default::default();
 
-    let blob_service_url = Url::parse(blob_service_addr.as_ref())?;
-    let directory_service_url = Url::parse(directory_service_addr.as_ref())?;
-    let path_info_service_url = Url::parse(path_info_service_addr.as_ref())?;
+    let blob_service_url = Url::parse(&urls.blob_service_addr)?;
+    let directory_service_url = Url::parse(&urls.directory_service_addr)?;
+    let path_info_service_url = Url::parse(&urls.path_info_service_addr)?;
 
     configs.blobservices.insert(
         "default".into(),
@@ -58,9 +123,7 @@ pub fn addrs_to_configs(
 
 /// Construct the store handles from their addrs.
 pub async fn construct_services(
-    blob_service_addr: impl AsRef<str>,
-    directory_service_addr: impl AsRef<str>,
-    path_info_service_addr: impl AsRef<str>,
+    urls: impl Into<ServiceUrls>,
 ) -> Result<
     (
         Arc<dyn BlobService>,
@@ -70,11 +133,7 @@ pub async fn construct_services(
     ),
     Box<dyn std::error::Error + Send + Sync>,
 > {
-    let configs = addrs_to_configs(
-        blob_service_addr,
-        directory_service_addr,
-        path_info_service_addr,
-    )?;
+    let configs = addrs_to_configs(urls)?;
     construct_services_from_configs(configs).await
 }