diff --git a/.gitignore b/.gitignore index 2f19938a..eccbd3d8 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ target /benchmark/output .DS_Store .trae +.codex diff --git a/Cargo.lock b/Cargo.lock index 346ed9b6..a02954b9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -458,7 +458,7 @@ version = "3.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fde0e0ec90c9dfb3b4b1a0891a7dcd0e2bffde2f7efed5fe7c9bb00e5bfb915e" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -670,7 +670,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -758,7 +758,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1574,7 +1574,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2003,7 +2003,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2651,7 +2651,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.59.0", ] [[package]] @@ -3003,7 +3003,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.11.0", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3168,7 +3168,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b55fb86dfd3a2f5f76ea78310a88f96c4ea21a3031f8d212443d56123fd0521" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3282,9 +3282,9 @@ checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" [[package]] name = "shmipc" -version = "0.1.1" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5f842e4e077bb5719fd45b7a6dd42b90c5697c46e852a4bec6415ab57142ea0" +checksum = "e46b66b2fbc3ed49e95bd15d9ff81017cabe482b0ef9de4591df0e9bb2430a53" dependencies = [ "anyhow", "arc-swap", @@ -3296,6 +3296,7 @@ dependencies = [ "nix", "pin-project", "serde", + "serde_json", "thiserror 2.0.17", "tokio", "tracing", @@ -3528,7 +3529,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.1.3", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -4163,7 +4164,7 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "volo" -version = "0.12.3" +version = "0.12.4" dependencies = [ "arc-swap", "async-broadcast", diff --git a/Cargo.toml b/Cargo.toml index 3c62a8ec..4a6481e1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -141,7 +141,7 @@ tokio-native-tls = "0.3" tungstenite = "0.28" tokio-tungstenite = "0.28" -shmipc = "0.1" +shmipc = "0.2.1" [profile.release] opt-level = 3 @@ -154,7 +154,6 @@ panic = 'unwind' incremental = false overflow-checks = false -# [patch.crates-io] # pilota = { git = "https://github.com/cloudwego/pilota.git", branch = "main" } # pilota-build = { git = "https://github.com/cloudwego/pilota.git", branch = "main" } # pilota-thrift-parser = { git = "https://github.com/cloudwego/pilota.git", branch = "main" } diff --git a/examples/src/thrift/shmipc/client.rs b/examples/src/thrift/shmipc/client.rs index 7d5b106a..bee18f06 100644 --- a/examples/src/thrift/shmipc/client.rs +++ b/examples/src/thrift/shmipc/client.rs @@ -1,11 +1,92 @@ use std::sync::LazyLock; -use volo_thrift::client::CallOpt; +use motore::{layer::Layer, service::Service}; +use volo::{context::Context, net::Address}; +use volo_thrift::{ + ClientError, + client::CallOpt, + context::ClientContext, + transport::{DialPlan, SelectedTransport}, +}; + +/// A per-request layer that dials shmipc first and falls back to a UDS/TCP address. +/// +/// It follows the `callee.address()` contract: it sets the primary (shmipc) address on the callee +/// *before* injecting the [`DialPlan`], so `volo-thrift` can validate that the plan primary matches +/// the callee address. The pool key is chosen per-attempt from the actually selected address, so a +/// UDS/TCP fallback transport never pollutes the shmipc key. There is no `.dial_plan(...)` builder; +/// the plan is always injected through the request context. +#[derive(Clone)] +struct ShmipcFallbackLayer { + plan: DialPlan, +} + +impl ShmipcFallbackLayer { + fn new(shmipc_addr: Address, fallback_addr: Address) -> Self { + Self { + plan: DialPlan::with_fallback(shmipc_addr, fallback_addr), + } + } +} + +impl Layer for ShmipcFallbackLayer { + type Service = ShmipcFallbackService; + + fn layer(self, inner: S) -> Self::Service { + ShmipcFallbackService { + inner, + plan: self.plan, + } + } +} + +#[derive(Clone)] +struct ShmipcFallbackService { + inner: S, + plan: DialPlan, +} + +impl Service for ShmipcFallbackService +where + S: Service + Send + Sync + 'static, + Req: Send + 'static, +{ + type Response = S::Response; + type Error = S::Error; + + async fn call(&self, cx: &mut ClientContext, req: Req) -> Result { + // 1. Set the primary address first, then inject the plan (order matters). + cx.rpc_info_mut() + .callee_mut() + .set_address(self.plan.primary().clone()); + cx.extensions_mut().insert(self.plan.clone()); + + let resp = self.inner.call(cx, req).await; + + // The actually selected transport is written back by volo-thrift. + if let Some(selected) = cx.extensions().get::() { + println!( + "selected transport: {} (attempt {})", + selected.address(), + selected.attempt() + ); + } + resp + } +} static CLIENT: LazyLock = LazyLock::new(|| { - let uds_path = std::os::unix::net::SocketAddr::from_pathname("/tmp/hello_test.sock").unwrap(); + let shmipc_path = + std::os::unix::net::SocketAddr::from_pathname("/tmp/hello_test.sock").unwrap(); + let shmipc_addr = Address::from(volo::net::ShmipcAddr(shmipc_path)); + // Fallback to a plain UDS address when shmipc is unavailable. + let fallback_path = + std::os::unix::net::SocketAddr::from_pathname("/tmp/hello_fallback.sock").unwrap(); + let fallback_addr = Address::from(fallback_path); + volo_gen::thrift_gen::hello::HelloServiceClientBuilder::new("hello") - .address(volo::net::ShmipcAddr(uds_path)) + .address(shmipc_addr.clone()) + .layer_outer_front(ShmipcFallbackLayer::new(shmipc_addr, fallback_addr)) .build() }); diff --git a/scripts/clippy-and-test.sh b/scripts/clippy-and-test.sh index 4bde703d..7528cd8e 100644 --- a/scripts/clippy-and-test.sh +++ b/scripts/clippy-and-test.sh @@ -31,6 +31,10 @@ run_clippy() { echo_command cargo clippy -p volo-thrift --no-default-features -- --deny warnings echo_command cargo clippy -p volo-thrift --no-default-features --features multiplex -- --deny warnings echo_command cargo clippy -p volo-thrift --no-default-features --features unsafe-codec -- --deny warnings + if [ "${RUN_SHMIPC}" = "yes" ]; then + echo_command cargo clippy -p volo-thrift --no-default-features --features shmipc -- --deny warnings + echo_command cargo clippy -p volo-thrift --no-default-features --features shmipc,multiplex -- --deny warnings + fi echo_command cargo clippy -p volo-grpc --no-default-features -- --deny warnings echo_command cargo clippy -p volo-grpc --no-default-features --features rustls -- --deny warnings echo_command cargo clippy -p volo-grpc --no-default-features --features native-tls -- --deny warnings @@ -53,6 +57,7 @@ run_clippy() { run_test() { echo_command cargo test -p volo-thrift echo_command cargo test -p volo-thrift --features shmipc + echo_command cargo test -p volo-thrift --features shmipc,multiplex echo_command cargo test -p volo-grpc --features rustls echo_command cargo test -p volo-http --features client,server,http1,query,form,json,tls,cookie,multipart,ws echo_command cargo test -p volo-http --features client,server,http2,query,form,json,tls,cookie,multipart,ws diff --git a/volo-thrift/src/client/mod.rs b/volo-thrift/src/client/mod.rs index c9b232c9..59cab9b3 100644 --- a/volo-thrift/src/client/mod.rs +++ b/volo-thrift/src/client/mod.rs @@ -17,8 +17,6 @@ use motore::{ }; use pilota::thrift::TMessageType; use tokio::time::Duration; -#[cfg(feature = "shmipc")] -use volo::net::shmipc_fallback::ShmipcMakeTransportWithFallback; use volo::{ FastStr, client::WithOptService, @@ -497,42 +495,6 @@ impl ClientBuilder>( - self, - fallback_addr: A, - ) -> ClientBuilder { - let shmipc_addr = self - .address - .expect("Must call .address() before .with_fallback_address()"); - - ClientBuilder { - config: self.config, - pool: self.pool, - caller_name: self.caller_name, - callee_name: self.callee_name, - address: Some(shmipc_addr), - inner_layer: self.inner_layer, - outer_layer: self.outer_layer, - mk_client: self.mk_client, - _marker: PhantomData, - make_transport: ShmipcMakeTransportWithFallback::new( - DefaultMakeTransport::default(), - DefaultMakeTransport::default(), - fallback_addr.into(), - ), - make_codec: self.make_codec, - mk_lb: self.mk_lb, - disable_timeout_layer: self.disable_timeout_layer, - enable_biz_error: self.enable_biz_error, - #[cfg(feature = "multiplex")] - multiplex: self.multiplex, - } - } - #[doc(hidden)] pub fn get_callee_name(&self) -> &FastStr { &self.callee_name diff --git a/volo-thrift/src/codec/default/mod.rs b/volo-thrift/src/codec/default/mod.rs index c9b819f3..d3bbbecf 100644 --- a/volo-thrift/src/codec/default/mod.rs +++ b/volo-thrift/src/codec/default/mod.rs @@ -43,6 +43,9 @@ pub mod framed; pub mod thrift; pub mod ttheader; +#[cfg(feature = "shmipc")] +const SHMIPC_DECODE_BUFFER_CAPACITY: usize = 512; + /// Trait for encoding a [`ThriftMessage`] in place. /// /// [`ZeroCopyEncoder`] tries to encode a message without copying large data taking the advantage @@ -340,22 +343,38 @@ where #[inline] fn make_codec(&self, reader: R, writer: W) -> (Self::Encoder, Self::Decoder) { let (encoder, decoder) = self.make_zero_copy_codec.make_codec(); + + #[cfg(feature = "shmipc")] + let reader = if reader.is_shmipc() { + BufReader::with_capacity(SHMIPC_DECODE_BUFFER_CAPACITY, reader) + } else { + BufReader::new(reader) + }; + #[cfg(not(feature = "shmipc"))] + let reader = BufReader::new(reader); + ( DefaultEncoder { encoder, writer, linked_bytes: LinkedBytes::new(), }, - DefaultDecoder { - decoder, - reader: BufReader::new(reader), - }, + DefaultDecoder { decoder, reader }, ) } } #[cfg(test)] mod tests { + #[cfg(feature = "shmipc")] + use std::{ + cell::RefCell, + collections::VecDeque, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + }; use std::{ io, pin::Pin, @@ -363,7 +382,9 @@ mod tests { }; use bytes::Bytes; - use tokio::io::{AsyncBufRead, AsyncRead, ReadBuf}; + use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; + #[cfg(feature = "shmipc")] + use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt}; use volo::context::RpcInfo; use super::*; @@ -385,6 +406,7 @@ mod tests { enum EofBehavior { EmptyBuffer, UnexpectedEof, + #[cfg(feature = "shmipc")] OtherError, } @@ -400,6 +422,7 @@ mod tests { io::ErrorKind::UnexpectedEof, "unexpected eof", ))), + #[cfg(feature = "shmipc")] EofBehavior::OtherError => Poll::Ready(Err(io::Error::new( io::ErrorKind::ConnectionReset, "connection reset", @@ -416,6 +439,7 @@ mod tests { io::ErrorKind::UnexpectedEof, "unexpected eof", ))), + #[cfg(feature = "shmipc")] EofBehavior::OtherError => Poll::Ready(Err(io::Error::new( io::ErrorKind::ConnectionReset, "connection reset", @@ -431,6 +455,11 @@ mod tests { Ok(tokio::io::Ready::READABLE | tokio::io::Ready::WRITABLE) } + #[cfg(feature = "shmipc")] + fn is_shmipc(&self) -> bool { + self.shmipc_stream.is_some() + } + #[cfg(feature = "shmipc")] fn shmipc_helper(&self) -> volo::net::shmipc::ShmipcHelper { if let Some(stream) = &self.shmipc_stream { @@ -482,6 +511,74 @@ mod tests { } } + #[cfg(feature = "shmipc")] + struct CapabilityReader { + data: Bytes, + pos: usize, + max_read_sizes: VecDeque, + is_shmipc: bool, + read_sizes: Arc>>, + capability_calls: Arc, + helper_calls: Arc, + } + + #[cfg(feature = "shmipc")] + impl CapabilityReader { + fn new(data: Bytes, is_shmipc: bool, max_read_sizes: impl Into>) -> Self { + Self { + data, + pos: 0, + max_read_sizes: max_read_sizes.into(), + is_shmipc, + read_sizes: Arc::default(), + capability_calls: Arc::default(), + helper_calls: Arc::default(), + } + } + } + + #[cfg(feature = "shmipc")] + impl AsyncRead for CapabilityReader { + fn poll_read( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + this.read_sizes.lock().unwrap().push(buf.remaining()); + + if this.pos == this.data.len() { + return Poll::Ready(Ok(())); + } + + let max_read_size = this.max_read_sizes.pop_front().unwrap_or(usize::MAX); + let read_size = max_read_size + .min(buf.remaining()) + .min(this.data.len() - this.pos); + assert!(read_size > 0); + buf.put_slice(&this.data[this.pos..this.pos + read_size]); + this.pos += read_size; + Poll::Ready(Ok(())) + } + } + + #[cfg(feature = "shmipc")] + impl volo::net::ext::AsyncExt for CapabilityReader { + async fn ready(&self, _interest: tokio::io::Interest) -> io::Result { + Ok(tokio::io::Ready::READABLE | tokio::io::Ready::WRITABLE) + } + + fn is_shmipc(&self) -> bool { + self.capability_calls.fetch_add(1, Ordering::Relaxed); + self.is_shmipc + } + + fn shmipc_helper(&self) -> volo::net::shmipc::ShmipcHelper { + self.helper_calls.fetch_add(1, Ordering::Relaxed); + volo::net::shmipc::ShmipcHelper::none() + } + } + /// A writer that records everything written to it, so a test can inspect /// the exact bytes the encoder produced. #[derive(Clone, Default)] @@ -696,6 +793,113 @@ mod tests { ); } + #[cfg(feature = "shmipc")] + async fn encode_ttheader_frames(payloads: &[Bytes]) -> Bytes { + let encoded = Arc::>>::default(); + let writer = RecordingWriter(encoded.clone()); + let (mut encoder, _) = DefaultMakeCodec::default() + .make_codec(CapabilityReader::new(Bytes::new(), false, []), writer); + + for (index, payload) in payloads.iter().enumerate() { + let mut cx = crate::context::ClientContext::new( + index as i32 + 1, + RpcInfo::with_role(volo::context::Role::Client), + pilota::thrift::TMessageType::Call, + ); + let msg = ThriftMessage::mk_client_msg(&cx, payload.clone()); + encoder.encode(&mut cx, msg).await.unwrap(); + } + + let encoded = encoded.lock().unwrap().clone(); + Bytes::from(encoded) + } + + #[cfg(feature = "shmipc")] + #[tokio::test] + async fn test_make_codec_selects_reader_capacity_without_helper_probe() { + for (is_shmipc, expected_capacity) in + [(true, SHMIPC_DECODE_BUFFER_CAPACITY), (false, 8 * 1024)] + { + let reader = CapabilityReader::new(Bytes::new(), is_shmipc, []); + let read_sizes = reader.read_sizes.clone(); + let capability_calls = reader.capability_calls.clone(); + let helper_calls = reader.helper_calls.clone(); + + let (_, mut decoder) = + DefaultMakeCodec::buffered().make_codec(reader, RecordingWriter::default()); + + assert_eq!(capability_calls.load(Ordering::Relaxed), 1); + assert_eq!(helper_calls.load(Ordering::Relaxed), 0); + assert!(decoder.reader.fill_buf().await.unwrap().is_empty()); + assert_eq!(*read_sizes.lock().unwrap(), [expected_capacity]); + assert_eq!(helper_calls.load(Ordering::Relaxed), 0); + } + } + + #[cfg(feature = "shmipc")] + #[test] + fn test_shmipc_capacity_covers_builtin_codec_probes() { + let max_probe = [ + thrift::HEADER_DETECT_LENGTH, + framed::HEADER_DETECT_LENGTH, + ttheader::HEADER_DETECT_LENGTH, + ] + .into_iter() + .max() + .unwrap(); + assert!(max_probe <= SHMIPC_DECODE_BUFFER_CAPACITY); + } + + #[cfg(feature = "shmipc")] + #[tokio::test] + async fn test_shmipc_reader_can_fill_its_full_capacity() { + let data = Bytes::from(vec![0x5a; SHMIPC_DECODE_BUFFER_CAPACITY]); + let reader = CapabilityReader::new(data.clone(), true, []); + let read_sizes = reader.read_sizes.clone(); + let (_, mut decoder) = + DefaultMakeCodec::buffered().make_codec(reader, RecordingWriter::default()); + + let buffered = decoder + .reader + .fill_buf_at_least(SHMIPC_DECODE_BUFFER_CAPACITY) + .await + .unwrap(); + assert_eq!(buffered, data); + assert_eq!(*read_sizes.lock().unwrap(), [SHMIPC_DECODE_BUFFER_CAPACITY]); + } + + #[cfg(feature = "shmipc")] + #[tokio::test] + async fn test_shmipc_capacity_uses_cache_then_direct_read_for_ttheader_frames() { + metainfo::METAINFO + .scope(RefCell::new(metainfo::MetaInfo::default()), async { + let payloads = [Bytes::from(vec![0x31; 2048]), Bytes::from(vec![0x72; 1024])]; + let encoded = encode_ttheader_frames(&payloads).await; + let reader = CapabilityReader::new(encoded, true, [3, 2, 1]); + let read_sizes = reader.read_sizes.clone(); + let helper_calls = reader.helper_calls.clone(); + let (_, mut decoder) = + DefaultMakeCodec::default().make_codec(reader, RecordingWriter::default()); + + for expected in payloads { + let mut cx = crate::context::ServerContext::default(); + let decoded: ThriftMessage = + decoder.decode(&mut cx).await.unwrap().unwrap(); + assert_eq!(decoded.data.unwrap(), expected); + } + + let read_sizes = read_sizes.lock().unwrap(); + assert_eq!(read_sizes[0], SHMIPC_DECODE_BUFFER_CAPACITY); + assert!( + read_sizes + .iter() + .any(|size| *size > SHMIPC_DECODE_BUFFER_CAPACITY) + ); + assert_eq!(helper_calls.load(Ordering::Relaxed), 0); + }) + .await; + } + #[tokio::test] async fn test_decode_empty_buffer_returns_none() { let reader = MockReader { @@ -743,6 +947,9 @@ mod tests { assert!(err.to_string().contains("unexpected eof")); } + #[cfg(feature = "shmipc")] + static SHMIPC_TEST_COUNTER: AtomicUsize = AtomicUsize::new(0); + #[cfg(feature = "shmipc")] struct ShmipcTestEnv { path: std::path::PathBuf, @@ -750,11 +957,17 @@ mod tests { #[cfg(feature = "shmipc")] impl ShmipcTestEnv { + fn next_socket_path() -> std::path::PathBuf { + let id = SHMIPC_TEST_COUNTER.fetch_add(1, Ordering::Relaxed); + std::env::temp_dir().join(format!( + "volo_shmipc_test_{}_{}.sock", + std::process::id(), + id + )) + } + async fn new() -> (Self, volo::net::shmipc::Stream) { - use std::{ - os::unix::net::SocketAddr, - sync::atomic::{AtomicUsize, Ordering}, - }; + use std::os::unix::net::SocketAddr; use motore::service::UnaryService; use volo::net::shmipc::{ @@ -762,29 +975,16 @@ mod tests { addr::{Address, ShmipcMakeTransport}, }; - static COUNTER: AtomicUsize = AtomicUsize::new(0); - let id = COUNTER.fetch_add(1, Ordering::Relaxed); - - let dir = std::env::temp_dir(); - let path = dir.join(format!( - "volo_shmipc_test_{}_{}.sock", - std::process::id(), - id - )); + let path = Self::next_socket_path(); let _ = std::fs::remove_file(&path); let addr_val = SocketAddr::from_pathname(&path).expect("failed to create socket addr"); let addr = Address::from(addr_val); - let addr_clone = addr.clone(); - - tokio::spawn(async move { - if let Ok(mut listener) = Listener::listen(addr_clone, None).await { - while let Ok(_stream) = listener.accept().await {} - } - }); + let mut listener = Listener::listen(addr.clone(), None) + .await + .expect("failed to listen on shmipc socket"); - // Give listener time to start - tokio::time::sleep(std::time::Duration::from_millis(50)).await; + tokio::spawn(async move { while let Ok(_stream) = listener.accept().await {} }); let svc = ShmipcMakeTransport::new(); let stream = svc @@ -794,6 +994,71 @@ mod tests { (Self { path }, stream) } + + async fn new_with_data( + data: Bytes, + write_sizes: Vec, + ) -> ( + Self, + volo::net::shmipc::Stream, + tokio::task::JoinHandle>, + ) { + use std::os::unix::net::SocketAddr; + + use motore::service::UnaryService; + use volo::net::shmipc::{ + Listener, + addr::{Address, ShmipcMakeTransport}, + }; + + let path = Self::next_socket_path(); + let _ = std::fs::remove_file(&path); + + let addr_val = SocketAddr::from_pathname(&path).expect("failed to create socket addr"); + let addr = Address::from(addr_val); + let mut listener = Listener::listen(addr.clone(), None) + .await + .expect("failed to listen on shmipc socket"); + + let server = tokio::spawn(async move { + let mut stream = listener.accept().await?; + let mut open_marker = [0; 1]; + stream.read_exact(&mut open_marker).await?; + let mut offset = 0; + for write_size in write_sizes { + let end = (offset + write_size).min(data.len()); + if end == offset { + continue; + } + stream.write_all(&data[offset..end]).await?; + stream.flush().await?; + offset = end; + tokio::task::yield_now().await; + } + if offset < data.len() { + stream.write_all(&data[offset..]).await?; + stream.flush().await?; + } + stream.shutdown().await?; + Ok(()) + }); + + let svc = ShmipcMakeTransport::new(); + let mut stream = svc + .call(addr) + .await + .expect("failed to connect to shmipc listener"); + stream + .write_all(&[0]) + .await + .expect("failed to open shmipc stream"); + stream + .flush() + .await + .expect("failed to flush shmipc stream open marker"); + + (Self { path }, stream, server) + } } #[cfg(feature = "shmipc")] @@ -803,6 +1068,63 @@ mod tests { } } + #[cfg(all(feature = "shmipc", target_os = "linux"))] + #[tokio::test] + async fn test_builtin_shmipc_capability_propagates_through_split_and_bufreader() { + let (_env, stream) = ShmipcTestEnv::new().await; + let conn: volo::net::conn::Conn = stream.into(); + + assert!(conn.is_shmipc()); + assert!(conn.shmipc_helper().available()); + + let (reader, writer) = conn.stream.into_split(); + assert!(reader.is_shmipc()); + assert!(writer.is_shmipc()); + + let reader = BufReader::new(reader); + assert!(reader.is_shmipc()); + assert!(reader.shmipc_helper().available()); + } + + #[cfg(all(feature = "shmipc", target_os = "linux"))] + #[tokio::test] + async fn test_real_shmipc_decodes_consecutive_ttheader_frames_without_pinning_slices() { + metainfo::METAINFO + .scope(RefCell::new(metainfo::MetaInfo::default()), async { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + let payloads = [Bytes::from(vec![0x19; 2048]), Bytes::from(vec![0xa7; 1024])]; + let encoded = encode_ttheader_frames(&payloads).await; + let (_env, stream, server) = + ShmipcTestEnv::new_with_data(encoded, vec![3, 2, 1, 700, 17]).await; + + let conn: volo::net::conn::Conn = stream.into(); + assert!(conn.is_shmipc()); + let (reader, writer) = conn.stream.into_split(); + let (_encoder, mut decoder) = + DefaultMakeCodec::default().make_codec(reader, writer); + let helper = decoder.shmipc_helper(); + assert!(helper.available()); + + let mut decoded_payloads = Vec::new(); + for expected in &payloads { + let mut cx = crate::context::ServerContext::default(); + let decoded: ThriftMessage = + decoder.decode(&mut cx).await.unwrap().unwrap(); + let decoded = decoded.data.unwrap(); + assert_eq!(&decoded, expected); + decoded_payloads.push(decoded); + } + + server.await.unwrap().unwrap(); + helper.release_read_and_reuse(); + assert_eq!(decoded_payloads, payloads); + }) + .await + .expect("real shmipc decode timed out"); + }) + .await; + } + #[cfg(all(feature = "shmipc", target_os = "linux"))] #[tokio::test] async fn test_decode_unexpected_eof_returns_none_when_shmipc_available() { diff --git a/volo-thrift/src/transport/dial.rs b/volo-thrift/src/transport/dial.rs new file mode 100644 index 00000000..9a048276 --- /dev/null +++ b/volo-thrift/src/transport/dial.rs @@ -0,0 +1,800 @@ +//! Pool-aware transport acquisition with in-order address fallback. +//! +//! [`DialPlan`] describes an ordered list of candidate [`Address`]es to try for a single request. +//! It is injected per-request through [`ClientContext`] extensions. The crate-private +//! `TransportAcquirer` performs the candidate walk *before* the request is sent, so every attempt +//! uses the real address as its pool key and no request object is ever cloned. +//! +//! The candidate walk lives only here; both the pingpong and multiplex clients delegate to +//! `TransportAcquirer::acquire` and keep only their protocol-specific send / EOF / reuse logic. + +use std::{fmt::Write as _, io, sync::Arc}; + +use motore::service::UnaryService; +use volo::{context::Context, net::Address}; + +use crate::{ + ClientError, + context::ClientContext, + transport::pool::{Config, Key, Poolable, Pooled, PooledMakeTransport, Ver}, +}; + +/// An ordered plan of transport addresses to try for a single request. +/// +/// The first address is the primary; any following addresses are fallbacks that are only attempted +/// when acquiring the previous candidate fails. Construct it through [`DialPlan::new`] or +/// [`DialPlan::with_fallback`] and inject it per-request via `ClientContext::extensions_mut`. +/// +/// The internal storage is an [`Arc`]-backed slice, so cloning a `DialPlan` only bumps a reference +/// count and never allocates or clones the request. +#[derive(Clone, Debug)] +pub struct DialPlan { + // Invariant: always contains at least one address, with no consecutive duplicates. + attempts: Arc<[Address]>, +} + +impl DialPlan { + /// Creates a plan that only ever dials `primary`. + pub fn new(primary: Address) -> Self { + Self { + attempts: Arc::from(vec![primary]), + } + } + + /// Creates a plan that dials `primary` first and falls back to `fallback`. + /// + /// If `fallback` equals `primary`, the duplicate is dropped and the plan degrades to a single + /// address. + pub fn with_fallback(primary: Address, fallback: Address) -> Self { + let attempts = if primary == fallback { + vec![primary] + } else { + vec![primary, fallback] + }; + Self { + attempts: Arc::from(attempts), + } + } + + /// Returns the primary (first) address of the plan. + pub fn primary(&self) -> &Address { + // The constructors guarantee at least one address. + &self.attempts[0] + } + + /// Returns an iterator over all candidate addresses in dial order. + pub fn attempts(&self) -> impl ExactSizeIterator { + self.attempts.iter() + } +} + +/// The transport actually selected by the transport acquirer. +/// +/// It is written into the request's `ClientContext` extensions so downstream layers (logging, +/// metrics) can report the real transport instead of the configured one. +#[derive(Clone, Debug)] +pub struct SelectedTransport { + address: Address, + attempt: usize, +} + +impl SelectedTransport { + /// Constructs a `SelectedTransport`. + /// + /// This is only intended for the acquire layer and for downstream tests that need to simulate a + /// selected transport; regular code should read it back from the context extensions. + #[doc(hidden)] + pub fn new(address: Address, attempt: usize) -> Self { + Self { address, attempt } + } + + /// The address that was actually connected. + pub fn address(&self) -> &Address { + &self.address + } + + /// The zero-based index of the successful attempt in the [`DialPlan`]. + pub fn attempt(&self) -> usize { + self.attempt + } +} + +/// A leased transport, either owned directly (shmipc, which manages its own stream pool) or checked +/// out from the outer connection [`Pool`](crate::transport::pool::Pool). +pub(crate) enum TransportLease { + /// A transport that bypasses the outer pool. Used for shmipc, whose `SessionManager` already + /// pools streams; the caller is responsible for stream reuse via the shmipc helper. + #[cfg(feature = "shmipc")] + Direct(T), + /// A transport checked out from the outer pool, keyed by its real address. + Pooled(Pooled), +} + +impl TransportLease { + /// Shared access for protocols whose `send` takes `&self` (multiplex), and for the pingpong + /// shmipc helper reuse path. + #[cfg(any(feature = "shmipc", feature = "multiplex"))] + pub(crate) fn transport(&self) -> &T { + match self { + #[cfg(feature = "shmipc")] + TransportLease::Direct(t) => t, + TransportLease::Pooled(p) => p.as_ref(), + } + } + + /// Exclusive access for protocols whose `send` takes `&mut self` (pingpong). + pub(crate) fn transport_mut(&mut self) -> &mut T { + match self { + #[cfg(feature = "shmipc")] + TransportLease::Direct(t) => t, + TransportLease::Pooled(p) => p.as_mut(), + } + } + + /// Consumes the lease, returning a pooled transport to its pool. + /// + /// A [`TransportLease::Direct`] transport is simply dropped here: shmipc stream reuse is driven + /// by the caller through the shmipc helper before the lease is dropped. + pub(crate) async fn reuse(self) { + match self { + #[cfg(feature = "shmipc")] + TransportLease::Direct(_t) => {} + TransportLease::Pooled(p) => p.reuse().await, + } + } +} + +/// The result of a successful [`TransportAcquirer::acquire`]: a leased transport ready to `send`. +pub(crate) struct AcquiredTransport { + lease: TransportLease, +} + +impl AcquiredTransport { + /// Shared access to the underlying transport. + #[cfg(any(feature = "shmipc", feature = "multiplex"))] + pub(crate) fn transport(&self) -> &T { + self.lease.transport() + } + + /// Exclusive access to the underlying transport. + pub(crate) fn transport_mut(&mut self) -> &mut T { + self.lease.transport_mut() + } + + /// Consumes the acquired transport, returning any pooled connection to its pool. + pub(crate) async fn reuse(self) { + self.lease.reuse().await + } +} + +/// Resolved candidate source for a single request: either the per-request [`DialPlan`] or a single +/// address fast path derived from `callee.address()`. +enum ResolvedPlan { + /// No `DialPlan` in the context: dial the single callee address. + Single(Address), + /// A per-request `DialPlan` from the context extensions. + Plan(DialPlan), +} + +impl ResolvedPlan { + fn len(&self) -> usize { + match self { + ResolvedPlan::Single(_) => 1, + ResolvedPlan::Plan(plan) => plan.attempts.len(), + } + } + + fn candidate(&self, idx: usize) -> &Address { + match self { + ResolvedPlan::Single(addr) => addr, + ResolvedPlan::Plan(plan) => &plan.attempts[idx], + } + } +} + +/// Resolves the candidate source for this request. +/// +/// The `callee.address()` contract is enforced by the caller before this is invoked; here we only +/// validate that a context `DialPlan` (if any) agrees with the already-read `target`. +fn resolve_dial_plan(cx: &ClientContext, target: &Address) -> Result { + match cx.extensions().get::() { + Some(plan) => { + if plan.primary() != target { + return Err(ClientError::Transport( + io::Error::new( + io::ErrorKind::InvalidData, + format!( + "DialPlan primary ({}) does not match callee address ({target})", + plan.primary(), + ), + ) + .into(), + )); + } + // Cloning only bumps the Arc refcount; ends the borrow of `cx` immediately. + Ok(ResolvedPlan::Plan(plan.clone())) + } + None => Ok(ResolvedPlan::Single(target.clone())), + } +} + +fn missing_address_error(cx: &ClientContext) -> ClientError { + ClientError::Transport( + io::Error::new( + io::ErrorKind::InvalidData, + format!("address is required, rpc_info: {:?}", cx.rpc_info()), + ) + .into(), + ) +} + +/// Acquires a transport for a single request, walking the [`DialPlan`] candidates in order. +/// +/// The same generic type is monomorphized separately for the pingpong and multiplex makers; the two +/// clients never share a maker instance, only this acquire algorithm. +pub(crate) struct TransportAcquirer +where + MT: UnaryService
, + MT::Response: Poolable, +{ + // Direct maker used to bypass the outer pool for shmipc candidates. + #[cfg(feature = "shmipc")] + direct: MT, + // Pool-aware maker used for UDS/TCP candidates, keyed by the real address. + pooled: PooledMakeTransport, +} + +impl Clone for TransportAcquirer +where + MT: UnaryService
+ Clone, + MT::Response: Poolable, +{ + fn clone(&self) -> Self { + Self { + #[cfg(feature = "shmipc")] + direct: self.direct.clone(), + pooled: self.pooled.clone(), + } + } +} + +impl TransportAcquirer +where + MT: UnaryService
+ Send + Sync + Clone + 'static, + MT::Response: Poolable + Send + 'static, + MT::Error: Into + Send, +{ + /// Builds an acquirer from a maker and an optional pool config. + pub(crate) fn new(direct: MT, pool_cfg: Option) -> Self { + #[cfg(feature = "shmipc")] + { + Self { + pooled: PooledMakeTransport::new(direct.clone(), pool_cfg), + direct, + } + } + #[cfg(not(feature = "shmipc"))] + { + Self { + pooled: PooledMakeTransport::new(direct, pool_cfg), + } + } + } + + /// Acquires a transport, trying each [`DialPlan`] candidate until one succeeds. + /// + /// This records the make-transport start/end stats, rewrites `callee.address` to the actually + /// attempted address, and writes a [`SelectedTransport`] into the context extensions. It never + /// touches the request object. + pub(crate) async fn acquire( + &self, + cx: &mut ClientContext, + ver: Ver, + ) -> Result, ClientError> { + // A retry layer may reuse this context across attempts. Clear any SelectedTransport left by + // a previous acquire so that if this acquire fails entirely, downstream layers (metrics, + // logging) do not report the previous success's transport. It is re-inserted below only + // when a candidate succeeds. + cx.extensions_mut().remove::(); + + // `callee.address()` is mandatory and represents the primary transport target. + let target = cx + .rpc_info() + .callee() + .address() + .ok_or_else(|| missing_address_error(cx))?; + + // Resolve candidates and validate any context plan; this ends the borrow of `cx`. + let resolved = resolve_dial_plan(cx, &target)?; + + cx.stats.record_make_transport_start_at(); + + let n = resolved.len(); + let mut last_err: Option = None; + let mut summary = String::new(); + + for attempt in 0..n { + let candidate = resolved.candidate(attempt).clone(); + // Point the callee at the address we are about to dial, so metrics/errors reflect the + // real attempt. The borrow ends before the `await` below. + cx.rpc_info_mut() + .callee_mut() + .set_address(candidate.clone()); + + // Only time fallback attempts. The healthy `attempt == 0` success path (including a + // shmipc primary that succeeds while a fallback is configured) reads no clock at all; a + // first-attempt failure derives its latency from `make_transport_start_at` in the cold + // log branch below. + let fallback_start = (attempt > 0).then(std::time::Instant::now); + + match self.acquire_candidate(candidate.clone(), ver).await { + Ok(Some(lease)) => { + cx.stats.record_make_transport_end_at(); + if let Some(started) = fallback_start { + // A fallback candidate succeeded after earlier attempts failed/skipped. + // Record the transition so fallback usage is observable even on success. + tracing::debug!( + attempt, + candidate = %candidate, + elapsed_us = started.elapsed().as_micros(), + "[VOLO] dial plan fell back to a later candidate" + ); + } + cx.extensions_mut() + .insert(SelectedTransport::new(candidate, attempt)); + return Ok(AcquiredTransport { lease }); + } + Ok(None) => { + // multiplex + shmipc candidate: skip and try the next candidate. Emit a + // tracing event now, because the skip reason is lost once a later candidate + // succeeds and the local summary is dropped. + tracing::debug!( + attempt, + candidate = %candidate, + "[VOLO] dial plan skipped candidate: shmipc does not support multiplex" + ); + let _ = write!( + summary, + "; attempt {attempt} ({candidate}): skipped (shmipc does not support \ + multiplex)" + ); + } + Err(e) => { + // For a real multi-candidate plan, record each failed attempt (with its + // acquire latency): without it, a primary failure that is later masked by a + // successful fallback would leave no trace. For a plain single-address request + // (`n == 1`) we stay silent and let the error propagate, preserving the + // pre-existing log behavior and avoiding duplicate failure logs during outages. + if n > 1 { + // Reuse `make_transport_start_at` for the first attempt so the hot path + // never pays for an extra `Instant`; later attempts use their own timer. + let elapsed_us = match fallback_start { + Some(started) => started.elapsed().as_micros(), + None => cx + .stats + .make_transport_start_at() + .map(|start| { + (chrono::Local::now() - start) + .num_microseconds() + .unwrap_or(0) + .max(0) as u128 + }) + .unwrap_or(0), + }; + tracing::debug!( + attempt, + candidate = %candidate, + elapsed_us, + error = %e, + "[VOLO] dial plan candidate failed" + ); + } + let _ = write!(summary, "; attempt {attempt} ({candidate}): {e}"); + last_err = Some(e); + } + } + } + + // No candidate succeeded. `make_transport_end_at` stays unset to preserve the existing + // "never acquired a transport" semantics. + match last_err { + Some(mut e) => { + e.append_msg(&format!(", dial plan exhausted{summary}")); + Err(e) + } + // Every candidate was skipped: a multiplex plan containing only shmipc candidates. + None => Err(ClientError::Transport( + io::Error::new( + io::ErrorKind::Unsupported, + "shmipc does not support multiplex", + ) + .into(), + )), + } + } + + /// Acquires a single candidate. + /// + /// Returns `Ok(Some(lease))` on success, `Ok(None)` if the candidate must be skipped (a shmipc + /// candidate under multiplex), or `Err(_)` if acquiring the candidate failed. + async fn acquire_candidate( + &self, + candidate: Address, + ver: Ver, + ) -> Result>, ClientError> { + #[cfg(feature = "shmipc")] + if candidate.is_shmipc() { + return match ver { + // shmipc bypasses the outer pool: its SessionManager already pools streams. + Ver::PingPong => self + .direct + .call(candidate) + .await + .map(|t| Some(TransportLease::Direct(t))) + .map_err(Into::into), + // shmipc does not support multiplex; skip so a non-shmipc fallback can be used. + Ver::Multiplex => Ok(None), + }; + } + + self.pooled + .call((candidate, ver)) + .await + .map(|p| Some(TransportLease::Pooled(p))) + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::HashSet, + net::SocketAddr, + sync::{Arc, Mutex}, + }; + + use motore::service::UnaryService; + use volo::{ + context::{Context, Role, RpcInfo}, + net::Address, + }; + + use super::{DialPlan, SelectedTransport, TransportAcquirer}; + use crate::{ + ClientError, + context::{ClientContext, Config}, + protocol::TMessageType, + transport::pool::{Poolable, Ver}, + }; + + fn ip(port: u16) -> Address { + Address::Ip(SocketAddr::from(([127, 0, 0, 1], port))) + } + + #[cfg(feature = "shmipc")] + fn shmipc(port: u16) -> Address { + Address::Shmipc(volo::net::shmipc::Address::Tcp(SocketAddr::from(( + [127, 0, 0, 1], + port, + )))) + } + + fn make_cx(address: Option
) -> ClientContext { + let mut cx = ClientContext::new( + 1, + RpcInfo::::with_role(Role::Client), + TMessageType::Call, + ); + if let Some(addr) = address { + cx.rpc_info_mut().callee_mut().set_address(addr); + } + cx + } + + fn io_kind(e: &ClientError) -> Option { + match e { + ClientError::Transport(te) => Some(te.io_error().kind()), + _ => None, + } + } + + /// Extracts the error from an acquire result whose `Ok` type is not `Debug`. + fn expect_err(res: Result) -> ClientError { + match res { + Ok(_) => panic!("expected an error, got Ok"), + Err(e) => e, + } + } + + /// A minimal poolable transport for exercising the acquire algorithm. + struct MockTransport; + + impl Poolable for MockTransport { + async fn reusable(&self) -> bool { + true + } + } + + /// A mock maker that records every address it is asked to dial and fails for a configured set. + /// The fail set is mutable so a test can change it between successive acquires on one context. + #[derive(Clone)] + struct MockMaker { + calls: Arc>>, + fail: Arc>>, + } + + impl MockMaker { + fn new(fail: impl IntoIterator) -> Self { + Self { + calls: Arc::new(Mutex::new(Vec::new())), + fail: Arc::new(Mutex::new(fail.into_iter().collect())), + } + } + + fn calls(&self) -> Vec
{ + self.calls.lock().unwrap().clone() + } + + /// Replaces the set of addresses that will fail on subsequent dials. + fn set_fail(&self, fail: impl IntoIterator) { + *self.fail.lock().unwrap() = fail.into_iter().collect(); + } + } + + impl UnaryService
for MockMaker { + type Response = MockTransport; + type Error = std::io::Error; + + async fn call(&self, addr: Address) -> Result { + self.calls.lock().unwrap().push(addr.clone()); + if self.fail.lock().unwrap().contains(&addr) { + Err(std::io::Error::new( + std::io::ErrorKind::ConnectionRefused, + format!("mock fail for {addr}"), + )) + } else { + Ok(MockTransport) + } + } + } + + // ---- DialPlan / SelectedTransport pure logic ---- + + #[test] + fn dial_plan_single() { + let plan = DialPlan::new(ip(1)); + assert_eq!(plan.primary(), &ip(1)); + assert_eq!(plan.attempts().len(), 1); + assert_eq!(plan.attempts().cloned().collect::>(), vec![ip(1)]); + } + + #[test] + fn dial_plan_with_fallback() { + let plan = DialPlan::with_fallback(ip(1), ip(2)); + assert_eq!(plan.primary(), &ip(1)); + assert_eq!( + plan.attempts().cloned().collect::>(), + vec![ip(1), ip(2)] + ); + } + + #[test] + fn dial_plan_dedup_identical_fallback() { + let plan = DialPlan::with_fallback(ip(1), ip(1)); + assert_eq!(plan.attempts().len(), 1); + assert_eq!(plan.primary(), &ip(1)); + } + + #[test] + fn selected_transport_accessors() { + let selected = SelectedTransport::new(ip(7), 3); + assert_eq!(selected.address(), &ip(7)); + assert_eq!(selected.attempt(), 3); + } + + // ---- acquire: validation ---- + + #[tokio::test] + async fn acquire_missing_address_is_invalid_data() { + let maker = MockMaker::new([]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(None); + // Even with a plan present, a missing callee address must fail fast. + cx.extensions_mut().insert(DialPlan::new(ip(1))); + + let err = expect_err(acquirer.acquire(&mut cx, Ver::PingPong).await); + assert_eq!(io_kind(&err), Some(std::io::ErrorKind::InvalidData)); + assert!(format!("{err}").contains("address is required")); + assert!(maker.calls().is_empty(), "maker must not be called"); + } + + #[tokio::test] + async fn acquire_plan_primary_mismatch_is_invalid_data() { + let maker = MockMaker::new([]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(ip(1))); + // Plan primary (ip(2)) disagrees with callee address (ip(1)). + cx.extensions_mut() + .insert(DialPlan::with_fallback(ip(2), ip(3))); + + let err = expect_err(acquirer.acquire(&mut cx, Ver::PingPong).await); + assert_eq!(io_kind(&err), Some(std::io::ErrorKind::InvalidData)); + assert!(maker.calls().is_empty(), "no dial before validation"); + } + + // ---- acquire: candidate walk ---- + + #[tokio::test] + async fn acquire_single_fast_path_without_plan() { + let maker = MockMaker::new([]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(ip(1))); + + acquirer.acquire(&mut cx, Ver::PingPong).await.unwrap(); + + assert_eq!(maker.calls(), vec![ip(1)]); + let selected = cx.extensions().get::().unwrap(); + assert_eq!(selected.address(), &ip(1)); + assert_eq!(selected.attempt(), 0); + assert!(cx.stats.make_transport_start_at().is_some()); + assert!(cx.stats.make_transport_end_at().is_some()); + } + + #[tokio::test] + async fn acquire_primary_success_skips_fallback() { + let maker = MockMaker::new([]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(ip(1))); + cx.extensions_mut() + .insert(DialPlan::with_fallback(ip(1), ip(2))); + + acquirer.acquire(&mut cx, Ver::PingPong).await.unwrap(); + + assert_eq!(maker.calls(), vec![ip(1)], "fallback must not be dialed"); + let selected = cx.extensions().get::().unwrap(); + assert_eq!(selected.address(), &ip(1)); + assert_eq!(selected.attempt(), 0); + } + + #[tokio::test] + async fn acquire_falls_back_on_primary_failure() { + let maker = MockMaker::new([ip(1)]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(ip(1))); + cx.extensions_mut() + .insert(DialPlan::with_fallback(ip(1), ip(2))); + + acquirer.acquire(&mut cx, Ver::PingPong).await.unwrap(); + + assert_eq!(maker.calls(), vec![ip(1), ip(2)]); + let selected = cx.extensions().get::().unwrap(); + assert_eq!(selected.address(), &ip(2)); + assert_eq!(selected.attempt(), 1); + // callee address is rewritten to the actually selected fallback. + assert_eq!(cx.rpc_info().callee().address(), Some(ip(2))); + assert!(cx.stats.make_transport_end_at().is_some()); + } + + #[tokio::test] + async fn acquire_all_fail_aggregates_errors() { + let maker = MockMaker::new([ip(1), ip(2)]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(ip(1))); + cx.extensions_mut() + .insert(DialPlan::with_fallback(ip(1), ip(2))); + + let err = expect_err(acquirer.acquire(&mut cx, Ver::PingPong).await); + let msg = format!("{err}"); + assert!( + msg.contains("127.0.0.1:1"), + "first attempt in summary: {msg}" + ); + assert!( + msg.contains("127.0.0.1:2"), + "second attempt in summary: {msg}" + ); + // callee address reflects the last attempt. + assert_eq!(cx.rpc_info().callee().address(), Some(ip(2))); + // end stays unset on total failure. + assert!(cx.stats.make_transport_start_at().is_some()); + assert!(cx.stats.make_transport_end_at().is_none()); + } + + #[tokio::test] + async fn acquire_failure_after_prior_success_clears_stale_selection() { + // Model a retry layer reusing one context: first acquire succeeds on shmipc, the second + // fails on every candidate. The stale shmipc SelectedTransport from the first attempt must + // not linger, or downstream metrics would mislabel the failed UDS attempt as shmipc. + let maker = MockMaker::new([]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(ip(1))); + cx.extensions_mut() + .insert(DialPlan::with_fallback(ip(1), ip(2))); + + // First acquire: primary succeeds. + acquirer.acquire(&mut cx, Ver::PingPong).await.unwrap(); + assert_eq!( + cx.extensions() + .get::() + .map(|s| s.address().clone()), + Some(ip(1)) + ); + + // Second acquire on the same context: both candidates now fail. + maker.set_fail([ip(1), ip(2)]); + let _ = expect_err(acquirer.acquire(&mut cx, Ver::PingPong).await); + + // The stale selection is gone, so downstream reads fall back to the last attempted address. + assert!( + cx.extensions().get::().is_none(), + "stale SelectedTransport must be cleared on a fully-failed acquire" + ); + assert_eq!(cx.rpc_info().callee().address(), Some(ip(2))); + } + + #[tokio::test] + async fn context_reset_clears_plan_and_selected() { + let maker = MockMaker::new([]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(ip(1))); + cx.extensions_mut().insert(DialPlan::new(ip(1))); + acquirer.acquire(&mut cx, Ver::PingPong).await.unwrap(); + assert!(cx.extensions().get::().is_some()); + + cx.reset(2, TMessageType::Call); + + assert!(cx.extensions().get::().is_none()); + assert!(cx.extensions().get::().is_none()); + } + + // ---- shmipc / multiplex specific behavior ---- + + #[cfg(feature = "shmipc")] + #[tokio::test] + async fn acquire_multiplex_skips_shmipc_candidate() { + let maker = MockMaker::new([]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(shmipc(1))); + cx.extensions_mut() + .insert(DialPlan::with_fallback(shmipc(1), ip(2))); + + acquirer.acquire(&mut cx, Ver::Multiplex).await.unwrap(); + + // shmipc candidate is skipped without ever calling the maker. + assert_eq!(maker.calls(), vec![ip(2)]); + let selected = cx.extensions().get::().unwrap(); + assert_eq!(selected.address(), &ip(2)); + assert_eq!(selected.attempt(), 1); + } + + #[cfg(feature = "shmipc")] + #[tokio::test] + async fn acquire_multiplex_only_shmipc_is_unsupported() { + let maker = MockMaker::new([]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(shmipc(1))); + cx.extensions_mut().insert(DialPlan::new(shmipc(1))); + + let err = expect_err(acquirer.acquire(&mut cx, Ver::Multiplex).await); + assert_eq!(io_kind(&err), Some(std::io::ErrorKind::Unsupported)); + assert!(format!("{err}").contains("shmipc does not support multiplex")); + assert!(maker.calls().is_empty(), "shmipc maker must not be called"); + } + + #[cfg(feature = "shmipc")] + #[tokio::test] + async fn acquire_pingpong_shmipc_uses_direct() { + let maker = MockMaker::new([]); + let acquirer = TransportAcquirer::new(maker.clone(), None); + let mut cx = make_cx(Some(shmipc(1))); + cx.extensions_mut() + .insert(DialPlan::with_fallback(shmipc(1), ip(2))); + + acquirer.acquire(&mut cx, Ver::PingPong).await.unwrap(); + + // pingpong dials the shmipc primary directly and succeeds without touching the fallback. + assert_eq!(maker.calls(), vec![shmipc(1)]); + let selected = cx.extensions().get::().unwrap(); + assert_eq!(selected.address(), &shmipc(1)); + assert_eq!(selected.attempt(), 0); + } +} diff --git a/volo-thrift/src/transport/mod.rs b/volo-thrift/src/transport/mod.rs index 1167c9d2..9d77da20 100644 --- a/volo-thrift/src/transport/mod.rs +++ b/volo-thrift/src/transport/mod.rs @@ -1,8 +1,10 @@ +pub mod dial; pub(crate) mod incoming; #[cfg(feature = "multiplex")] pub mod multiplex; pub mod pingpong; pub mod pool; +pub use dial::{DialPlan, SelectedTransport}; use pilota::thrift::ThriftException; pub use pool::Config; diff --git a/volo-thrift/src/transport/multiplex/client.rs b/volo-thrift/src/transport/multiplex/client.rs index d251e52c..b57cbb01 100644 --- a/volo-thrift/src/transport/multiplex/client.rs +++ b/volo-thrift/src/transport/multiplex/client.rs @@ -9,8 +9,9 @@ use crate::{ context::ClientContext, protocol::TMessageType, transport::{ + dial::TransportAcquirer, multiplex::thrift_transport::ThriftTransport, - pool::{Config, PooledMakeTransport, Ver}, + pool::{Config, Ver}, }, }; @@ -85,8 +86,7 @@ where MkC: MakeCodec + Sync, Resp: EntryMessage + Send + 'static, { - #[allow(clippy::type_complexity)] - make_transport: PooledMakeTransport, Address>, + acquirer: TransportAcquirer>, _marker: PhantomData, } @@ -98,7 +98,7 @@ where { fn clone(&self) -> Self { Self { - make_transport: self.make_transport.clone(), + acquirer: self.acquirer.clone(), _marker: self._marker, } } @@ -112,9 +112,9 @@ where { pub fn new(make_transport: MkT, pool_cfg: Option, make_codec: MkC) -> Self { let make_transport = MakeClientTransport::new(make_transport, make_codec); - let make_transport = PooledMakeTransport::new(make_transport, pool_cfg); + let acquirer = TransportAcquirer::new(make_transport, pool_cfg); Client { - make_transport, + acquirer, _marker: PhantomData, } } @@ -136,16 +136,11 @@ where cx: &mut ClientContext, req: ThriftMessage, ) -> Result { - let rpc_info = &cx.rpc_info; - let target = rpc_info.callee().address().ok_or_else(|| { - let msg = format!("address is required, rpcinfo: {rpc_info:?}"); - ClientError::Transport(io::Error::new(io::ErrorKind::InvalidData, msg).into()) - })?; let oneway = cx.message_type == TMessageType::OneWay; - cx.stats.record_make_transport_start_at(); - let transport = self.make_transport.call((target, Ver::Multiplex)).await?; - cx.stats.record_make_transport_end_at(); - let resp = transport.send(cx, req, oneway).await; + // Acquire happens before `send`; the candidate walk and stats live in the acquirer. + // A multiplex acquire always yields a pooled lease (shmipc candidates are skipped). + let acquired = self.acquirer.acquire(cx, Ver::Multiplex).await?; + let resp = acquired.transport().send(cx, req, oneway).await; if let Ok(None) = resp { if !oneway { return Err(ClientError::Transport( @@ -157,7 +152,7 @@ where } } if cx.transport.should_reuse && resp.is_ok() { - transport.reuse().await; + acquired.reuse().await; } resp } diff --git a/volo-thrift/src/transport/pingpong/client.rs b/volo-thrift/src/transport/pingpong/client.rs index d500cdd6..ce0a957e 100644 --- a/volo-thrift/src/transport/pingpong/client.rs +++ b/volo-thrift/src/transport/pingpong/client.rs @@ -1,7 +1,6 @@ use std::{io, marker::PhantomData}; use motore::service::{Service, UnaryService}; -use pilota::thrift::TransportException; use volo::net::{Address, dial::MakeTransport}; use crate::{ @@ -10,8 +9,9 @@ use crate::{ context::ClientContext, protocol::TMessageType, transport::{ + dial::TransportAcquirer, pingpong::thrift_transport::ThriftTransport, - pool::{Config, PooledMakeTransport, Ver}, + pool::{Config, Ver}, }, }; @@ -61,8 +61,7 @@ where MkT: MakeTransport, MkC: MakeCodec + Sync, { - #[allow(clippy::type_complexity)] - make_transport: PooledMakeTransport, Address>, + acquirer: TransportAcquirer>, _marker: PhantomData, } @@ -73,7 +72,7 @@ where { fn clone(&self) -> Self { Self { - make_transport: self.make_transport.clone(), + acquirer: self.acquirer.clone(), _marker: self._marker, } } @@ -86,9 +85,9 @@ where { pub fn new(make_transport: MkT, pool_cfg: Option, make_codec: MkC) -> Self { let make_transport = MakeClientTransport::new(make_transport, make_codec); - let make_transport = PooledMakeTransport::new(make_transport, pool_cfg); + let acquirer = TransportAcquirer::new(make_transport, pool_cfg); Client { - make_transport, + acquirer, _marker: PhantomData, } } @@ -111,27 +110,19 @@ where cx: &mut ClientContext, req: ThriftMessage, ) -> Result { - let rpc_info = &cx.rpc_info; - let target = rpc_info.callee().address().ok_or_else(|| { - TransportException::from(io::Error::new( - io::ErrorKind::InvalidData, - format!("address is required, rpc_info: {rpc_info:?}"), - )) - })?; let oneway = cx.message_type == TMessageType::OneWay; - cx.stats.record_make_transport_start_at(); - let mut transport = self.make_transport.call((target, Ver::PingPong)).await?; - cx.stats.record_make_transport_end_at(); - let resp = transport.send(cx, req, oneway).await; + // Acquire happens before `send`; the candidate walk and stats live in the acquirer. + let mut acquired = self.acquirer.acquire(cx, Ver::PingPong).await?; + // `send` may be cancelled by an outer timeout. Keep a clone of the shmipc stream in a + // guard so cancellation closes it instead of leaving it in the session map forever. + #[cfg(feature = "shmipc")] + let mut shmipc_close_guard = { + let helper = acquired.transport().shmipc_helper(); + helper.available().then(|| helper.close_guard()) + }; + let resp = acquired.transport_mut().send(cx, req, oneway).await; if let Ok(None) = resp { if !oneway { - #[cfg(feature = "shmipc")] - { - let helper = transport.shmipc_helper(); - if helper.available() { - helper.reuse().await; - } - } return Err(crate::ClientError::Transport( pilota::thrift::TransportException::from(io::Error::new( io::ErrorKind::UnexpectedEof, @@ -143,21 +134,299 @@ where )); } } - // if shmipc enabled and is shmipc: close + // Shmipc manages its own stream pool; recycle only completed calls. #[cfg(feature = "shmipc")] { - let helper = transport.shmipc_helper(); + let helper = acquired.transport().shmipc_helper(); if helper.available() { - helper.reuse().await; + // A failed send can leave data in shmipc's send buffer (for example after a + // persistent QueueFull). Only a completed RPC is safe to recycle; on error the + // armed guard closes the stream and reclaims its buffers. + if resp.is_ok() { + helper.reuse().await; + if let Some(guard) = shmipc_close_guard.take() { + guard.disarm(); + } + } } else if cx.transport.should_reuse && resp.is_ok() { - transport.reuse().await; + acquired.reuse().await; } } #[cfg(not(feature = "shmipc"))] if cx.transport.should_reuse && resp.is_ok() { - transport.reuse().await; + acquired.reuse().await; } resp } } + +#[cfg(all(test, feature = "shmipc", target_os = "linux"))] +mod tests { + use std::{ + cell::RefCell, + io, + os::unix::net::SocketAddr, + path::PathBuf, + sync::{Arc, Mutex}, + time::Duration, + }; + + use motore::service::{Service, UnaryService}; + use pilota::thrift::ThriftException; + use tokio::{ + io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, + time::timeout, + }; + use volo::{ + context::{Context, Role, RpcInfo}, + net::{ + Address, + conn::{OwnedReadHalf, OwnedWriteHalf}, + dial::MakeTransport, + ext::AsyncExt, + shmipc::{Listener, addr::ShmipcMakeTransport}, + }, + }; + + use super::Client; + use crate::{ + Bytes, EntryMessage, ThriftMessage, + codec::{Decoder, DefaultMakeCodec, Encoder, MakeCodec}, + context::{ClientContext, Config, ThriftContext}, + protocol::TMessageType, + }; + + #[derive(Clone)] + struct OneShotTransport { + halves: Arc>>, + } + + impl MakeTransport for OneShotTransport { + type ReadHalf = OwnedReadHalf; + type WriteHalf = OwnedWriteHalf; + + async fn make_transport( + &self, + _addr: Address, + ) -> io::Result<(Self::ReadHalf, Self::WriteHalf)> { + self.halves.lock().unwrap().take().ok_or_else(|| { + io::Error::new(io::ErrorKind::NotConnected, "transport already taken") + }) + } + + fn set_connect_timeout(&mut self, _timeout: Option) {} + + fn set_read_timeout(&mut self, _timeout: Option) {} + + fn set_write_timeout(&mut self, _timeout: Option) {} + } + + struct SocketCleanup(PathBuf); + + impl Drop for SocketCleanup { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + } + } + + #[derive(Clone)] + struct WriteThenFailMakeCodec; + + struct WriteThenFailEncoder { + writer: W, + } + + impl Encoder for WriteThenFailEncoder + where + W: AsyncWrite + AsyncExt + Unpin + Send + Sync + 'static, + { + async fn encode( + &mut self, + _cx: &mut Cx, + _msg: ThriftMessage, + ) -> Result<(), ThriftException> { + self.writer.write_all(b"buffered but unsent").await?; + Err(io::Error::other("injected send failure").into()) + } + + fn shmipc_helper(&self) -> volo::net::shmipc::ShmipcHelper { + self.writer.shmipc_helper() + } + } + + struct UnreachableDecoder { + reader: R, + } + + impl Decoder for UnreachableDecoder + where + R: AsyncRead + AsyncExt + Unpin + Send + Sync + 'static, + { + async fn decode( + &mut self, + _cx: &mut Cx, + ) -> Result>, ThriftException> { + unreachable!("send failure must bypass decode") + } + + fn shmipc_helper(&self) -> volo::net::shmipc::ShmipcHelper { + self.reader.shmipc_helper() + } + } + + impl MakeCodec for WriteThenFailMakeCodec + where + R: AsyncRead + AsyncExt + Unpin + Send + Sync + 'static, + W: AsyncWrite + AsyncExt + Unpin + Send + Sync + 'static, + { + type Encoder = WriteThenFailEncoder; + type Decoder = UnreachableDecoder; + + fn make_codec(&self, reader: R, writer: W) -> (Self::Encoder, Self::Decoder) { + ( + WriteThenFailEncoder { writer }, + UnreachableDecoder { reader }, + ) + } + } + + #[tokio::test] + async fn failed_call_does_not_reuse_shmipc_stream_with_unsent_data() { + let path = std::env::temp_dir().join(format!( + "volo_failed_shmipc_call_{}.sock", + std::process::id() + )); + let _cleanup = SocketCleanup(path.clone()); + let _ = std::fs::remove_file(&path); + let shmipc_addr = volo::net::shmipc::Address::from( + SocketAddr::from_pathname(&path).expect("valid socket path"), + ); + let _listener = Listener::listen(shmipc_addr.clone(), None) + .await + .expect("listen on shmipc socket"); + + let stream = ShmipcMakeTransport::new() + .call(shmipc_addr.clone()) + .await + .expect("connect shmipc stream"); + let failed_stream_addr = stream.peer_addr(); + let (read_half, write_half) = stream.into_split(); + let halves = Arc::new(Mutex::new(Some(( + OwnedReadHalf::Shmipc(read_half), + OwnedWriteHalf::Shmipc(write_half), + )))); + let client: Client = + Client::new(OneShotTransport { halves }, None, WriteThenFailMakeCodec); + + let mut cx = ClientContext::new( + 1, + RpcInfo::::with_role(Role::Client), + TMessageType::Call, + ); + cx.rpc_info_mut() + .callee_mut() + .set_address(Address::Shmipc(shmipc_addr.clone())); + let req = ThriftMessage::mk_client_msg(&cx, Bytes::new()); + + let result = metainfo::METAINFO + .scope( + RefCell::new(metainfo::MetaInfo::default()), + client.call(&mut cx, req), + ) + .await; + assert!(result.is_err(), "the injected send failure must propagate"); + + let next_stream = ShmipcMakeTransport::new() + .call(shmipc_addr) + .await + .expect("connect another shmipc stream"); + assert_ne!( + next_stream.peer_addr(), + failed_stream_addr, + "a stream with unsent data must not return to the session pool" + ); + next_stream + .helper() + .close() + .await + .expect("close replacement stream"); + } + + #[tokio::test] + async fn cancelled_call_closes_shmipc_stream() { + let path = std::env::temp_dir().join(format!( + "volo_cancelled_shmipc_call_{}.sock", + std::process::id() + )); + let _cleanup = SocketCleanup(path.clone()); + let _ = std::fs::remove_file(&path); + let shmipc_addr = volo::net::shmipc::Address::from( + SocketAddr::from_pathname(&path).expect("valid socket path"), + ); + let mut listener = Listener::listen(shmipc_addr.clone(), None) + .await + .expect("listen on shmipc socket"); + let mut accept = tokio::spawn(async move { listener.accept().await }); + + let stream = ShmipcMakeTransport::new() + .call(shmipc_addr.clone()) + .await + .expect("connect shmipc stream"); + let (read_half, write_half) = stream.into_split(); + let halves = Arc::new(Mutex::new(Some(( + OwnedReadHalf::Shmipc(read_half), + OwnedWriteHalf::Shmipc(write_half), + )))); + let client: Client = Client::new( + OneShotTransport { halves }, + None, + DefaultMakeCodec::default(), + ); + + let mut cx = ClientContext::new( + 1, + RpcInfo::::with_role(Role::Client), + TMessageType::Call, + ); + cx.rpc_info_mut() + .callee_mut() + .set_address(Address::Shmipc(shmipc_addr)); + let req = ThriftMessage::mk_client_msg(&cx, Bytes::new()); + + let mut call = Box::pin(metainfo::METAINFO.scope( + RefCell::new(metainfo::MetaInfo::default()), + client.call(&mut cx, req), + )); + let mut peer = timeout(Duration::from_secs(1), async { + tokio::select! { + result = &mut call => panic!("RPC completed before cancellation: {result:?}"), + accepted = &mut accept => accepted + .expect("accept task should not panic") + .expect("accept shmipc stream"), + } + }) + .await + .expect("server should receive the request"); + + assert!( + timeout(Duration::ZERO, call).await.is_err(), + "the RPC must still be waiting for a response when it is cancelled" + ); + + let mut request = Vec::new(); + let read_result = timeout(Duration::from_secs(1), peer.read_to_end(&mut request)) + .await + .expect("cancelled client should close the stream"); + if let Err(err) = read_result { + // Older shmipc releases report their EndOfStream marker as UnexpectedEof. + assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + } + assert!( + !request.is_empty(), + "server should receive the encoded request" + ); + + peer.helper().close().await.expect("close server stream"); + } +} diff --git a/volo/Cargo.toml b/volo/Cargo.toml index 7afbec53..4c8f7b25 100644 --- a/volo/Cargo.toml +++ b/volo/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "volo" -version = "0.12.3" +version = "0.12.4" edition.workspace = true homepage.workspace = true repository.workspace = true diff --git a/volo/src/net/dial.rs b/volo/src/net/dial.rs index f1ba1878..63176a3d 100644 --- a/volo/src/net/dial.rs +++ b/volo/src/net/dial.rs @@ -16,6 +16,17 @@ use super::{ }; /// [`MakeTransport`] creates an [`AsyncRead`] and an [`AsyncWrite`] for the given [`Address`]. +/// +/// # Single-address contract +/// +/// An implementation MUST connect to exactly the [`Address`] it is given, or return an error. +/// It MUST NOT silently redial a different address (for example, falling back from a shmipc +/// address to a UDS/TCP address inside `make_transport`). +/// +/// This invariant is required by the connection pool: `volo-thrift` builds the pool key from the +/// address *before* calling `make_transport`, so redialing another address here would insert the +/// resulting transport under the wrong key. Address fallback must instead be performed by the +/// acquire layer, which can pick the correct pool key for each attempt. pub trait MakeTransport: Clone + Send + Sync + 'static { type ReadHalf: AsyncRead + Send + Sync + Unpin + 'static; type WriteHalf: AsyncWrite + Send + Sync + Unpin + 'static; @@ -165,3 +176,48 @@ impl UnaryService
for DefaultMakeTransport { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::net::Address; + + /// The single-address contract: `DefaultMakeTransport` connects to exactly the address it is + /// given and returns the error directly on failure. It never silently redials another address. + #[tokio::test] + async fn default_make_transport_tcp_failure_is_reported() { + let mkt = DefaultMakeTransport::default(); + // An address with no listener; connecting must fail rather than fall back elsewhere. + let addr = Address::Ip("127.0.0.1:1".parse().unwrap()); + let res = mkt.make_transport(addr).await; + assert!(res.is_err(), "connecting to a closed TCP port must fail"); + } + + #[cfg(target_family = "unix")] + #[tokio::test] + async fn default_make_transport_uds_failure_is_reported() { + let mkt = DefaultMakeTransport::default(); + // A path that does not exist; connecting must fail rather than fall back elsewhere. + let addr = Address::Unix( + std::os::unix::net::SocketAddr::from_pathname("/tmp/volo-nonexistent-dial-test.sock") + .unwrap(), + ); + let res = mkt.make_transport(addr).await; + assert!(res.is_err(), "connecting to a missing UDS must fail"); + } + + #[cfg(feature = "shmipc")] + #[tokio::test] + async fn default_make_transport_shmipc_failure_is_reported() { + let mkt = DefaultMakeTransport::default(); + // An unroutable shmipc TCP address; the shmipc maker must surface the error, not fall back. + let addr = Address::Shmipc(crate::net::shmipc::Address::Tcp( + "127.0.0.1:1".parse().unwrap(), + )); + let res = mkt.make_transport(addr).await; + assert!( + res.is_err(), + "connecting to an unreachable shmipc must fail" + ); + } +} diff --git a/volo/src/net/ext.rs b/volo/src/net/ext.rs index e63a0808..e6536ee8 100644 --- a/volo/src/net/ext.rs +++ b/volo/src/net/ext.rs @@ -11,6 +11,12 @@ pub trait AsyncExt { /// See [`tokio::net::TcpStream::ready`] for details. fn ready(&self, interest: Interest) -> impl Future> + Send; + /// Returns whether this IO object is backed by ShmIPC. + #[cfg(feature = "shmipc")] + fn is_shmipc(&self) -> bool { + false + } + /// Get helper of ShmIPC. #[cfg(feature = "shmipc")] fn shmipc_helper(&self) -> super::shmipc::ShmipcHelper { @@ -37,6 +43,11 @@ impl AsyncExt for Conn { } } + #[cfg(feature = "shmipc")] + fn is_shmipc(&self) -> bool { + self.stream.is_shmipc() + } + #[cfg(feature = "shmipc")] fn shmipc_helper(&self) -> super::shmipc::ShmipcHelper { match &self.stream { @@ -65,6 +76,11 @@ impl AsyncExt for OwnedReadHalf { } } + #[cfg(feature = "shmipc")] + fn is_shmipc(&self) -> bool { + matches!(self, OwnedReadHalf::Shmipc(_)) + } + #[cfg(feature = "shmipc")] fn shmipc_helper(&self) -> super::shmipc::ShmipcHelper { match self { @@ -93,6 +109,11 @@ impl AsyncExt for OwnedWriteHalf { } } + #[cfg(feature = "shmipc")] + fn is_shmipc(&self) -> bool { + matches!(self, OwnedWriteHalf::Shmipc(_)) + } + #[cfg(feature = "shmipc")] fn shmipc_helper(&self) -> super::shmipc::ShmipcHelper { match self { diff --git a/volo/src/net/shmipc/addr.rs b/volo/src/net/shmipc/addr.rs index 5464cc4f..77535e04 100644 --- a/volo/src/net/shmipc/addr.rs +++ b/volo/src/net/shmipc/addr.rs @@ -1,14 +1,44 @@ #[cfg(target_os = "linux")] use std::os::linux::net::SocketAddrExt; -use std::{collections::HashMap, fmt, hash::Hash, io, sync::LazyLock}; +use std::{ + collections::HashMap, + fmt, + hash::Hash, + io, + sync::{Arc, LazyLock}, +}; use motore::service::UnaryService; use shmipc::session::SessionManager; -use tokio::sync::RwLock; +use tokio::sync::{OnceCell, RwLock}; -pub(crate) static SESSION_MANAGERS: LazyLock>>> = +type SessionManagerCell = Arc>>; + +pub(crate) static SESSION_MANAGERS: LazyLock>> = LazyLock::new(Default::default); +async fn get_or_insert_cell( + cells: &RwLock>>>, + key: K, +) -> Arc> +where + K: Eq + Hash, +{ + { + let read = cells.read().await; + if let Some(cell) = read.get(&key) { + return Arc::clone(cell); + } + } + + let mut write = cells.write().await; + Arc::clone( + write + .entry(key) + .or_insert_with(|| Arc::new(OnceCell::new())), + ) +} + #[derive(Clone, Debug)] pub enum Address { // The address must be loopback addr @@ -195,21 +225,60 @@ impl UnaryService
for ShmipcMakeTransport { )); } - { - let read = SESSION_MANAGERS.read().await; - if let Some(sm) = read.get(&addr) { - return sm.get_stream().map(super::Stream::new).map_err(Into::into); - } - } + let cell = get_or_insert_cell(&SESSION_MANAGERS, addr.clone()).await; + let sm = cell + .get_or_try_init(|| async { + let config = super::config::session_manager_config(); + tracing::debug!("ShmipcMakeTransport: config: {config:?}"); + SessionManager::new(config, Connector, addr) + .await + .map_err(Into::::into) + }) + .await?; + + sm.get_stream().map(super::Stream::new).map_err(Into::into) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use tokio::sync::Barrier; + + use super::*; - let config = super::config::session_manager_config(); - tracing::debug!("ShmipcMakeTransport: config: {config:?}"); - let sm = SessionManager::new(config, Connector, addr.clone()) - .await - .map_err(Into::::into)?; - let ret = sm.get_stream().map(super::Stream::new).map_err(Into::into); - SESSION_MANAGERS.write().await.insert(addr, sm); + #[tokio::test] + async fn same_key_initialization_is_singleflight() { + const CONCURRENCY: usize = 32; - ret + let cells = Arc::new(RwLock::new(HashMap::>>::new())); + let barrier = Arc::new(Barrier::new(CONCURRENCY)); + let initializations = Arc::new(AtomicUsize::new(0)); + let mut tasks = Vec::with_capacity(CONCURRENCY); + + for _ in 0..CONCURRENCY { + let cells = Arc::clone(&cells); + let barrier = Arc::clone(&barrier); + let initializations = Arc::clone(&initializations); + tasks.push(tokio::spawn(async move { + barrier.wait().await; + let cell = get_or_insert_cell(&cells, 1).await; + *cell + .get_or_try_init(|| async { + initializations.fetch_add(1, Ordering::Relaxed); + tokio::task::yield_now().await; + Ok::<_, ()>(42) + }) + .await + .unwrap() + })); + } + + for task in tasks { + assert_eq!(task.await.unwrap(), 42); + } + assert_eq!(initializations.load(Ordering::Relaxed), 1); + assert_eq!(cells.read().await.len(), 1); } } diff --git a/volo/src/net/shmipc/helpers.rs b/volo/src/net/shmipc/helpers.rs index e980e7ba..e05e001f 100644 --- a/volo/src/net/shmipc/helpers.rs +++ b/volo/src/net/shmipc/helpers.rs @@ -64,6 +64,13 @@ pub struct ShmipcCloseGuard { inner: Option>, } +impl ShmipcCloseGuard { + /// Prevent the guard from closing the stream after it has been explicitly recycled. + pub fn disarm(mut self) { + self.inner = None; + } +} + impl Drop for ShmipcCloseGuard { fn drop(&mut self) { if let Some(mut stream) = self.inner.take() { diff --git a/volo/src/net/shmipc_fallback.rs b/volo/src/net/shmipc_fallback.rs index 2199e1a6..8c72bab6 100644 --- a/volo/src/net/shmipc_fallback.rs +++ b/volo/src/net/shmipc_fallback.rs @@ -1,11 +1,6 @@ use std::io; -use super::{ - Address, DefaultIncoming, MakeIncoming, - conn::{Conn, OwnedReadHalf, OwnedWriteHalf}, - dial::{DefaultMakeTransport, MakeTransport}, - incoming::Incoming, -}; +use super::{Address, DefaultIncoming, MakeIncoming, conn::Conn, incoming::Incoming}; pub struct ShmipcAddressWithFallback { pub shmipc_addr: Address, @@ -60,63 +55,3 @@ where } } } - -#[derive(Clone, Debug)] -pub struct ShmipcMakeTransportWithFallback { - pub shmipc_mkt: DefaultMakeTransport, - pub default_mkt: DefaultMakeTransport, - pub fallback_addr: Address, -} - -impl ShmipcMakeTransportWithFallback { - pub fn new( - shmipc: DefaultMakeTransport, - default_mkt: DefaultMakeTransport, - fallback_addr: Address, - ) -> Self { - Self { - shmipc_mkt: shmipc, - default_mkt, - fallback_addr, - } - } -} - -impl MakeTransport for ShmipcMakeTransportWithFallback { - type ReadHalf = OwnedReadHalf; - type WriteHalf = OwnedWriteHalf; - - async fn make_transport( - &self, - mut addr: Address, - ) -> io::Result<(Self::ReadHalf, Self::WriteHalf)> { - if addr.is_shmipc() { - match self.shmipc_mkt.make_transport(addr).await { - Ok(ret) => return Ok(ret), - Err(e) => { - tracing::info!( - "failed to connect to shmipc target: {e}, fallback to default target" - ); - addr = self.fallback_addr.clone(); - } - } - } - - self.default_mkt.make_transport(addr).await - } - - fn set_connect_timeout(&mut self, timeout: Option) { - self.default_mkt.set_connect_timeout(timeout); - self.shmipc_mkt.set_connect_timeout(timeout); - } - - fn set_read_timeout(&mut self, timeout: Option) { - self.default_mkt.set_read_timeout(timeout); - self.shmipc_mkt.set_read_timeout(timeout); - } - - fn set_write_timeout(&mut self, timeout: Option) { - self.default_mkt.set_write_timeout(timeout); - self.shmipc_mkt.set_write_timeout(timeout); - } -} diff --git a/volo/src/util/buf_reader.rs b/volo/src/util/buf_reader.rs index 0aae72a1..9d8d6003 100644 --- a/volo/src/util/buf_reader.rs +++ b/volo/src/util/buf_reader.rs @@ -192,6 +192,12 @@ impl AsyncBufRead for BufReader { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); + // Always return buffered data before polling the underlying reader. In particular, an + // EOF or error from the underlying reader must not hide data that has already been read. + if *me.pos < *me.len { + return Poll::Ready(Ok(&me.buf[*me.pos..*me.len])); + } + // If we've reached the end of our internal buffer then we need to fetch // some more data from the underlying reader. // Branch using `>=` instead of the more correct `==` @@ -203,17 +209,8 @@ impl AsyncBufRead for BufReader { *me.len = buf.filled().len(); *me.pos = 0; } else if *me.len < *me.cap { - // We have some buffer let mut buf = ReadBuf::new(&mut me.buf[*me.len..*me.cap]); - match me.inner.poll_read(cx, &mut buf) { - Poll::Ready(t) => t, - Poll::Pending => { - if *me.pos < *me.len { - return Poll::Ready(Ok(&me.buf[*me.pos..*me.len])); - } - return Poll::Pending; - } - }?; + ready!(me.inner.poll_read(cx, &mut buf))?; *me.len += buf.filled().len(); } Poll::Ready(Ok(&me.buf[*me.pos..*me.len])) @@ -230,6 +227,11 @@ impl AsyncExt for BufReader { self.inner.ready(interest).await } + #[cfg(feature = "shmipc")] + fn is_shmipc(&self) -> bool { + self.inner.is_shmipc() + } + #[cfg(feature = "shmipc")] fn shmipc_helper(&self) -> crate::net::shmipc::ShmipcHelper { self.inner.shmipc_helper() @@ -250,8 +252,37 @@ impl fmt::Debug for BufReader { #[cfg(test)] mod tests { + use tokio::io::AsyncBufReadExt; + use super::*; + struct DataThenError { + data: Option>, + reads: usize, + } + + impl AsyncRead for DataThenError { + fn poll_read( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + this.reads += 1; + + if let Some(data) = this.data.take() { + assert!(buf.remaining() >= data.len()); + buf.put_slice(&data); + Poll::Ready(Ok(())) + } else { + Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "stream closed", + ))) + } + } + } + #[cfg(test)] fn is_unpin() {} @@ -302,4 +333,40 @@ mod tests { assert_eq!(buf.len, 5); assert_eq!(buf.buf[buf.pos..buf.len], [6, 7, 8, 9, 10]); } + + #[tokio::test] + async fn test_buffered_ttheader_data_is_delivered_before_close_error() { + // Simulate a complete 1123-byte fallback TTHeader event immediately followed by close. + const FRAME_LEN: usize = 1123; + const PAYLOAD_LEN: usize = FRAME_LEN - 4; + + let mut frame = vec![0x5a; FRAME_LEN]; + frame[..4].copy_from_slice(&(PAYLOAD_LEN as u32).to_be_bytes()); + frame[4..6].copy_from_slice(&[0x10, 0x00]); + let expected_payload = frame[4..].to_vec(); + + let inner = DataThenError { + data: Some(frame), + reads: 0, + }; + let mut reader = BufReader::new(inner); + + let header = reader.fill_buf_at_least(6).await.unwrap(); + assert_eq!( + u32::from_be_bytes(header[..4].try_into().unwrap()) as usize, + PAYLOAD_LEN + ); + assert_eq!(&header[4..6], &[0x10, 0x00]); + + reader.consume(4); + let mut payload = vec![0; PAYLOAD_LEN]; + reader.read_exact(&mut payload).await.unwrap(); + + assert_eq!(payload, expected_payload); + assert_eq!(reader.get_ref().reads, 1); + + let err = reader.read_u8().await.unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + assert_eq!(reader.get_ref().reads, 2); + } }