From ebfe45625190f9920125aa82b1c6c0ec63abac0d Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Wed, 15 Nov 2023 17:40:42 +0200 Subject: refactor(tvix/castore/tonic): use match in channel_from_url Having random if blocks and returning from them is error-prone. Also, turns out we only need the unprefixed scheme in the fallback case, so move it down to there. Change-Id: Ifcb09279c963f8a39e0dbabe145990263f3d7cf9 Reviewed-on: https://cl.tvl.fyi/c/depot/+/10041 Autosubmit: flokli Tested-by: BuildkiteCI Reviewed-by: raitobezarius --- tvix/castore/src/tonic.rs | 93 ++++++++++++++++++++++++----------------------- 1 file changed, 48 insertions(+), 45 deletions(-) (limited to 'tvix/castore/src') diff --git a/tvix/castore/src/tonic.rs b/tvix/castore/src/tonic.rs index 96f7c7174141..8ea417548ce9 100644 --- a/tvix/castore/src/tonic.rs +++ b/tvix/castore/src/tonic.rs @@ -9,54 +9,57 @@ fn url_wants_wait_connect(url: &url::Url) -> bool { } /// 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, for example `grpc+http://[::1]:8000`. -/// It supports wait-connect=1 as a URL parameter, in which case we don't connect lazily. +/// It supports the following schemes (and URLs): +/// - `grpc+http://[::1]:8000`, connecting over unencrypted HTTP/2 (h2c) +/// - `grpc+https://[::1]:8000`, connecting over encrypted HTTP/2 +/// - `grpc+unix:/path/to/socket`, connecting to a unix domain socket +/// +/// All URLs support adding `wait-connect=1` as a URL parameter, in which case +/// the connection is established lazily. pub async fn channel_from_url(url: &url::Url) -> Result { - // 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 unprefixed_url_str = match url.to_string().strip_prefix("grpc+") { - None => return Err(Error::MissingGRPCPrefix()), - Some(url_str) => url_str.to_owned(), - }; - - if url.scheme() == "grpc+unix" { - if url.host_str().is_some() { - return Err(Error::HostSetForUnixSocket()); + match url.scheme() { + "grpc+unix" => { + if url.host_str().is_some() { + return Err(Error::HostSetForUnixSocket()); + } + + let connector = tower::service_fn({ + let url = url.clone(); + move |_: tonic::transport::Uri| UnixStream::connect(url.path().to_string().clone()) + }); + + // the URL doesn't matter + let endpoint = Endpoint::from_static("http://[::]:50051"); + if url_wants_wait_connect(url) { + Ok(endpoint.connect_with_connector(connector).await?) + } else { + Ok(endpoint.connect_with_connector_lazy(connector)) + } + } + _ => { + // 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 unprefixed_url_str = match url.to_string().strip_prefix("grpc+") { + None => return Err(Error::MissingGRPCPrefix()), + Some(url_str) => url_str.to_owned(), + }; + + // Use the regular tonic transport::Endpoint logic, but unprefixed_url_str, + // as tonic doesn't know about grpc+http[s]. + let endpoint = Endpoint::try_from(unprefixed_url_str)?; + if url_wants_wait_connect(url) { + Ok(endpoint.connect().await?) + } else { + Ok(endpoint.connect_lazy()) + } } - - let connector = tower::service_fn({ - let url = url.clone(); - move |_: tonic::transport::Uri| UnixStream::connect(url.path().to_string().clone()) - }); - - let channel = if url_wants_wait_connect(url) { - Endpoint::from_static("http://[::]:50051") - .connect_with_connector(connector) - .await? - } else { - Endpoint::from_static("http://[::]:50051").connect_with_connector_lazy(connector) - }; - - return Ok(channel); - } - - // ensure path is empty, not supported with gRPC. - if !url.path().is_empty() { - return Err(Error::PathMayNotBeSet()); } - - // Use the regular tonic transport::Endpoint logic, but unprefixed_url_str, - // as tonic doesn't know about grpc+http[s]. - let endpoint = Endpoint::try_from(unprefixed_url_str)?; - let channel = if url_wants_wait_connect(url) { - endpoint.connect().await? - } else { - endpoint.connect_lazy() - }; - - Ok(channel) } /// Errors occuring when trying to connect to a backend -- cgit 1.4.1