feat(tvix/build): add from_addr method r/7386

This allows constructing a BuildService from a URI, similar to how it's
done in tvix-[ca]store.

Change-Id: Ib962b329535c6c7e378ab7ac7f4dd254366497b3
Reviewed-on: https://cl.tvl.fyi/c/depot/+/10630
Tested-by: BuildkiteCI
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Autosubmit: flokli <flokli@flokli.de>

5 files changed, 99 insertions, 0 deletions

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 95cc50a4072b..a37c04f974d6 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -3288,11 +3288,13 @@ dependencies = [
  "prost-build",
  "test-case",
  "thiserror",
+ "tokio",
  "tonic 0.10.2",
  "tonic-build",
  "tonic-reflection",
  "tracing",
  "tvix-castore",
+ "url",
 ]
 
 [[package]]
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index 71f456290d1d..d1ffbf326470 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -10238,6 +10238,10 @@ rec {
             packageId = "thiserror";
           }
           {
+            name = "tokio";
+            packageId = "tokio";
+          }
+          {
             name = "tonic";
             packageId = "tonic 0.10.2";
             features = [ "tls" "tls-roots" ];
@@ -10255,6 +10259,10 @@ rec {
             name = "tvix-castore";
             packageId = "tvix-castore";
           }
+          {
+            name = "url";
+            packageId = "url";
+          }
         ];
         buildDependencies = [
           {
diff --git a/tvix/build/Cargo.toml b/tvix/build/Cargo.toml
index 52dcf0abec0c..08102a66be50 100644
--- a/tvix/build/Cargo.toml
+++ b/tvix/build/Cargo.toml
@@ -8,9 +8,11 @@ bytes = "1.4.0"
 itertools = "0.12.0"
 prost = "0.12.1"
 thiserror = "1.0.56"
+tokio = { version = "1.32.0" }
 tonic = { version = "0.10.2", features = ["tls", "tls-roots"] }
 tvix-castore = { path = "../castore" }
 tracing = "0.1.37"
+url = "2.4.0"
 
 [dependencies.tonic-reflection]
 optional = true
diff --git a/tvix/build/src/buildservice/from_addr.rs b/tvix/build/src/buildservice/from_addr.rs
new file mode 100644
index 000000000000..ee2b4e50b4da
--- /dev/null
+++ b/tvix/build/src/buildservice/from_addr.rs
@@ -0,0 +1,85 @@
+use super::{grpc::GRPCBuildService, BuildService, DummyBuildService};
+use tvix_castore::{blobservice::BlobService, directoryservice::DirectoryService};
+use url::Url;
+
+/// Constructs a new instance of a [BuildService] from an URI.
+///
+/// The following schemes are supported by the following services:
+/// - `dummy://` ([DummyBuildService])
+/// - `grpc+*://` ([GRPCBuildService])
+///
+/// As some of these [BuildService] need to talk to a [BlobService] and
+/// [DirectoryService], these also need to be passed in.
+pub async fn from_addr<BS, DS>(
+    uri: &str,
+    _blob_service: BS,
+    _directory_service: DS,
+) -> std::io::Result<Box<dyn BuildService>>
+where
+    BS: AsRef<dyn BlobService> + Send + Sync + Clone + 'static,
+    DS: AsRef<dyn DirectoryService> + Send + Sync + Clone + 'static,
+{
+    let url = Url::parse(uri)
+        .map_err(|e| std::io::Error::other(format!("unable to parse url: {}", e)))?;
+
+    Ok(match url.scheme() {
+        // dummy doesn't care about parameters.
+        "dummy" => Box::<DummyBuildService>::default(),
+        scheme => {
+            if scheme.starts_with("grpc+") {
+                let client = crate::proto::build_service_client::BuildServiceClient::new(
+                    tvix_castore::tonic::channel_from_url(&url)
+                        .await
+                        .map_err(std::io::Error::other)?,
+                );
+                // FUTUREWORK: also allow responding to {blob,directory}_service
+                // requests from the remote BuildService?
+                Box::new(GRPCBuildService::from_client(client))
+            } else {
+                Err(std::io::Error::other(format!(
+                    "unknown scheme: {}",
+                    url.scheme()
+                )))?
+            }
+        }
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use super::from_addr;
+    use test_case::test_case;
+    use tvix_castore::utils::{gen_blob_service, gen_directory_service};
+
+    /// This uses an unsupported scheme.
+    #[test_case("http://foo.example/test", false; "unsupported scheme")]
+    /// This configures dummy
+    #[test_case("dummy://", true; "valid dummy")]
+    /// Correct scheme to connect to a unix socket.
+    #[test_case("grpc+unix:///path/to/somewhere", true; "grpc valid unix socket")]
+    /// Correct scheme for unix socket, but setting a host too, which is invalid.
+    #[test_case("grpc+unix://host.example/path/to/somewhere", false; "grpc invalid unix socket and host")]
+    /// Correct scheme to connect to localhost, with port 12345
+    #[test_case("grpc+http://[::1]:12345", true; "grpc valid IPv6 localhost port 12345")]
+    /// Correct scheme to connect to localhost over http, without specifying a port.
+    #[test_case("grpc+http://localhost", true; "grpc valid http host without port")]
+    /// Correct scheme to connect to localhost over http, without specifying a port.
+    #[test_case("grpc+https://localhost", true; "grpc valid https host without port")]
+    /// Correct scheme to connect to localhost over http, but with additional path, which is invalid.
+    #[test_case("grpc+http://localhost/some-path", false; "grpc valid invalid host and path")]
+    #[tokio::test]
+    async fn test_from_addr(uri_str: &str, is_ok: bool) {
+        assert_eq!(
+            from_addr(
+                uri_str,
+                Arc::from(gen_blob_service()),
+                Arc::from(gen_directory_service())
+            )
+            .await
+            .is_ok(),
+            is_ok
+        )
+    }
+}
diff --git a/tvix/build/src/buildservice/mod.rs b/tvix/build/src/buildservice/mod.rs
index cbee221a5e23..a61d782919b9 100644
--- a/tvix/build/src/buildservice/mod.rs
+++ b/tvix/build/src/buildservice/mod.rs
@@ -3,9 +3,11 @@ use tonic::async_trait;
 use crate::proto::{Build, BuildRequest};
 
 mod dummy;
+mod from_addr;
 mod grpc;
 
 pub use dummy::DummyBuildService;
+pub use from_addr::from_addr;
 
 #[async_trait]
 pub trait BuildService: Send + Sync {