diff --git a/CLAUDE.md b/CLAUDE.md index af7d127d..2625e9b4 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -76,6 +76,7 @@ All built on the `motore` crate's `Service` and `Layer` traits. - **Service Discovery**: `Discover` trait (volo) -- `StaticDiscover`, `WeightedStaticDiscover` - **Load Balancing**: `LoadBalance` trait (volo) -- weighted random, consistent hashing +- **Transport Pool**: `volo::pool::Pool` -- unique or shared (multiplexed) transports per peer; used by volo-grpc ## Design Patterns diff --git a/Cargo.lock b/Cargo.lock index 46e894a1..86192192 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4172,6 +4172,7 @@ dependencies = [ "faststr", "futures", "libc", + "linked-hash-map", "metainfo", "motore", "mur3", @@ -4277,7 +4278,6 @@ dependencies = [ "futures", "futures-util", "h2", - "hex", "hickory-resolver", "http", "http-body", diff --git a/docs/grpc-client-transport.md b/docs/grpc-client-transport.md new file mode 100644 index 00000000..c6748ce0 --- /dev/null +++ b/docs/grpc-client-transport.md @@ -0,0 +1,160 @@ +# volo-grpc Client Transport: `:authority` and Connection Pooling + +## 1. Overview + +The volo-grpc client transport (`volo-grpc/src/transport/`) turns a call that service discovery +and load balancing have already routed to an `Address` into an HTTP/2 request. Two things about +it are worth knowing as a user: + +- **Which `:authority` the server sees.** It is derived from what the client was configured + with — the TLS server name or the service name — not from the socket that happens to be dialed. + A proxy or ingress that routes on virtual host therefore sees the same host name the client + presented for SNI, with no extra configuration. +- **How connections are kept.** One multiplexed HTTP/2 connection per callee address, held in the + generic `volo::pool::Pool` from the core crate, with idle eviction. + +## 2. `:authority` selection + +HTTP/2 has no `Host` header; the request's authority is carried in the `:authority` +pseudo-header, which hyper takes from the request URI. volo-grpc builds that URI per call from +the callee `Endpoint` in the context, using the first rule that applies: + +| # | Condition | `:authority` | +| - | -------------------------------------------------- | ---------------------------------------------------- | +| 0 | callee `Endpoint` carries a `volo_grpc::client::Authority` tag | the tag's value verbatim | +| 1 | `tls_config(ClientTlsConfig::new(server_name, ..))` | `server_name`, plus `:` of the dialed address unless the port is 443 | +| 2 | callee `service_name` is a valid authority | `service_name` verbatim | +| 3 | otherwise | the dialed `ip:port` (`localhost` for Unix sockets) | + +`:scheme` is `https` when TLS is configured and `http` otherwise. + +### Examples + +Talking to a TLS-terminating proxy that routes on virtual host: + +```rust +let client = GreeterClientBuilder::new("grpc.internal.example.com:50051") + .tls_config(ClientTlsConfig::new("grpc.internal.example.com", connector)) + .build(); +// SNI: grpc.internal.example.com +// :authority: grpc.internal.example.com:50051 (rule 1) +// :scheme: https +``` + +Plain gRPC with the default DNS resolver — the service name *is* the host: + +```rust +let client = GreeterClientBuilder::new("grpc.internal.example.com:50051").build(); +// :authority: grpc.internal.example.com:50051 (rule 2) +``` + +Plain gRPC with a custom `Discover` and a logical service name: + +```rust +let client = GreeterClientBuilder::new("user-service") + .discover(my_discover) + .build(); +// :authority: user-service (rule 2, same as grpc-go / tonic would send) +``` + +Fixed address, no usable name: + +```rust +let client = GreeterClientBuilder::new("").address(addr).build(); +// :authority: 10.0.0.7:8080 (rule 3) +``` + +Overriding, when the name the server must be addressed by differs from all of the above (a mesh +sidecar routing on a cluster name while TLS terminates at the sidecar, say). The tag lives on the +callee `Endpoint`'s `faststr_tags`, so it can be set per call through `CallOpt`, or per client by +a layer that runs before the transport: + +```rust +use volo_grpc::client::{Authority, CallOpt}; + +let mut opt = CallOpt::new(); +opt.callee_faststr_tags + .insert::(FastStr::from_static_str("users.mesh.local:50051")); +client.clone().with_opt(opt).get_user(req).await?; +// :authority: users.mesh.local:50051 (rule 0) +``` + +A tag that is not a valid authority (empty, userinfo, a scheme) is ignored and the derived +default applies. + +Rule 1 wins over rule 2 because SNI and `:authority` must agree for any TLS-terminating proxy; +the default port 443 is omitted the way HTTP clients omit it for `https`. An IP literal as +`server_name` is formatted as a socket address (`[::1]:8080`). + +### Behaviour change + +Before this design, `:authority` was always the resolved `ip:port`, because the same string was +also hyper's pool key and dial target. Plaintext clients whose `service_name` is a logical name +(e.g. `hello`) now send that name instead of `127.0.0.1:8080`. gRPC servers ignore `:authority` +unless they route on it; if yours does, this is the value it will see. + +## 3. Connection pooling + +Connections are keyed by the callee `Address` — the instance picked by the load balancer for this +call — so client-side load balancing across several instances keeps one connection per instance, +whatever `:authority` says. The transport holds a `volo::pool::Pool` +in `Mode::Shared`: + +- one HTTP/2 connection per address, all calls to that address multiplexed on it; +- concurrent first calls to a new address wait for the single in-flight handshake rather than + each dialing — and that handshake runs as its own task, so a caller hitting its rpc timeout + does not cancel the connection everyone else is waiting for; +- a connection the peer closed is noticed on the next call and replaced; a request that hyper + hands back untouched because the connection died underneath it is retried once on a fresh one; +- connections idle for 90 seconds are closed by the pool's background task. + +Timeouts (`connect_timeout`, `read_timeout`, `write_timeout`) and the `http2_*` settings on +`ClientBuilder` apply per connection exactly as before. + +## 4. `volo::pool` for extension developers + +The pool lives in the core crate (`volo/src/pool/`) so every protocol crate can share it. Using +it takes two impls: + +```rust +use volo::pool::{Mode, Pool, Poolable, Reservation}; + +// 1. The transport that lives in the pool. +#[derive(Clone)] +struct MyConn(/* ... */); + +impl Poolable for MyConn { + async fn reusable(&self) -> bool { /* still usable? */ true } + + // Only for multiplexed transports; the default is exclusive use. + fn reserve(self) -> Reservation { Reservation::Shared(self.clone(), self) } + fn can_share(&self) -> bool { true } + fn try_checkout(&self) -> Option { Some(self.clone()) } +} + +// 2. Something that makes one: any `motore::UnaryService`. +#[derive(Clone)] +struct MyConnector; + +impl motore::UnaryService for MyConnector { + type Response = MyConn; + type Error = MyError; + async fn call(&self, addr: volo::net::Address) -> Result { /* dial */ } +} + +let pool = Pool::new(volo::pool::Config::default().idle_timeout(Duration::from_secs(90))); +let conn = pool.get(addr, Mode::Shared, MyConnector).await?; // Result, pool::Error> +``` + +- `Mode::Unique` transports (one request at a time, like thrift ping-pong) are returned with + `Pooled::reuse().await` after use; dropping the `Pooled` without it discards the transport. +- `Mode::Shared` transports need nothing after use; the pool keeps its own copy. +- `pool::Error` is generic over the connector's error: `Connect(E)` when making a transport + failed, `Canceled` when this caller was waiting on someone else's connect that failed (or the + pool was dropped). Map it to your crate's error type; volo-grpc does + `impl From> for Status`. +- `Pool::new` is safe to call outside a runtime (in a `LazyLock`, say); the idle-eviction task is + started on the first `get`. + +`volo-thrift` and `volo-http` still carry their own copies of this pool design; migrating them to +`volo::pool` is a follow-up. diff --git a/volo-grpc/CLAUDE.md b/volo-grpc/CLAUDE.md index 79e4bff2..2285e852 100644 --- a/volo-grpc/CLAUDE.md +++ b/volo-grpc/CLAUDE.md @@ -33,7 +33,7 @@ volo-grpc/src/ ├── codec/ # Codec trait, encode/decode, compression (gzip/zlib/zstd) ├── metadata/ # MetadataMap, MetadataKey, MetadataValue (binary keys use `-bin` suffix) ├── layer/ # Shared layers: loadbalance, grpc_timeout, grpc_web, user_agent, CORS -└── transport/ # Client transport, connection, TLS config +└── transport/ # Client transport: per-address HTTP/2 connections, `:authority` selection ``` ## Key Components @@ -90,3 +90,4 @@ Server HTTP/2 settings: 2. **gRPC-Web requires additional configuration**: Enable `grpc-web` feature and set `accept_http1(true)` 3. **Compression is optional**: gzip and zlib enabled by default, zstd requires manual enabling 4. **TLS requires backend selection**: rustls or native-tls +5. **`:authority` is derived, not configured**: the client transport pools HTTP/2 connections in `volo::pool::Pool` (shared mode, 90s idle timeout) keyed by the callee `Address`, and builds the request URI from the callee `Endpoint`: an explicit `client::Authority` faststr tag on the callee if present, else TLS `server_name` (+ non-443 port) if TLS is on, else the callee `service_name`, else the address. See `transport/client.rs::authority`. diff --git a/volo-grpc/Cargo.toml b/volo-grpc/Cargo.toml index 0904c244..e498aa00 100644 --- a/volo-grpc/Cargo.toml +++ b/volo-grpc/Cargo.toml @@ -39,7 +39,6 @@ faststr.workspace = true futures-util.workspace = true futures.workspace = true h2.workspace = true -hex.workspace = true hickory-resolver.workspace = true http.workspace = true http-body.workspace = true @@ -49,7 +48,6 @@ hyper-timeout.workspace = true hyper-util = { workspace = true, features = [ "tokio", "client", - "client-legacy", "server", "http2", ] } diff --git a/volo-grpc/src/client/mod.rs b/volo-grpc/src/client/mod.rs index 6f511565..603c3295 100644 --- a/volo-grpc/src/client/mod.rs +++ b/volo-grpc/src/client/mod.rs @@ -37,6 +37,25 @@ use crate::{ }; pub mod layer; +/// Tag type for overriding the HTTP/2 `:authority` of a call. +/// +/// By default the authority is derived from the client configuration (the TLS server name, else +/// the service name, else the callee address). When the server has to be addressed by a name +/// that differs from all of those -- a mesh sidecar routing on a cluster name, say -- set the +/// value on the callee [`Endpoint`]'s `faststr_tags` under this type, per call through +/// [`CallOpt::callee_faststr_tags`] or per client from a layer: +/// +/// ```rust,ignore +/// let mut opt = CallOpt::new(); +/// opt.callee_faststr_tags +/// .insert::(FastStr::from_static_str("users.mesh.local:50051")); +/// client.clone().with_opt(opt).get_user(req).await?; +/// ``` +/// +/// The value must be a valid HTTP authority (`host[:port]`, no scheme, path or userinfo); +/// anything else is ignored in favour of the derived default. +pub struct Authority; + /// [`ClientBuilder`] provides a builder-like interface to construct a [`Client`]. pub struct ClientBuilder { http2_config: Http2Config, @@ -66,6 +85,12 @@ impl > { /// Creates a new [`ClientBuilder`]. + /// + /// With the default [`DnsResolver`] the `service_name` is the `host[:port]` to resolve and + /// connect to. It is also what the client sends as the HTTP/2 `:authority` (unless TLS is + /// configured, in which case the TLS server name is used, see [`Self::tls_config`]), so a + /// proxy or ingress in front of the server sees the same host name the client was + /// configured with. See [`Authority`] to override it. pub fn new(service_client: C, service_name: impl AsRef) -> Self { Self { http2_config: Default::default(), @@ -466,6 +491,11 @@ impl ClientBuilder { /// Sets the [`ClientTlsConfig`][ClientTlsConfig] for the client. /// + /// Besides being used as the TLS SNI, the `server_name` of the config is sent as the HTTP/2 + /// `:authority` of every call (with the port of the callee address appended unless it is + /// 443), so a TLS-terminating proxy that routes on the virtual host sees the same name it + /// presented a certificate for. See [`Authority`] to override it. + /// /// [ClientTlsConfig]: volo::net::tls::ClientTlsConfig #[cfg(feature = "__tls")] #[cfg_attr(docsrs, doc(cfg(any(feature = "rustls", feature = "native-tls"))))] diff --git a/volo-grpc/src/status.rs b/volo-grpc/src/status.rs index 54fa2b6d..1468dded 100644 --- a/volo-grpc/src/status.rs +++ b/volo-grpc/src/status.rs @@ -764,6 +764,15 @@ impl From for Status { } } +impl From> for Status { + fn from(err: volo::pool::Error) -> Self { + match err { + volo::pool::Error::Connect(status) => status, + volo::pool::Error::Canceled => Status::unavailable(err.to_string()), + } + } +} + impl From for Status { fn from(err: anyhow::Error) -> Self { Self::from_error(err.into()) diff --git a/volo-grpc/src/transport/client.rs b/volo-grpc/src/transport/client.rs index 759d9b6c..8c1f7508 100644 --- a/volo-grpc/src/transport/client.rs +++ b/volo-grpc/src/transport/client.rs @@ -1,16 +1,30 @@ -use std::{io, marker::PhantomData}; +use std::{ + io, + marker::PhantomData, + net::{IpAddr, SocketAddr}, + str::FromStr, + time::Duration, +}; use bytes::Bytes; use http::{ HeaderValue, header::{CONTENT_TYPE, TE}, + uri::{Authority, Scheme}, }; use http_body::Frame; use http_body_util::StreamBody; -use hyper_util::rt::{TokioExecutor, TokioTimer}; -use motore::Service; -use tower::{Service as TowerService, util::ServiceExt}; -use volo::net::Address; +use hyper::{ + body::Incoming, + client::conn::http2::{Builder as Http2Builder, SendRequest}, +}; +use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer}; +use motore::{Service, make::MakeConnection, service::UnaryService}; +use volo::{ + context::Endpoint, + net::Address, + pool::{Mode, Pool, Poolable, Pooled, Reservation}, +}; use super::connect::Connector; use crate::{ @@ -24,21 +38,85 @@ use crate::{ context::{ClientContext, Config}, }; -/// A simple wrapper of [`hyper_util::client::legacy::Client`] that implements [`Service`] -/// to make outgoing requests. -#[allow(clippy::type_complexity)] +type Body = StreamBody, Status>>>; + +/// Idle connections are dropped after this long; matches what hyper's client pool did. +const IDLE_TIMEOUT: Duration = Duration::from_secs(90); + +/// A multiplexed HTTP/2 connection to one peer, shared by every call to that peer. +#[derive(Clone)] +struct Http2Connection(SendRequest); + +impl Poolable for Http2Connection { + async fn reusable(&self) -> bool { + !self.0.is_closed() + } + + fn reserve(self) -> Reservation { + Reservation::Shared(self.clone(), self) + } + + fn can_share(&self) -> bool { + true + } + + fn try_checkout(&self) -> Option { + (!self.0.is_closed()).then(|| self.clone()) + } +} + +/// Dials an address and runs the HTTP/2 handshake on it. +#[derive(Clone)] +struct Http2Connector { + connector: Connector, + http2_builder: Http2Builder, +} + +impl UnaryService
for Http2Connector { + type Response = Http2Connection; + type Error = Status; + + async fn call(&self, addr: Address) -> Result { + let io = self + .connector + .make_connection(addr) + .await + .map_err(|err| Status::from_error(err.into()))?; + let (mut tx, conn) = self + .http2_builder + .handshake::<_, Body>(TokioIo::new(io)) + .await + .map_err(|err| Status::from_error(err.into()))?; + tokio::spawn(async move { + if let Err(err) = conn.await { + tracing::debug!("[VOLO] http2 client connection error: {err}"); + } + }); + // Wait for the connection to accept requests before handing it out. + tx.ready() + .await + .map_err(|err| Status::from_error(err.into()))?; + Ok(Http2Connection(tx)) + } +} + +/// gRPC client transport: one HTTP/2 connection per callee address, requests multiplexed on it. +/// +/// Connections live in a [`volo::pool::Pool`] keyed by the [`Address`] chosen by service +/// discovery / load balancing, while the request URI is built from the callee [`Endpoint`], so +/// the HTTP/2 `:authority` names the server (the SNI name or the service name) rather than the +/// socket that happens to be dialed. See `authority` in this module for the exact rule. pub struct ClientTransport { - http_client: hyper_util::client::legacy::Client< - Connector, - StreamBody, crate::Status>>>, - >, + connector: Http2Connector, + pool: Pool, _marker: PhantomData, } impl Clone for ClientTransport { fn clone(&self) -> Self { Self { - http_client: self.http_client.clone(), + connector: self.connector.clone(), + pool: self.pool.clone(), _marker: self._marker, } } @@ -47,63 +125,91 @@ impl Clone for ClientTransport { impl ClientTransport { /// Creates a new [`ClientTransport`] by setting the underlying connection /// with the given config. + #[must_use] pub fn new(http2_config: &Http2Config, rpc_config: &Config) -> Self { - let config = volo::net::dial::Config::new( - rpc_config.connect_timeout, - rpc_config.read_timeout, - rpc_config.write_timeout, - ); - let http_client = hyper_util::client::legacy::Client::builder(TokioExecutor::new()) - .timer(TokioTimer::new()) - .http2_only(true) - .http2_initial_stream_window_size(http2_config.init_stream_window_size) - .http2_initial_connection_window_size(http2_config.init_connection_window_size) - .http2_max_frame_size(http2_config.max_frame_size) - .http2_adaptive_window(http2_config.adaptive_window) - .http2_keep_alive_interval(http2_config.http2_keepalive_interval) - .http2_keep_alive_timeout(http2_config.http2_keepalive_timeout) - .http2_keep_alive_while_idle(http2_config.http2_keepalive_while_idle) - .http2_max_concurrent_reset_streams(http2_config.max_concurrent_reset_streams) - .http2_max_send_buf_size(http2_config.max_send_buf_size) - .build(Connector::new(Some(config))); - - ClientTransport { - http_client, - _marker: PhantomData, - } + Self::with_connector(http2_config, Connector::new(Some(dial_config(rpc_config)))) } #[cfg(feature = "__tls")] #[cfg_attr(docsrs, doc(cfg(any(feature = "rustls", feature = "native-tls"))))] + #[must_use] pub fn new_with_tls( http2_config: &Http2Config, rpc_config: &Config, tls_config: volo::net::tls::ClientTlsConfig, ) -> Self { - let config = volo::net::dial::Config::new( - rpc_config.connect_timeout, - rpc_config.read_timeout, - rpc_config.write_timeout, - ); - let http_client = hyper_util::client::legacy::Client::builder(TokioExecutor::new()) - .timer(TokioTimer::new()) - .http2_only(true) - .http2_initial_stream_window_size(http2_config.init_stream_window_size) - .http2_initial_connection_window_size(http2_config.init_connection_window_size) - .http2_max_frame_size(http2_config.max_frame_size) - .http2_adaptive_window(http2_config.adaptive_window) - .http2_keep_alive_interval(http2_config.http2_keepalive_interval) - .http2_keep_alive_timeout(http2_config.http2_keepalive_timeout) - .http2_keep_alive_while_idle(http2_config.http2_keepalive_while_idle) - .http2_max_concurrent_reset_streams(http2_config.max_concurrent_reset_streams) - .http2_max_send_buf_size(http2_config.max_send_buf_size) - .build(Connector::new_with_tls(Some(config), tls_config)); - - ClientTransport { - http_client, + Self::with_connector( + http2_config, + Connector::new_with_tls(Some(dial_config(rpc_config)), tls_config), + ) + } + + fn with_connector(http2_config: &Http2Config, connector: Connector) -> Self { + Self { + connector: Http2Connector { + connector, + http2_builder: http2_builder(http2_config), + }, + pool: Pool::new(volo::pool::Config::default().idle_timeout(IDLE_TIMEOUT)), _marker: PhantomData, } } + + async fn connection(&self, addr: &Address) -> Result, Status> { + self.pool + .get(addr.clone(), Mode::Shared, self.connector.clone()) + .await + .map_err(Status::from) + } + + async fn send( + &self, + addr: &Address, + req: http::Request, + ) -> Result, Status> { + let mut conn = self.connection(addr).await?; + let mut err = match conn.0.try_send_request(req).await { + Ok(resp) => return Ok(resp), + Err(err) => err, + }; + // The connection went away between the pool's liveness check and the dispatch. hyper + // hands the request back untouched in that case, so retry it once on a fresh connection; + // the pool sees the closed one and replaces it. + if let Some(req) = err.take_message() { + tracing::debug!("[VOLO] http2 connection to {addr} was closed, reconnecting"); + let mut conn = self.connection(addr).await?; + return conn + .0 + .send_request(req) + .await + .map_err(|err| Status::from_error(err.into())); + } + Err(Status::from_error(err.into_error().into())) + } +} + +fn dial_config(rpc_config: &Config) -> volo::net::dial::Config { + volo::net::dial::Config::new( + rpc_config.connect_timeout, + rpc_config.read_timeout, + rpc_config.write_timeout, + ) +} + +fn http2_builder(config: &Http2Config) -> Http2Builder { + let mut builder = Http2Builder::new(TokioExecutor::new()); + builder + .timer(TokioTimer::new()) + .initial_stream_window_size(config.init_stream_window_size) + .initial_connection_window_size(config.init_connection_window_size) + .max_frame_size(config.max_frame_size) + .adaptive_window(config.adaptive_window) + .keep_alive_interval(config.http2_keepalive_interval) + .keep_alive_timeout(config.http2_keepalive_timeout) + .keep_alive_while_idle(config.http2_keepalive_while_idle) + .max_concurrent_reset_streams(config.max_concurrent_reset_streams) + .max_send_buf_size(config.max_send_buf_size); + builder } impl Service> for ClientTransport @@ -121,7 +227,6 @@ where cx: &mut ClientContext, volo_req: Request, ) -> Result { - let mut http_client = self.http_client.clone(); // SAFETY: parameters controlled by volo-grpc are guaranteed to be valid. // get the call address from the context let target = cx.rpc_info.callee().address().ok_or_else(|| { @@ -141,10 +246,19 @@ where let body = http_body_util::StreamBody::new(message.into_body(send_compression)); + let uri = build_uri( + self.connector.connector.scheme(), + authority( + self.connector.connector.tls_server_name(), + cx.rpc_info.callee(), + &target, + ), + path, + ); let mut req = http::Request::builder() .version(http::Version::HTTP_2) .method(http::Method::POST) - .uri(build_uri(target.clone(), path)) + .uri(uri) .extension(extensions) .body(body) .map_err(|err| Status::from_error(err.into()))?; @@ -171,13 +285,7 @@ where } cx.stats.record_make_transport_start_at(); - let resp = http_client - .ready() - .await - .map_err(|err| Status::from_error(err.into()))? - .call(req) - .await - .map_err(|err| Status::from_error(err.into()))?; + let resp = self.send(&target, req).await?; cx.stats.record_make_transport_end_at(); @@ -214,60 +322,197 @@ where } } -fn build_uri(addr: Address, path: &str) -> hyper::Uri { +/// Picks the HTTP/2 `:authority` for a call. +/// +/// gRPC uses `:authority` as the virtual host of the callee, so it has to name the server rather +/// than the socket that happens to be dialed: TLS-terminating proxies and ingresses route on it, +/// and it has to agree with the TLS SNI. In order of preference: +/// +/// 0. an explicit override on the callee endpoint, tagged [`crate::client::Authority`]; +/// 1. the TLS server name (the SNI), with the dialed port appended unless it is the default 443; +/// 2. the callee's service name, which is the `host[:port]` given to the DNS resolver, or the +/// logical name used with a custom `Discover`; +/// 3. the dialed address itself. +// Which `Address` variants exist depends on the target and on volo's features, so the +// non-IP arms have to be wildcards. +#[allow(clippy::match_wildcard_for_single_variants)] +fn authority(tls_server_name: Option<&str>, callee: &Endpoint, addr: &Address) -> Authority { + if let Some(authority) = callee + .get_faststr::() + .and_then(|name| parse_authority(name)) + { + return authority; + } + + let port = match addr { + Address::Ip(addr) => Some(addr.port()), + #[allow(unreachable_patterns)] + _ => None, + }; + + if let Some(name) = tls_server_name { + let host_port = match (name.parse::(), port) { + (Ok(ip), Some(port)) => SocketAddr::new(ip, port).to_string(), + (Ok(IpAddr::V6(ip)), None) => format!("[{ip}]"), + (Err(_), Some(port)) if port != 443 => format!("{name}:{port}"), + _ => name.to_owned(), + }; + if let Ok(authority) = Authority::from_str(&host_port) { + return authority; + } + } + + if let Some(authority) = parse_authority(callee.service_name_ref()) { + return authority; + } + match addr { - Address::Ip(ip) => hyper::Uri::builder() - .scheme(http::uri::Scheme::HTTP) - .authority(ip.to_string()) - .path_and_query(path) - .build() - .expect("fail to build ip uri"), - #[cfg(target_family = "unix")] - Address::Unix(unix) => hyper::Uri::builder() - .scheme("http+unix") - .authority(hex::encode( - unix.as_pathname() - .expect("target address is an invalid unix socket") - .to_string_lossy() - .as_bytes(), - )) - .path_and_query(path) - .build() - .expect("fail to build unix uri"), + Address::Ip(addr) => { + Authority::from_str(&addr.to_string()).expect("socket addr is a valid authority") + } #[allow(unreachable_patterns)] - _ => unimplemented!("unsupported type of address"), + _ => Authority::from_static("localhost"), + } +} + +/// Parses a user-supplied name as an `:authority`, rejecting what HTTP/2 forbids there: an +/// empty value, or one carrying userinfo. +fn parse_authority(name: &str) -> Option { + if name.is_empty() || name.contains('@') { + return None; } + Authority::from_str(name).ok() +} + +fn build_uri(scheme: Scheme, authority: Authority, path: &str) -> hyper::Uri { + hyper::Uri::builder() + .scheme(scheme) + .authority(authority) + .path_and_query(path) + .build() + .expect("fail to build uri") } #[cfg(test)] mod tests { + use std::net::SocketAddr; + + use http::uri::{Authority, Scheme}; + use volo::{FastStr, context::Endpoint, net::Address}; + + use super::{authority, build_uri}; + + fn endpoint(service_name: &str) -> Endpoint { + Endpoint::new(FastStr::new(service_name)) + } + + fn ip(addr: &str) -> Address { + Address::from(addr.parse::().unwrap()) + } #[test] - fn test_build_uri_ip() { - let addr = "127.0.0.1:8000".parse::().unwrap(); - let path = "/path?query=1"; - let uri = "http://127.0.0.1:8000/path?query=1" - .parse::() - .unwrap(); - assert_eq!(super::build_uri(volo::net::Address::from(addr), path), uri); + fn authority_tag_overrides_everything() { + let mut callee = endpoint("grpc.example.com:50051"); + callee.insert_faststr::(FastStr::from_static_str( + "users.mesh.local:50051", + )); + let a = authority(Some("grpc.example.com"), &callee, &ip("10.0.0.1:50051")); + assert_eq!(a, Authority::from_static("users.mesh.local:50051")); + } + + #[test] + fn authority_tag_that_is_not_an_authority_is_ignored() { + for bad in ["", "user@host", "http://host", "not a host"] { + let mut callee = endpoint("grpc.example.com:50051"); + callee.insert_faststr::(FastStr::new(bad)); + let a = authority(None, &callee, &ip("10.0.0.1:50051")); + assert_eq!( + a, + Authority::from_static("grpc.example.com:50051"), + "{bad:?}" + ); + } + } + + #[test] + fn authority_prefers_tls_server_name_with_dialed_port() { + let a = authority( + Some("grpc.example.com"), + &endpoint("some-logical-name"), + &ip("10.0.0.1:50051"), + ); + assert_eq!(a, Authority::from_static("grpc.example.com:50051")); + } + + #[test] + fn authority_omits_default_https_port_for_tls_server_name() { + let a = authority( + Some("grpc.example.com"), + &endpoint("grpc.example.com:443"), + &ip("10.0.0.1:443"), + ); + assert_eq!(a, Authority::from_static("grpc.example.com")); + } + + #[test] + fn authority_handles_ip_server_names() { + let a = authority(Some("::1"), &endpoint("svc"), &ip("[::1]:8080")); + assert_eq!(a, Authority::from_static("[::1]:8080")); + let a = authority(Some("127.0.0.1"), &endpoint("svc"), &ip("127.0.0.1:8080")); + assert_eq!(a, Authority::from_static("127.0.0.1:8080")); + } + + #[test] + fn authority_falls_back_to_service_name() { + let a = authority( + None, + &endpoint("grpc.example.com:50051"), + &ip("10.0.0.1:50051"), + ); + assert_eq!(a, Authority::from_static("grpc.example.com:50051")); + let a = authority(None, &endpoint("user-service"), &ip("10.0.0.1:50051")); + assert_eq!(a, Authority::from_static("user-service")); + } + + #[test] + fn authority_falls_back_to_address_for_unusable_service_names() { + for name in ["", "not a host", "user@host", "http://host:80"] { + let a = authority(None, &endpoint(name), &ip("10.0.0.1:50051")); + assert_eq!(a, Authority::from_static("10.0.0.1:50051"), "{name:?}"); + } } #[cfg(target_family = "unix")] #[test] - fn test_build_uri_unix() { - let addr = "/tmp/rpc.sock".parse::().unwrap(); - let path = "/path?query=1"; - let uri = "http+unix://2f746d702f7270632e736f636b/path?query=1" - .parse::() - .unwrap(); + fn authority_for_unix_socket_without_service_name() { + let addr = std::os::unix::net::SocketAddr::from_pathname("/tmp/rpc.sock").unwrap(); + let a = authority(None, &endpoint(""), &Address::from(addr)); + assert_eq!(a, Authority::from_static("localhost")); + } + + #[test] + fn test_build_uri() { + let uri = build_uri( + Scheme::HTTP, + Authority::from_static("127.0.0.1:8000"), + "/path?query=1", + ); assert_eq!( - super::build_uri( - volo::net::Address::from( - std::os::unix::net::SocketAddr::from_pathname(addr).unwrap() - ), - path - ), - uri + uri, + "http://127.0.0.1:8000/path?query=1" + .parse::() + .unwrap() + ); + let uri = build_uri( + Scheme::HTTPS, + Authority::from_static("grpc.example.com:50051"), + "/pkg.Svc/Method", + ); + assert_eq!( + uri, + "https://grpc.example.com:50051/pkg.Svc/Method" + .parse::() + .unwrap() ); } @@ -277,4 +522,214 @@ mod tests { fn test_is_unpin() { is_unpin::>(); } + + mod wire { + //! Drive the transport against a real HTTP/2 server and look at what arrives. + + use std::{ + net::SocketAddr, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + }; + + use bytes::Bytes; + use futures::StreamExt; + use http::header::CONTENT_TYPE; + use http_body::Frame; + use http_body_util::Empty; + use hyper::service::service_fn; + use hyper_util::rt::{TokioExecutor, TokioIo}; + use motore::Service; + use tokio::net::TcpListener; + use volo::{ + FastStr, + context::{Endpoint, Role, RpcInfo}, + net::Address, + }; + + use crate::{ + Code, Request, Status, + body::BoxBody, + client::Http2Config, + codec::{compression::CompressionEncoding, decode::Kind}, + context::{ClientContext, Config}, + message::{RecvEntryMessage, SendEntryMessage}, + transport::ClientTransport, + }; + + struct Empty_; + + impl SendEntryMessage for Empty_ { + fn into_body( + self, + _: Option, + ) -> crate::BoxStream<'static, Result, Status>> { + futures::stream::empty().boxed() + } + } + + impl RecvEntryMessage for Empty_ { + fn from_body( + _: Option<&str>, + _: BoxBody, + _: Kind, + _: Option, + ) -> Result { + Ok(Self) + } + } + + #[derive(Default)] + struct Observed { + authorities: Mutex>, + connections: AtomicUsize, + } + + /// A plain HTTP/2 server that records `:authority` of every request and answers each + /// with `grpc-status: UNIMPLEMENTED` in the headers, which the client surfaces as an + /// error without ever touching the body. + async fn serve() -> (SocketAddr, Arc) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let observed = Arc::new(Observed::default()); + let observed_ = observed.clone(); + tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.unwrap(); + observed_.connections.fetch_add(1, Ordering::SeqCst); + let observed = observed_.clone(); + tokio::spawn(async move { + let svc = service_fn(move |req: http::Request| { + let observed = observed.clone(); + async move { + let authority = req + .uri() + .authority() + .map(|a| a.to_string()) + .unwrap_or_default(); + observed.authorities.lock().unwrap().push(authority); + http::Response::builder() + .header(CONTENT_TYPE, "application/grpc") + .header("grpc-status", Code::Unimplemented as i32) + .body(Empty::::new()) + } + }); + let _ = hyper::server::conn::http2::Builder::new(TokioExecutor::new()) + .serve_connection(TokioIo::new(stream), svc) + .await; + }); + } + }); + (addr, observed) + } + + fn cx(service_name: &str, addr: SocketAddr) -> ClientContext { + let mut callee = Endpoint::new(FastStr::new(service_name)); + callee.set_address(Address::from(addr)); + ClientContext::new(RpcInfo::new( + Role::Client, + FastStr::from_static_str("/pkg.Svc/Method"), + Endpoint::new(FastStr::from_static_str("caller")), + callee, + Config::default(), + )) + } + + async fn call(transport: &ClientTransport, cx: &mut ClientContext) { + let err = transport + .call(cx, Request::new(Empty_)) + .await + .err() + .expect("server answers with a non-OK grpc-status"); + assert_eq!(err.code(), Code::Unimplemented, "{err:?}"); + } + + #[tokio::test] + async fn authority_is_the_service_name_and_connections_are_reused() { + let (addr, observed) = serve().await; + let transport = + ClientTransport::::new(&Http2Config::default(), &Config::default()); + + call(&transport, &mut cx("grpc.example.com:50051", addr)).await; + call(&transport, &mut cx("grpc.example.com:50051", addr)).await; + // A different logical callee behind the same address still shares the connection. + call(&transport, &mut cx("user-service", addr)).await; + + assert_eq!( + *observed.authorities.lock().unwrap(), + [ + "grpc.example.com:50051", + "grpc.example.com:50051", + "user-service" + ] + ); + assert_eq!(observed.connections.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn authority_tag_is_sent_on_the_wire() { + let (addr, observed) = serve().await; + let transport = + ClientTransport::::new(&Http2Config::default(), &Config::default()); + + let mut cx = cx("grpc.example.com:50051", addr); + cx.rpc_info + .callee_mut() + .insert_faststr::(FastStr::from_static_str( + "users.mesh.local:50051", + )); + call(&transport, &mut cx).await; + + assert_eq!( + *observed.authorities.lock().unwrap(), + ["users.mesh.local:50051"] + ); + } + + #[tokio::test] + async fn authority_falls_back_to_the_address() { + let (addr, observed) = serve().await; + let transport = + ClientTransport::::new(&Http2Config::default(), &Config::default()); + + call(&transport, &mut cx("", addr)).await; + + assert_eq!(*observed.authorities.lock().unwrap(), [addr.to_string()]); + } + + #[tokio::test] + async fn concurrent_first_calls_share_one_handshake() { + let (addr, observed) = serve().await; + let transport = + ClientTransport::::new(&Http2Config::default(), &Config::default()); + + futures::future::join_all((0..32).map(|_| { + let transport = transport.clone(); + async move { call(&transport, &mut cx("user-service", addr)).await } + })) + .await; + + assert_eq!(observed.authorities.lock().unwrap().len(), 32); + assert_eq!(observed.connections.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn distinct_addresses_get_distinct_connections() { + let (addr_a, observed_a) = serve().await; + let (addr_b, observed_b) = serve().await; + let transport = + ClientTransport::::new(&Http2Config::default(), &Config::default()); + + call(&transport, &mut cx("user-service", addr_a)).await; + call(&transport, &mut cx("user-service", addr_b)).await; + call(&transport, &mut cx("user-service", addr_a)).await; + + assert_eq!(observed_a.connections.load(Ordering::SeqCst), 1); + assert_eq!(observed_b.connections.load(Ordering::SeqCst), 1); + assert_eq!(observed_a.authorities.lock().unwrap().len(), 2); + assert_eq!(observed_b.authorities.lock().unwrap().len(), 1); + } + } } diff --git a/volo-grpc/src/transport/connect.rs b/volo-grpc/src/transport/connect.rs index 579311f6..daca4373 100644 --- a/volo-grpc/src/transport/connect.rs +++ b/volo-grpc/src/transport/connect.rs @@ -1,17 +1,7 @@ -#[cfg(target_family = "unix")] -use std::os::unix::net::SocketAddr as UnixSocketAddr; -use std::{ - io, - net::SocketAddr, - pin::Pin, - task::{Context, Poll}, -}; +use std::io; -use futures_util::future::BoxFuture; -use hyper::rt::ReadBufCursor; -use hyper_util::client::legacy::connect::{Connected, Connection}; +use http::uri::Scheme; use motore::{make::MakeConnection, service::UnaryService}; -use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; #[cfg(feature = "__tls")] #[cfg_attr(docsrs, doc(cfg(any(feature = "rustls", feature = "native-tls"))))] use volo::net::tls::{ClientTlsConfig, TlsMakeTransport}; @@ -21,12 +11,21 @@ use volo::net::{ dial::{Config, DefaultMakeTransport, MakeTransport}, }; +/// Dials the callee address picked by service discovery and load balancing. +/// +/// The connector only knows about [`Address`]es; the request URI (and therefore the HTTP/2 +/// `:authority`) is chosen independently by the transport, see +/// [`ClientTransport`][super::ClientTransport]. #[derive(Clone, Debug)] pub enum Connector { Default(DefaultMakeTransport), #[cfg(feature = "__tls")] #[cfg_attr(docsrs, doc(cfg(any(feature = "rustls", feature = "native-tls"))))] - Tls(TlsMakeTransport), + Tls { + transport: TlsMakeTransport, + /// The SNI server name, kept so the transport can reuse it as `:authority`. + server_name: volo::FastStr, + }, } impl Connector { @@ -43,13 +42,35 @@ impl Connector { #[cfg(feature = "__tls")] #[cfg_attr(docsrs, doc(cfg(any(feature = "rustls", feature = "native-tls"))))] pub fn new_with_tls(cfg: Option, tls_config: ClientTlsConfig) -> Self { + let server_name = volo::FastStr::new(&tls_config.server_name); let mut mt = TlsMakeTransport::new(cfg.unwrap_or_default(), tls_config); if let Some(cfg) = cfg { mt.set_connect_timeout(cfg.connect_timeout); mt.set_read_timeout(cfg.read_timeout); mt.set_write_timeout(cfg.write_timeout); } - Self::Tls(mt) + Self::Tls { + transport: mt, + server_name, + } + } + + /// The URI scheme matching the transport security, sent as the `:scheme` pseudo-header. + pub fn scheme(&self) -> Scheme { + match self { + Self::Default(_) => Scheme::HTTP, + #[cfg(feature = "__tls")] + Self::Tls { .. } => Scheme::HTTPS, + } + } + + /// The server name used for SNI, if the connector speaks TLS. + pub fn tls_server_name(&self) -> Option<&str> { + match self { + Self::Default(_) => None, + #[cfg(feature = "__tls")] + Self::Tls { server_name, .. } => Some(server_name), + } } } @@ -67,161 +88,7 @@ impl UnaryService
for Connector { match self { Self::Default(mkt) => mkt.make_connection(addr).await, #[cfg(feature = "__tls")] - Self::Tls(mkt) => mkt.make_connection(addr).await, - } - } -} - -impl tower::Service for Connector { - type Response = ConnectionWrapper; - - type Error = io::Error; - - type Future = BoxFuture<'static, Result>; - - fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn call(&mut self, uri: hyper::Uri) -> Self::Future { - let connector = self.clone(); - Box::pin(async move { - let authority = uri.authority().expect("authority required").as_str(); - let target: Address = match uri.scheme_str() { - Some("http") => Address::Ip(authority.parse::().map_err(|_| { - io::Error::new( - io::ErrorKind::InvalidInput, - "authority must be valid SocketAddr", - ) - })?), - #[cfg(target_family = "unix")] - Some("http+unix") => { - use hex::FromHex; - - let bytes = Vec::from_hex(authority).map_err(|_| { - io::Error::new( - io::ErrorKind::InvalidInput, - "authority must be hex-encoded path", - ) - })?; - Address::Unix(UnixSocketAddr::from_pathname( - String::from_utf8(bytes).map_err(|_| { - io::Error::new( - io::ErrorKind::InvalidInput, - "authority must be valid UTF-8", - ) - })?, - )?) - } - _ => unimplemented!(), - }; - - Ok(ConnectionWrapper { - inner: connector.make_connection(target).await?, - }) - }) - } -} - -#[pin_project::pin_project] -pub struct ConnectionWrapper { - #[pin] - inner: Conn, -} - -impl hyper::rt::Read for ConnectionWrapper { - fn poll_read( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - mut buf: ReadBufCursor<'_>, - ) -> Poll> { - let n = unsafe { - let mut tbuf = tokio::io::ReadBuf::uninit(buf.as_mut()); - match tokio::io::AsyncRead::poll_read(self.project().inner, cx, &mut tbuf) { - Poll::Ready(Ok(())) => tbuf.filled().len(), - other => return other, - } - }; - - unsafe { - buf.advance(n); + Self::Tls { transport, .. } => transport.make_connection(addr).await, } - Poll::Ready(Ok(())) - } -} - -impl AsyncRead for ConnectionWrapper { - #[inline] - fn poll_read( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &mut ReadBuf<'_>, - ) -> Poll> { - Pin::new(&mut self.inner).poll_read(cx, buf) - } -} - -impl hyper::rt::Write for ConnectionWrapper { - fn poll_write( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - Pin::new(&mut self.inner).poll_write(cx, buf) - } - - fn poll_flush( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - Pin::new(&mut self.inner).poll_flush(cx) - } - - fn poll_shutdown( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - Pin::new(&mut self.inner).poll_shutdown(cx) - } -} - -impl AsyncWrite for ConnectionWrapper { - #[inline] - fn poll_write( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - Pin::new(&mut self.inner).poll_write(cx, buf) - } - - #[inline] - fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.inner).poll_flush(cx) - } - - #[inline] - fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.inner).poll_shutdown(cx) - } -} - -impl Connection for ConnectionWrapper { - fn connected(&self) -> Connected { - Connected::new() - } -} - -#[cfg(test)] -mod tests { - use hex::FromHex; - - #[test] - fn test_convert() { - let authority = "2f746d702f7270632e736f636b"; - assert_eq!( - String::from_utf8(Vec::from_hex(authority).unwrap()).unwrap(), - "/tmp/rpc.sock" - ); } } diff --git a/volo/CLAUDE.md b/volo/CLAUDE.md index f218e500..0f1be36c 100644 --- a/volo/CLAUDE.md +++ b/volo/CLAUDE.md @@ -25,6 +25,8 @@ volo/src/ │ ├── random.rs # WeightedRandomBalance │ └── consistent_hash.rs # ConsistentHashBalance (requires RequestHash) │ +├── pool/ # Generic transport pool (Pool, Poolable, Mode::{Unique, Shared}); used by volo-grpc +│ ├── net/ # Network transport layer │ ├── mod.rs # Address enum (Ip, Unix, Shmipc) │ ├── conn.rs # ConnStream, Conn, OwnedReadHalf/OwnedWriteHalf @@ -59,6 +61,10 @@ volo/src/ Unified transport abstraction. `Address` enum supports TCP (`Ip`), Unix sockets (`Unix`), and shared memory (`Shmipc`). `ConnStream` enum wraps all connection types. +### Transport Pool (`pool`) + +`Pool` keyed by peer (usually `net::Address`). `Mode::Unique` transports are checked out exclusively and handed back with `Pooled::reuse`; `Mode::Shared` (multiplexed, e.g. HTTP/2) keeps one transport per key, hands out clones, and dedupes concurrent connects. New transports come from a `UnaryService`; errors surface as `pool::Error`. Idle eviction runs in a task started lazily on first `get`. Currently used by `volo-grpc`; `volo-thrift` and `volo-http` still carry their own copies of the same design and are to be migrated. + ### Hot Restart (`hotrestart`, Unix only) Zero-downtime restarts via Unix Domain Socket. Parent passes listening socket FDs to child process via `SCM_RIGHTS`, then child signals parent to terminate. Global instance: `DEFAULT_HOT_RESTART`. diff --git a/volo/Cargo.toml b/volo/Cargo.toml index 4c8f7b25..de973da6 100644 --- a/volo/Cargo.toml +++ b/volo/Cargo.toml @@ -32,6 +32,7 @@ dashmap.workspace = true faststr.workspace = true futures.workspace = true libc.workspace = true +linked-hash-map.workspace = true metainfo.workspace = true mur3.workspace = true nix = { workspace = true, features = [ diff --git a/volo/src/lib.rs b/volo/src/lib.rs index 6a1e1d36..f2cb927e 100644 --- a/volo/src/lib.rs +++ b/volo/src/lib.rs @@ -12,6 +12,7 @@ pub mod context; pub mod discovery; pub mod loadbalance; pub mod net; +pub mod pool; pub mod util; pub use hack::Unwrap; #[cfg(target_family = "unix")] diff --git a/volo/src/pool/mod.rs b/volo/src/pool/mod.rs new file mode 100644 index 00000000..c7e37b28 --- /dev/null +++ b/volo/src/pool/mod.rs @@ -0,0 +1,924 @@ +//! Generic transport pool shared by the protocol crates. +//! +//! A [`Pool`] keeps transports (connections, or handles multiplexing a connection) per [`Key`], +//! usually the callee [`Address`][crate::net::Address] chosen by service discovery and load +//! balancing. Transports come in two flavours, chosen per [`Pool::get`] with a [`Mode`]: +//! +//! - [`Mode::Unique`]: one request at a time. The caller gets exclusive use of the transport and +//! hands it back with [`Pooled::reuse`] once done; dropping it without `reuse` discards it. +//! - [`Mode::Shared`]: multiplexed. The pool keeps one transport per key and hands out clones (see +//! [`Poolable::reserve`]); concurrent callers hitting an empty key wait for the single in-flight +//! connect instead of each dialing. +//! +//! Transports idle for longer than [`Config::idle_timeout`] are evicted by a background task that +//! is started on first use and stops when the pool is dropped. +//! +//! Protocol crates plug in with two impls: [`Poolable`] on the transport, and a +//! [`UnaryService`] that makes a new one; the pool is agnostic to the error type and reports +//! it back through [`Error::Connect`]. + +use std::{ + collections::{HashMap, HashSet, VecDeque}, + fmt::Debug, + future::Future, + hash::Hash, + ops::{Deref, DerefMut}, + pin::Pin, + sync::{Arc, Mutex, MutexGuard, PoisonError, Weak}, + task::{Context, Poll, ready}, + time::Duration, +}; + +use futures::future::{self, Either}; +use linked_hash_map::LinkedHashMap; +use motore::service::UnaryService; +use pin_project::pin_project; +use tokio::{ + sync::oneshot, + task::JoinHandle, + time::{Instant, Interval, interval}, +}; + +/// Identifies the peer a transport belongs to; usually [`crate::net::Address`]. +pub trait Key: Eq + Hash + Clone + Debug + Unpin + Send + 'static {} + +impl Key for T where T: Eq + Hash + Clone + Debug + Unpin + Send + 'static {} + +/// How transports for a key are handed out, see the [module docs][self]. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum Mode { + /// Exclusive use; returned with [`Pooled::reuse`]. + Unique, + /// Multiplexed; one transport per key, handed out as clones. + Shared, +} + +/// A transport that can live in a [`Pool`]. +pub trait Poolable: Sized { + /// Whether the transport can still be used. Called before handing an idle transport out + /// and before putting one back. + fn reusable(&self) -> impl Future + Send; + + /// Splits the transport into the copy kept by the pool and the one handed to the caller. + /// + /// Shared transports return [`Reservation::Shared`]; the default is exclusive use. + fn reserve(self) -> Reservation { + Reservation::Unique(self) + } + + /// Whether [`Self::reserve`] returns [`Reservation::Shared`]. + fn can_share(&self) -> bool { + false + } + + /// Fast, synchronous checkout for shared transports: a clone if the transport is known to be + /// usable, `None` to fall back to the async [`Self::reusable`] check. + fn try_checkout(&self) -> Option { + None + } +} + +/// Result of [`Poolable::reserve`]. +#[allow(missing_debug_implementations)] +pub enum Reservation { + /// The first copy stays in the pool, the second goes to the caller. + Shared(T, T), + /// The transport is handed to the caller exclusively. + Unique(T), +} + +/// Error returned by [`Pool::get`]. +#[derive(Debug, thiserror::Error)] +pub enum Error { + /// Making a new transport failed. + #[error("failed to make a new transport: {0}")] + Connect(E), + /// Waiting for a transport was given up: the pool was dropped, or the in-flight connect + /// that this caller was waiting on failed. + #[error("waiting for a pooled transport was canceled")] + Canceled, +} + +/// Pool configuration. +#[derive(Clone, Copy, Debug)] +pub struct Config { + max_idle_per_key: usize, + idle_timeout: Duration, +} + +impl Default for Config { + fn default() -> Self { + Config { + max_idle_per_key: 10240, + idle_timeout: Duration::from_secs(15), + } + } +} + +impl Config { + #[must_use] + pub fn new(max_idle_per_key: usize, idle_timeout: Duration) -> Self { + Config { + max_idle_per_key, + idle_timeout, + } + } + + /// Maximum idle transports kept per key. Shared transports are always capped at one. + #[must_use] + pub fn max_idle_per_key(mut self, max_idle_per_key: usize) -> Self { + self.max_idle_per_key = max_idle_per_key; + self + } + + /// How long a transport may sit idle before being evicted; also the eviction check period. + #[must_use] + pub fn idle_timeout(mut self, idle_timeout: Duration) -> Self { + self.idle_timeout = idle_timeout; + self + } +} + +/// A transport pool, cheap to clone; all clones share the same transports. +pub struct Pool { + inner: Arc>>, +} + +impl Clone for Pool { + fn clone(&self) -> Self { + Pool { + inner: self.inner.clone(), + } + } +} + +impl Debug for Pool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Pool").finish_non_exhaustive() + } +} + +impl Pool { + /// Creates an empty pool. Safe to call outside a runtime; the idle-eviction task starts on + /// the first [`Self::get`]. + #[must_use] + pub fn new(cfg: Config) -> Self { + let (tx, rx) = oneshot::channel(); + Pool { + inner: Arc::new(Mutex::new(Inner { + connecting: HashSet::new(), + idle: HashMap::new(), + waiters: HashMap::new(), + idle_timeout: cfg.idle_timeout, + max_idle_per_key: cfg.max_idle_per_key, + idle_task_tx: Some(tx), + _pool_drop_rx: rx, + })), + } + } + + fn lock(&self) -> MutexGuard<'_, Inner> { + self.inner.lock().unwrap_or_else(PoisonError::into_inner) + } + + fn ensure_idle_task(&self, inner: &mut Inner) { + if let Some(tx) = inner.idle_task_tx.take() { + tokio::spawn(IdleTask { + interval: interval(inner.idle_timeout), + inner: Arc::downgrade(&self.inner), + pool_drop_tx: tx, + }); + } + } + + /// Takes the "connecting lock" for `key`. + /// + /// Shared transports allow a single in-flight connect per key, so this returns `None` if + /// one is already running; unique transports never contend. + fn connecting(&self, key: &K, mode: Mode) -> Option> { + match mode { + Mode::Shared => { + let mut inner = self.lock(); + if inner.connecting.insert(key.clone()) { + tracing::trace!("[VOLO] shared connecting for {key:?}"); + Some(Connecting { + key: key.clone(), + pool: WeakOpt::downgrade(&self.inner), + }) + } else { + tracing::trace!("[VOLO] shared connecting already in progress for {key:?}"); + None + } + } + // Never locked, so there is nothing to release on drop. + Mode::Unique => Some(Connecting { + key: key.clone(), + pool: WeakOpt::none(), + }), + } + } + + /// Returns a transport for `key`: an idle one if there is a usable one, otherwise whichever + /// comes first of a transport handed back by another caller and a freshly made one via `mt`. + /// + /// # Errors + /// + /// [`Error::Connect`] carries `mt`'s error when making a new transport fails. + /// [`Error::Canceled`] means this caller was waiting on someone else's connect for a shared + /// key and that connect failed, or the pool was dropped meanwhile. + pub async fn get( + &self, + key: K, + mode: Mode, + mt: MT, + ) -> Result, Error> + where + MT: UnaryService + Send + Sync + 'static, + MT::Error: Send + 'static, + { + let checkout = { + let entry = 'outer: loop { + let entry = 'inner: { + let mut inner = self.lock(); + self.ensure_idle_task(&mut inner); + let idle_timeout = inner.idle_timeout; + + let Some(list) = inner.idle.get_mut(&key) else { + break 'outer None; + }; + + // Fast path: shared transports can be checked out synchronously under the + // lock, which also avoids a race where the list looks empty right after a + // pop and a second transport gets made needlessly. + while list.front().is_some_and(|e| e.inner.can_share()) { + if list[0].expired(idle_timeout) { + list.pop_front(); + continue; + } + if let Some(t) = list[0].inner.try_checkout() { + list[0].idle_at = Instant::now(); + return Ok(self.reuse(&key, t)); + } + // Unknown or broken: fall through to the async `reusable` check. + break; + } + + while let Some(entry) = list.pop_front() { + if entry.expired(idle_timeout) { + tracing::trace!("[VOLO] dropping expired idle transport for {key:?}"); + continue; + } + break 'inner entry; + } + break 'outer None; + }; + // Closed underneath us: drop it and keep looking. + if !entry.inner.reusable().await { + continue; + } + break 'outer Some(entry); + }; + + let mut inner = self.lock(); + if let Some(entry) = entry { + let t = match entry.inner.reserve() { + Reservation::Shared(to_keep, to_return) => { + if let Some(list) = inner.idle.get_mut(&key) { + list.push_back(Idle::new(to_keep)); + } + to_return + } + Reservation::Unique(t) => t, + }; + return Ok(self.reuse(&key, t)); + } + + // Nothing idle: queue as a waiter, then race that against making a new transport. + let (tx, rx) = oneshot::channel(); + let token = inner.waiters.entry(key.clone()).or_default().insert(tx); + Checkout { + key: key.clone(), + pool: self.clone(), + waiter: rx, + token, + clean: true, + } + // lock released here, before any await + }; + + let Some(connecting) = self.connecting(&key, mode) else { + // A shared connect is already in flight for this key; it will serve us as a waiter. + let t = checkout.await.map_err(|_| Error::Canceled)?; + return Ok(self.reuse(&key, t)); + }; + + // The connect runs as its own task so that this caller giving up (an rpc timeout, say) + // does not tear down the transport every other caller to this key is waiting for. + let connect = ConnectTask { + handle: { + let pool = self.clone(); + let key = key.clone(); + tokio::spawn(async move { + let t = mt.call(key).await?; + tracing::debug!("[VOLO] made transport for {:?}", connecting.key); + Ok(pool.pooled(connecting, t)) + }) + }, + // A unique transport nobody is going to use is not worth finishing. + abort_on_drop: mode == Mode::Unique, + }; + + let connect = match future::select(checkout, connect).await { + Either::Left((Ok(t), _connect)) => return Ok(self.reuse(&key, t)), + // The waiter side is gone (pool dropped); our own connect can still deliver. + Either::Left((Err(_), connect)) => connect, + Either::Right((result, _checkout)) => return result, + }; + connect.await + } + + fn pooled(&self, mut connecting: Connecting, t: T) -> Pooled { + let (t, pool) = match t.reserve() { + Reservation::Shared(to_keep, to_return) => { + let mut inner = self.lock(); + inner.put(connecting.key.clone(), to_keep); + inner.connected(&connecting.key); + connecting.pool = WeakOpt::none(); + // The pool keeps its own copy; the caller's clone needs no way back. + (to_return, WeakOpt::none()) + } + Reservation::Unique(t) => (t, WeakOpt::downgrade(&self.inner)), + }; + Pooled::new(connecting.key.clone(), t, pool) + } + + fn reuse(&self, key: &K, t: T) -> Pooled { + tracing::debug!("[VOLO] reusing idle transport for {key:?}"); + // Only unique transports need a way back into the pool. + let pool = if t.can_share() { + WeakOpt::none() + } else { + WeakOpt::downgrade(&self.inner) + }; + Pooled::new(key.clone(), t, pool) + } +} + +/// A transport checked out of a [`Pool`]; derefs to the transport. +/// +/// Unique transports go back to the pool with [`Self::reuse`]; dropping one without it discards +/// the transport. Shared transports need nothing, the pool keeps its own copy. +#[pin_project] +pub struct Pooled { + key: Option, + #[pin] + t: Option, + pool: WeakOpt>>, +} + +impl Pooled { + fn new(key: K, t: T, pool: WeakOpt>>) -> Self { + Pooled { + key: Some(key), + t: Some(t), + pool, + } + } + + /// Hands a unique transport back to the pool, if it is still usable. + pub async fn reuse(mut self) { + let Some(t) = self.t.take() else { return }; + if !t.reusable().await { + return; + } + let Some(key) = self.key.take() else { return }; + if let Some(pool) = self.pool.upgrade() { + pool.lock() + .unwrap_or_else(PoisonError::into_inner) + .put(key, t); + } + } +} + +impl AsRef for Pooled { + fn as_ref(&self) -> &T { + self.t.as_ref().expect("transport already handed back") + } +} + +impl AsMut for Pooled { + fn as_mut(&mut self) -> &mut T { + self.t.as_mut().expect("transport already handed back") + } +} + +impl Deref for Pooled { + type Target = T; + + fn deref(&self) -> &T { + self.as_ref() + } +} + +impl DerefMut for Pooled { + fn deref_mut(&mut self) -> &mut T { + self.as_mut() + } +} + +/// Holds the "connecting lock" for a shared key while a transport is being made. +struct Connecting { + key: K, + pool: WeakOpt>>, +} + +impl Drop for Connecting { + fn drop(&mut self) { + if let Some(pool) = self.pool.upgrade() { + // Never panic in drop. + if let Ok(mut inner) = pool.lock() { + inner.connected(&self.key); + } + } + } +} + +/// Waits for a transport handed back by another caller. +struct Checkout { + key: K, + pool: Pool, + waiter: oneshot::Receiver, + token: usize, + clean: bool, +} + +impl Future for Checkout { + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let out = ready!(Pin::new(&mut self.waiter).poll(cx)); + // The sender was popped from the waiter list to reach us; nothing to clean up. + self.clean = false; + Poll::Ready(out) + } +} + +impl Drop for Checkout { + fn drop(&mut self) { + if self.clean { + tracing::trace!("[VOLO] checkout dropped for {:?}", self.key); + if let Ok(mut inner) = self.pool.inner.lock() { + if let Some(waiters) = inner.waiters.get_mut(&self.key) { + waiters.remove(self.token); + } + } + } + } +} + +/// The spawned connect; optionally aborted when nobody wants its result any more. +struct ConnectTask { + handle: JoinHandle, + abort_on_drop: bool, +} + +impl Future for ConnectTask, E>> { + type Output = Result, Error>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let out = match ready!(Pin::new(&mut self.handle).poll(cx)) { + Ok(Ok(pooled)) => Ok(pooled), + Ok(Err(err)) => Err(Error::Connect(err)), + Err(err) => { + tracing::error!("[VOLO] connect task failed: {err}"); + Err(Error::Canceled) + } + }; + Poll::Ready(out) + } +} + +impl Drop for ConnectTask { + fn drop(&mut self) { + if self.abort_on_drop { + self.handle.abort(); + } + } +} + +/// An optional weak reference, so shared transports can carry "no way back" cheaply. +struct WeakOpt(Option>); + +impl WeakOpt { + fn none() -> Self { + WeakOpt(None) + } + + fn downgrade(arc: &Arc) -> Self { + WeakOpt(Some(Arc::downgrade(arc))) + } + + fn upgrade(&self) -> Option> { + self.0.as_ref().and_then(Weak::upgrade) + } +} + +struct Idle { + inner: T, + idle_at: Instant, +} + +impl Idle { + fn new(inner: T) -> Self { + Idle { + inner, + idle_at: Instant::now(), + } + } + + fn expired(&self, timeout: Duration) -> bool { + self.idle_at.elapsed() > timeout + } +} + +/// FIFO of waiters with O(1) removal by token, for callers that give up. +struct WaiterList { + inner: LinkedHashMap>, + counter: usize, +} + +impl Default for WaiterList { + fn default() -> Self { + Self { + inner: LinkedHashMap::new(), + counter: 0, + } + } +} + +impl WaiterList { + fn pop(&mut self) -> Option> { + self.inner.pop_front().map(|(_, v)| v) + } + + fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + fn insert(&mut self, sender: oneshot::Sender) -> usize { + let token = self.counter; + self.counter = self.counter.wrapping_add(1); + self.inner.insert(token, sender); + token + } + + fn remove(&mut self, token: usize) -> Option> { + self.inner.remove(&token) + } +} + +struct Inner { + /// Keys with a shared connect in flight; guards against dialing a peer twice. + connecting: HashSet, + idle: HashMap>>, + waiters: HashMap>, + idle_timeout: Duration, + max_idle_per_key: usize, + /// Taken by the idle task when it starts; `None` afterwards. + idle_task_tx: Option>, + /// Dropped with the pool, which wakes the idle task so it can stop. + _pool_drop_rx: oneshot::Receiver<()>, +} + +impl Inner { + fn clear_expired(&mut self) { + let timeout = self.idle_timeout; + self.idle.retain(|key, list| { + list.retain(|entry| { + let keep = !entry.expired(timeout); + if !keep { + tracing::trace!("[VOLO] idle task evicting expired transport for {key:?}"); + } + keep + }); + !list.is_empty() + }); + } + + /// Puts a transport back: to the first live waiter, else onto the idle list. + fn put(&mut self, key: K, t: T) { + let mut value = Some(t); + if let Some(waiters) = self.waiters.get_mut(&key) { + while let Some(waiter) = waiters.pop() { + if waiter.is_closed() { + continue; + } + let t = value + .take() + .expect("value is present until a unique send succeeds"); + let to_send = match t.reserve() { + Reservation::Shared(to_keep, to_send) => { + value = Some(to_keep); + to_send + } + Reservation::Unique(t) => t, + }; + match waiter.send(to_send) { + Ok(()) => { + tracing::trace!("[VOLO] put: served a waiter for {key:?}"); + if value.is_none() { + break; + } + } + Err(t) => value = Some(t), + } + } + if waiters.is_empty() { + self.waiters.remove(&key); + } + } + + if let Some(t) = value { + if t.can_share() && self.idle.contains_key(&key) { + tracing::trace!("[VOLO] put: shared transport for {key:?} already idle"); + return; + } + let idle = self.idle.entry(key).or_default(); + if idle.len() < self.max_idle_per_key { + idle.push_back(Idle::new(t)); + } + } + } + + /// A shared connect for `key` finished, successfully or not; release the lock. Any waiters + /// left at this point were waiting on a connect that failed, so they are told to give up. + fn connected(&mut self, key: &K) { + let existed = self.connecting.remove(key); + debug_assert!(existed, "Connecting dropped, key not in pool.connecting"); + self.waiters.remove(key); + } +} + +#[pin_project] +struct IdleTask { + #[pin] + interval: Interval, + inner: Weak>>, + #[pin] + pool_drop_tx: oneshot::Sender<()>, +} + +impl Future for IdleTask { + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let mut this = self.project(); + loop { + if this.pool_drop_tx.as_mut().poll_closed(cx).is_ready() { + tracing::trace!("[VOLO] pool dropped, stopping idle task"); + return Poll::Ready(()); + } + ready!(this.interval.as_mut().poll_tick(cx)); + let Some(inner) = this.inner.upgrade() else { + return Poll::Ready(()); + }; + if let Ok(mut inner) = inner.lock() { + inner.clear_expired(); + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::{ + Arc, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }; + + use super::*; + + #[derive(Clone, Debug)] + struct Conn { + id: usize, + shared: bool, + closed: Arc, + } + + impl Conn { + fn is_open(&self) -> bool { + !self.closed.load(Ordering::SeqCst) + } + } + + impl Poolable for Conn { + async fn reusable(&self) -> bool { + self.is_open() + } + + fn reserve(self) -> Reservation { + if self.shared { + Reservation::Shared(self.clone(), self) + } else { + Reservation::Unique(self) + } + } + + fn can_share(&self) -> bool { + self.shared + } + + fn try_checkout(&self) -> Option { + (self.shared && self.is_open()).then(|| self.clone()) + } + } + + #[derive(Clone)] + struct Maker { + made: Arc, + delay: Duration, + fail: bool, + shared: bool, + } + + impl Maker { + fn new(shared: bool) -> Self { + Maker { + made: Arc::new(AtomicUsize::new(0)), + delay: Duration::ZERO, + fail: false, + shared, + } + } + + fn made(&self) -> usize { + self.made.load(Ordering::SeqCst) + } + } + + impl UnaryService<&'static str> for Maker { + type Response = Conn; + type Error = &'static str; + + async fn call(&self, _key: &'static str) -> Result { + tokio::time::sleep(self.delay).await; + if self.fail { + return Err("boom"); + } + Ok(Conn { + id: self.made.fetch_add(1, Ordering::SeqCst), + shared: self.shared, + closed: Arc::default(), + }) + } + } + + fn pool() -> Pool<&'static str, Conn> { + Pool::new(Config::default()) + } + + #[tokio::test] + async fn unique_is_reused_after_being_handed_back() { + let pool = pool(); + let maker = Maker::new(false); + + let a = pool.get("k", Mode::Unique, maker.clone()).await.unwrap(); + let id = a.id; + a.reuse().await; + let b = pool.get("k", Mode::Unique, maker.clone()).await.unwrap(); + + assert_eq!(b.id, id); + assert_eq!(maker.made(), 1); + } + + #[tokio::test] + async fn unique_held_transports_do_not_block_others() { + let pool = pool(); + let maker = Maker::new(false); + + let a = pool.get("k", Mode::Unique, maker.clone()).await.unwrap(); + let b = pool.get("k", Mode::Unique, maker.clone()).await.unwrap(); + + assert_ne!(a.id, b.id); + assert_eq!(maker.made(), 2); + } + + #[tokio::test] + async fn unique_dropped_without_reuse_is_discarded() { + let pool = pool(); + let maker = Maker::new(false); + + drop(pool.get("k", Mode::Unique, maker.clone()).await.unwrap()); + pool.get("k", Mode::Unique, maker.clone()).await.unwrap(); + + assert_eq!(maker.made(), 2); + } + + #[tokio::test] + async fn shared_concurrent_callers_share_one_connect() { + let pool = pool(); + let mut maker = Maker::new(true); + maker.delay = Duration::from_millis(20); + + let conns = future::join_all((0..16).map(|_| { + let pool = pool.clone(); + let maker = maker.clone(); + async move { pool.get("k", Mode::Shared, maker).await.unwrap().id } + })) + .await; + + assert!(conns.iter().all(|id| *id == conns[0])); + assert_eq!(maker.made(), 1); + } + + #[tokio::test] + async fn shared_connect_survives_the_first_caller_giving_up() { + let pool = pool(); + let mut maker = Maker::new(true); + maker.delay = Duration::from_millis(50); + + let impatient = tokio::time::timeout( + Duration::from_millis(5), + pool.get("k", Mode::Shared, maker.clone()), + ); + let patient = pool.get("k", Mode::Shared, maker.clone()); + let (impatient, patient) = future::join(impatient, patient).await; + + assert!(impatient.is_err(), "the first caller timed out"); + patient.expect("the waiter is served by the connect the first caller started"); + assert_eq!(maker.made(), 1); + } + + #[tokio::test] + async fn shared_closed_transport_is_replaced() { + let pool = pool(); + let maker = Maker::new(true); + + let a = pool.get("k", Mode::Shared, maker.clone()).await.unwrap(); + a.closed.store(true, Ordering::SeqCst); + let b = pool.get("k", Mode::Shared, maker.clone()).await.unwrap(); + + assert_ne!(a.id, b.id); + assert_eq!(maker.made(), 2); + } + + #[tokio::test] + async fn keys_are_independent() { + let pool = pool(); + let maker = Maker::new(true); + + let a = pool.get("a", Mode::Shared, maker.clone()).await.unwrap(); + let b = pool.get("b", Mode::Shared, maker.clone()).await.unwrap(); + let a2 = pool.get("a", Mode::Shared, maker.clone()).await.unwrap(); + + assert_ne!(a.id, b.id); + assert_eq!(a.id, a2.id); + assert_eq!(maker.made(), 2); + } + + #[tokio::test] + async fn connect_failure_is_reported() { + let pool = pool(); + let mut maker = Maker::new(true); + maker.fail = true; + + let err = pool.get("k", Mode::Shared, maker).await.err().unwrap(); + assert!(matches!(err, Error::Connect("boom")), "{err:?}"); + } + + #[tokio::test] + async fn idle_transports_are_evicted() { + let pool: Pool<&'static str, Conn> = + Pool::new(Config::default().idle_timeout(Duration::from_millis(100))); + let maker = Maker::new(true); + + let a = pool.get("k", Mode::Shared, maker.clone()).await.unwrap(); + // Keeping the transport in use refreshes its idle timestamp. + tokio::time::sleep(Duration::from_millis(80)).await; + let a2 = pool.get("k", Mode::Shared, maker.clone()).await.unwrap(); + assert_eq!(a.id, a2.id); + + tokio::time::sleep(Duration::from_millis(250)).await; + assert!( + pool.lock().idle.is_empty(), + "idle task evicted the transport" + ); + let b = pool.get("k", Mode::Shared, maker.clone()).await.unwrap(); + assert_ne!(a.id, b.id); + } + + #[tokio::test] + async fn dropping_the_pool_stops_the_idle_task() { + let pool = pool(); + let inner = Arc::downgrade(&pool.inner); + pool.get("k", Mode::Shared, Maker::new(true)).await.unwrap(); + + drop(pool); + tokio::task::yield_now().await; + + assert!( + inner.upgrade().is_none(), + "the idle task holds no strong reference" + ); + } +}