refactor(tvix/*store): helper for channel creation from url r/6802

This moves the repetitive code to parse a URL and create a channel
connected to it into `tvix_castore::channel::from_url`.

Part of b/308

Change-Id: Idd342cd71cad5e78a9b258b38c1b227993e75310
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9707
Tested-by: BuildkiteCI
Reviewed-by: Connor Brewster <cbrewster@hey.com>

5 files changed, 84 insertions, 88 deletions

diff --git a/tvix/castore/src/blobservice/grpc.rs b/tvix/castore/src/blobservice/grpc.rs
index 627f011a5041..d0f619fcebcb 100644
--- a/tvix/castore/src/blobservice/grpc.rs
+++ b/tvix/castore/src/blobservice/grpc.rs
@@ -9,7 +9,7 @@ use std::{
     task::Poll,
 };
 use tokio::io::AsyncWriteExt;
-use tokio::{net::UnixStream, task::JoinHandle};
+use tokio::task::JoinHandle;
 use tokio_stream::{wrappers::ReceiverStream, StreamExt};
 use tokio_util::{
     io::{CopyToBytes, SinkWriter},
@@ -44,49 +44,10 @@ impl BlobService for GRPCBlobService {
     /// - In the case of unix sockets, there must be a path, but may not be a host.
     /// - In the case of non-unix sockets, there must be a host, but no path.
     fn from_url(url: &url::Url) -> Result<Self, crate::Error> {
-        // Start checking for the scheme to start with grpc+.
-        match url.scheme().strip_prefix("grpc+") {
-            None => Err(crate::Error::StorageError("invalid scheme".to_string())),
-            Some(rest) => {
-                let channel = if rest == "unix" {
-                    if url.host_str().is_some() {
-                        return Err(crate::Error::StorageError(
-                            "host may not be set".to_string(),
-                        ));
-                    }
-                    let path = url.path().to_string();
-                    tonic::transport::Endpoint::try_from("http://[::]:50051") // doesn't matter
-                        .unwrap()
-                        .connect_with_connector_lazy(tower::service_fn(
-                            move |_: tonic::transport::Uri| UnixStream::connect(path.clone()),
-                        ))
-                } else {
-                    // ensure path is empty, not supported with gRPC.
-                    if !url.path().is_empty() {
-                        return Err(crate::Error::StorageError(
-                            "path may not be set".to_string(),
-                        ));
-                    }
-
-                    // clone the uri, and drop the grpc+ from the scheme.
-                    // Recreate a new uri with the `grpc+` prefix dropped from the scheme.
-                    // We can't use `url.set_scheme(rest)`, as it disallows
-                    // setting something http(s) that previously wasn't.
-                    let url = {
-                        let url_str = url.to_string();
-                        let s_stripped = url_str.strip_prefix("grpc+").unwrap();
-                        url::Url::parse(s_stripped).unwrap()
-                    };
-                    tonic::transport::Endpoint::try_from(url.to_string())
-                        .unwrap()
-                        .connect_lazy()
-                };
-
-                Ok(Self::from_client(
-                    proto::blob_service_client::BlobServiceClient::new(channel),
-                ))
-            }
-        }
+        let channel = crate::channel::from_url(url)?;
+        Ok(Self::from_client(
+            proto::blob_service_client::BlobServiceClient::new(channel),
+        ))
     }
 
     #[instrument(skip(self, digest), fields(blob.digest=%digest))]
diff --git a/tvix/castore/src/channel.rs b/tvix/castore/src/channel.rs
new file mode 100644
index 000000000000..53f9bd33442f
--- /dev/null
+++ b/tvix/castore/src/channel.rs
@@ -0,0 +1,68 @@
+use tokio::net::UnixStream;
+use tonic::transport::Channel;
+
+/// Turn a [url::Url] to a [Channel] if it can be parsed successfully.
+/// It supports `grpc+unix:/path/to/socket`,
+/// as well as the regular schemes supported by tonic, prefixed with grpc+,
+/// for example `grpc+http://[::1]:8000`.
+pub fn from_url(url: &url::Url) -> Result<Channel, self::Error> {
+    // Start checking for the scheme to start with grpc+.
+    // If it doesn't start with that, bail out.
+    match url.scheme().strip_prefix("grpc+") {
+        None => Err(Error::MissingGRPCPrefix()),
+        Some(rest) => {
+            if rest == "unix" {
+                if url.host_str().is_some() {
+                    return Err(Error::HostSetForUnixSocket());
+                }
+
+                let url = url.clone();
+                Ok(
+                    tonic::transport::Endpoint::from_static("http://[::]:50051") // doesn't matter
+                        .connect_with_connector_lazy(tower::service_fn(
+                            move |_: tonic::transport::Uri| {
+                                UnixStream::connect(url.path().to_string().clone())
+                            },
+                        )),
+                )
+            } else {
+                // ensure path is empty, not supported with gRPC.
+                if !url.path().is_empty() {
+                    return Err(Error::PathMayNotBeSet());
+                }
+
+                // Stringify the URL and remove the grpc+ prefix.
+                // We can't use `url.set_scheme(rest)`, as it disallows
+                // setting something http(s) that previously wasn't.
+                let url = url.to_string().strip_prefix("grpc+").unwrap().to_owned();
+
+                // Use the regular tonic transport::Endpoint logic to
+                Ok(tonic::transport::Endpoint::try_from(url)
+                    .unwrap()
+                    .connect_lazy())
+            }
+        }
+    }
+}
+
+/// Errors occuring when trying to connect to a backend
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("grpc+ prefix is missing from Url")]
+    MissingGRPCPrefix(),
+
+    #[error("host may not be set for unix domain sockets")]
+    HostSetForUnixSocket(),
+
+    #[error("path may not be set")]
+    PathMayNotBeSet(),
+
+    #[error("transport error: {0}")]
+    TransportError(tonic::transport::Error),
+}
+
+impl From<tonic::transport::Error> for Error {
+    fn from(value: tonic::transport::Error) -> Self {
+        Self::TransportError(value)
+    }
+}
diff --git a/tvix/castore/src/directoryservice/grpc.rs b/tvix/castore/src/directoryservice/grpc.rs
index c3207ad064c3..43d460fae629 100644
--- a/tvix/castore/src/directoryservice/grpc.rs
+++ b/tvix/castore/src/directoryservice/grpc.rs
@@ -6,7 +6,6 @@ use crate::proto::{self, get_directory_request::ByWhat};
 use crate::{B3Digest, Error};
 use async_stream::try_stream;
 use futures::Stream;
-use tokio::net::UnixStream;
 use tokio::spawn;
 use tokio::sync::mpsc::UnboundedSender;
 use tokio::task::JoinHandle;
@@ -42,49 +41,10 @@ impl DirectoryService for GRPCDirectoryService {
     /// - In the case of unix sockets, there must be a path, but may not be a host.
     /// - In the case of non-unix sockets, there must be a host, but no path.
     fn from_url(url: &url::Url) -> Result<Self, crate::Error> {
-        // Start checking for the scheme to start with grpc+.
-        match url.scheme().strip_prefix("grpc+") {
-            None => Err(crate::Error::StorageError("invalid scheme".to_string())),
-            Some(rest) => {
-                let channel = if rest == "unix" {
-                    if url.host_str().is_some() {
-                        return Err(crate::Error::StorageError(
-                            "host may not be set".to_string(),
-                        ));
-                    }
-                    let path = url.path().to_string();
-                    tonic::transport::Endpoint::try_from("http://[::]:50051") // doesn't matter
-                        .unwrap()
-                        .connect_with_connector_lazy(tower::service_fn(
-                            move |_: tonic::transport::Uri| UnixStream::connect(path.clone()),
-                        ))
-                } else {
-                    // ensure path is empty, not supported with gRPC.
-                    if !url.path().is_empty() {
-                        return Err(crate::Error::StorageError(
-                            "path may not be set".to_string(),
-                        ));
-                    }
-
-                    // clone the uri, and drop the grpc+ from the scheme.
-                    // Recreate a new uri with the `grpc+` prefix dropped from the scheme.
-                    // We can't use `url.set_scheme(rest)`, as it disallows
-                    // setting something http(s) that previously wasn't.
-                    let url = {
-                        let url_str = url.to_string();
-                        let s_stripped = url_str.strip_prefix("grpc+").unwrap();
-                        url::Url::parse(s_stripped).unwrap()
-                    };
-                    tonic::transport::Endpoint::try_from(url.to_string())
-                        .unwrap()
-                        .connect_lazy()
-                };
-
-                Ok(Self::from_client(
-                    proto::directory_service_client::DirectoryServiceClient::new(channel),
-                ))
-            }
-        }
+        let channel = crate::channel::from_url(url)?;
+        Ok(Self::from_client(
+            proto::directory_service_client::DirectoryServiceClient::new(channel),
+        ))
     }
 
     async fn get(
diff --git a/tvix/castore/src/errors.rs b/tvix/castore/src/errors.rs
index 3b23f972b045..9f596afd4e26 100644
--- a/tvix/castore/src/errors.rs
+++ b/tvix/castore/src/errors.rs
@@ -34,6 +34,12 @@ impl From<Error> for Status {
     }
 }
 
+impl From<crate::channel::Error> for Error {
+    fn from(value: crate::channel::Error) -> Self {
+        Self::StorageError(value.to_string())
+    }
+}
+
 // TODO: this should probably go somewhere else?
 impl From<Error> for std::io::Error {
     fn from(value: Error) -> Self {
diff --git a/tvix/castore/src/lib.rs b/tvix/castore/src/lib.rs
index d51022428e5b..5a031aec7686 100644
--- a/tvix/castore/src/lib.rs
+++ b/tvix/castore/src/lib.rs
@@ -2,6 +2,7 @@ mod digests;
 mod errors;
 
 pub mod blobservice;
+pub mod channel;
 pub mod directoryservice;
 pub mod fixtures;
 pub mod import;