diff --git a/Cargo.lock b/Cargo.lock index eb81a62..b59eb98 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -493,6 +493,36 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "example-proto-spawn" +version = "0.1.0" +dependencies = [ + "env_logger", + "futures", + "histogram", + "log", + "prost", + "protosocket", + "protosocket-prost", + "protosocket-rpc", + "tokio", +] + +[[package]] +name = "example-proto-spawn-stream" +version = "0.1.0" +dependencies = [ + "env_logger", + "futures", + "histogram", + "log", + "prost", + "protosocket", + "protosocket-prost", + "protosocket-rpc", + "tokio", +] + [[package]] name = "example-proto-stream" version = "0.1.0" @@ -1430,6 +1460,7 @@ name = "protosocket-rpc" version = "1.1.0" dependencies = [ "ahash", + "bytes", "futures", "level-runtime", "log", diff --git a/Cargo.toml b/Cargo.toml index 45d6398..53d0688 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,8 @@ members = [ "example-messagepack", "example-pool", "example-proto", + "example-proto-spawn", + "example-proto-spawn-stream", "example-proto-stream", "example-proto-tls", "example-telnet", diff --git a/example-messagepack/src/server.rs b/example-messagepack/src/server.rs index 99f5821..f01b7aa 100644 --- a/example-messagepack/src/server.rs +++ b/example-messagepack/src/server.rs @@ -1,11 +1,11 @@ use std::sync::atomic::AtomicUsize; -use futures::Stream; +use futures::{stream::BoxStream, Stream, StreamExt}; use messages::{EchoRequest, EchoResponse, EchoStream, Request, Response, ResponseBehavior}; use protosocket::{PooledEncoder, StreamWithAddress, TcpSocketListener}; use protosocket_rpc::{ - server::{ConnectionService, RpcResponder, SocketService}, - Message, ProtosocketControlCode, + server::{ConnectionService, RpcKind, SocketService}, + ProtosocketControlCode, }; use tokio::net::TcpStream; @@ -55,20 +55,17 @@ async fn run_main() -> Result<(), Box> { /// ConnectionServices to application-wide state tracking. struct DemoRpcSocketService; impl SocketService for DemoRpcSocketService { - type Codec = ( - // Use a pooled encoder to amortize memory allocation cost. - // Each connection gets its own little memory pool. - PooledEncoder>, - protosocket_messagepack::MessagePackDecoder, - ); type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&self) -> Self::Codec { + fn codec(&mut self) -> ::Codec { Default::default() } - fn new_stream_service(&self, stream: &StreamWithAddress) -> Self::ConnectionService { + fn new_stream_service( + &mut self, + stream: &StreamWithAddress, + ) -> Self::ConnectionService { log::info!("new connection server {}", stream.address()); DemoRpcConnectionServer { address: stream.address(), @@ -82,32 +79,36 @@ struct DemoRpcConnectionServer { address: std::net::SocketAddr, } impl ConnectionService for DemoRpcConnectionServer { + type Codec = ( + // Use a pooled encoder to amortize memory allocation cost. + // Each connection gets its own little memory pool. + PooledEncoder>, + protosocket_messagepack::MessagePackDecoder, + ); type Request = Request; type Response = Response; + type UnaryFutureType = futures::future::Ready; + type StreamType = BoxStream<'static, Response>; fn new_rpc( &mut self, initiating_message: Self::Request, - responder: RpcResponder<'_, Self::Response>, - ) { + ) -> RpcKind { log::debug!("{} new rpc: {initiating_message:?}", self.address); let request_id = initiating_message.request_id; let behavior = initiating_message.response_behavior; match initiating_message.body { Some(echo) => match behavior { ResponseBehavior::Unary => { - responder.immediate(echo_request(request_id, echo)); + RpcKind::Unary(futures::future::ready(echo_request(request_id, echo))) } ResponseBehavior::Stream => { - tokio::spawn(responder.stream(echo_stream(request_id, echo))); + RpcKind::Streaming(echo_stream(request_id, echo).boxed()) } }, None => { - // No completion messages will be sent for this message - log::warn!( - "{request_id} no request in rpc body. This may cause a client memory leak." - ); - responder.immediate(Response::cancelled(request_id)); + log::warn!("{request_id} no request in rpc body"); + RpcKind::Cancelled } } } diff --git a/example-pool/src/server.rs b/example-pool/src/server.rs index 81c5f75..917987c 100644 --- a/example-pool/src/server.rs +++ b/example-pool/src/server.rs @@ -28,12 +28,12 @@ impl ServerConnector for ServerContext { type Codec = ByteBufferRingCodec; type Reactor = PooledReactor; - fn codec(&self) -> Self::Codec { + fn codec(&mut self) -> Self::Codec { ByteBufferRingCodec::default() } fn new_reactor( - &self, + &mut self, optional_outbound: spillway::Sender<::Message>, _address: &StreamWithAddress, ) -> Self::Reactor { @@ -46,7 +46,6 @@ impl ServerConnector for ServerContext { &self, connection: protosocket::Connection< ::Stream, - Self::Codec, Self::Reactor, >, ) { @@ -58,6 +57,7 @@ struct PooledReactor { outbound: spillway::Sender, } impl MessageReactor for PooledReactor { + type Codec = ByteBufferRingCodec; type Inbound = String; type Outbound = String; type LogicalOutbound = String; diff --git a/example-proto-spawn-stream/Cargo.toml b/example-proto-spawn-stream/Cargo.toml new file mode 100644 index 0000000..83638b7 --- /dev/null +++ b/example-proto-spawn-stream/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "example-proto-spawn-stream" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "example-proto-spawn-stream-server" +path = "src/server.rs" + +[[bin]] +name = "example-proto-spawn-stream-client" +path = "src/client.rs" + +[dependencies] +protosocket = { workspace = true } +protosocket-prost = { workspace = true } +protosocket-rpc = { workspace = true } + +env_logger = { workspace = true } +futures = { workspace = true } +log = { workspace = true } +prost = { workspace = true } +tokio = { workspace = true, features = ["full"] } + +histogram = { version = "0.11" } diff --git a/example-proto-spawn-stream/src/client.rs b/example-proto-spawn-stream/src/client.rs new file mode 100644 index 0000000..a7f5482 --- /dev/null +++ b/example-proto-spawn-stream/src/client.rs @@ -0,0 +1,189 @@ +use std::{ + sync::{ + atomic::{AtomicU64, AtomicUsize, Ordering}, + Arc, + }, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use futures::StreamExt; +use messages::{EchoRequest, Request, Response}; +use protosocket::PooledEncoder; +use protosocket_prost::{ProstDecoder, ProstSerializer}; +use protosocket_rpc::{ + client::{Configuration, RpcClient, TcpStreamConnector}, + ProtosocketControlCode, +}; + +mod messages; + +fn main() -> Result<(), Box> { + static I: AtomicUsize = AtomicUsize::new(0); + let runtime = tokio::runtime::Builder::new_multi_thread() + .thread_name_fn(|| { + format!( + "app-{}", + I.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + ) + }) + .worker_threads(4) + .event_interval(7) + .enable_all() + .build()?; + + runtime.block_on(run_main()) +} + +async fn run_main() -> Result<(), Box> { + env_logger::init(); + + let response_count = Arc::new(AtomicUsize::new(0)); + let latency = Arc::new(histogram::AtomicHistogram::new(7, 52).expect("histogram works")); + + let concurrency = 256; + let connections = 4; + + let concurrent_count = Arc::new(AtomicUsize::new(0)); + let request_ids = Arc::new(AtomicU64::new(1)); + let mut configuration = Configuration::new(TcpStreamConnector); + configuration.max_queued_outbound_messages(64); + for _i in 0..connections { + let (client, connection) = protosocket_rpc::client::connect::< + ( + PooledEncoder>, + ProstDecoder, + ), + _, + >( + std::env::var("ENDPOINT") + .unwrap_or_else(|_| "127.0.0.1:9000".to_string()) + .parse() + .expect("must use a valid socket address"), + &configuration, + ) + .await?; + let _connection_handle = tokio::spawn(connection); + let tasks = concurrency / connections; + for _ in 0..tasks { + let _client_handle = tokio::spawn(generate_traffic( + concurrent_count.clone(), + request_ids.clone(), + client.clone(), + response_count.clone(), + latency.clone(), + )); + } + } + + let metrics = tokio::spawn(print_periodic_metrics( + response_count, + latency, + concurrent_count, + )); + + tokio::select!( + _ = metrics => { + log::warn!("metrics runtime quit"); + } + ); + + Ok(()) +} + +async fn print_periodic_metrics( + response_count: Arc, + latency: Arc, + concurrent_count: Arc, +) { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + loop { + let start = Instant::now(); + interval.tick().await; + let total = response_count.swap(0, std::sync::atomic::Ordering::Relaxed); + let hz = (total as f64) / start.elapsed().as_secs_f64().max(0.1); + + let latency = latency.drain(); + let p90 = latency + .percentile(0.9) + .unwrap_or_default() + .map(|b| *b.range().end()) + .unwrap_or_default() as f64 + / 1000.0; + let p999 = latency + .percentile(0.999) + .unwrap_or_default() + .map(|b| *b.range().end()) + .unwrap_or_default() as f64 + / 1000.0; + let p9999 = latency + .percentile(0.9999) + .unwrap_or_default() + .map(|b| *b.range().end()) + .unwrap_or_default() as f64 + / 1000.0; + let concurrent = concurrent_count.load(std::sync::atomic::Ordering::Relaxed); + eprintln!( + "Messages: {total:10} rate: {hz:9.1}hz p90: {p90:6.1}µs p999: {p999:6.1}µs p9999: {p9999:6.1}µs concurrency: {concurrent}" + ); + } +} + +async fn generate_traffic( + concurrent_count: Arc, + request_ids: Arc, + client: RpcClient, + metrics_count: Arc, + metrics_latency: Arc, +) { + log::debug!("running traffic generator"); + loop { + let i = request_ids.fetch_add(1, Ordering::Relaxed); + let request = Request { + request_id: i, + code: ProtosocketControlCode::Normal as u32, + body: Some(EchoRequest { + message: i.to_string(), + nanotime: SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("time works") + .as_nanos() as u64, + }), + }; + concurrent_count.fetch_add(1, Ordering::Relaxed); + match client.send_streaming(request) { + Ok(mut completion) => { + while let Some(Ok(response)) = completion.next().await { + handle_stream_response(response, &metrics_count, &metrics_latency); + } + } + Err(e) => { + log::error!("send should work: {e:?}"); + return; + } + } + concurrent_count.fetch_sub(1, Ordering::Relaxed); + } +} + +fn handle_stream_response( + response: Response, + metrics_count: &AtomicUsize, + metrics_latency: &histogram::AtomicHistogram, +) { + metrics_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + + assert_ne!(response.request_id, 0, "received bad message"); + match response.body { + Some(char_response) => { + let latency = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("time works") + .as_nanos() as u64 + - char_response.nanotime; + let _ = metrics_latency.increment(latency); + } + None => { + log::warn!("no response body"); + } + } +} diff --git a/example-proto-spawn-stream/src/messages.rs b/example-proto-spawn-stream/src/messages.rs new file mode 100644 index 0000000..5859384 --- /dev/null +++ b/example-proto-spawn-stream/src/messages.rs @@ -0,0 +1,103 @@ +//! If you're only using rust, of course you can hand-write prost structs, but if you +//! want to use a protosocket server with clients in other languages you'll want to +//! generate from protos. + +use protosocket_rpc::ProtosocketControlCode; + +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct Request { + #[prost(uint64, tag = "1")] + pub request_id: u64, + #[prost(uint32, tag = "2")] + pub code: u32, + #[prost(message, tag = "3")] + pub body: Option, +} + +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct EchoRequest { + #[prost(string, tag = "1")] + pub message: String, + #[prost(uint64, tag = "2")] + pub nanotime: u64, +} + +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct Response { + #[prost(uint64, tag = "1")] + pub request_id: u64, + #[prost(uint32, tag = "2")] + pub code: u32, + #[prost(message, tag = "3")] + pub body: Option, +} + +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct EchoStream { + #[prost(string, tag = "1")] + pub message: String, + #[prost(uint64, tag = "2")] + pub nanotime: u64, + #[prost(uint64, tag = "3")] + pub sequence: u64, +} + +impl protosocket_rpc::Message for Request { + fn message_id(&self) -> u64 { + self.request_id + } + + fn control_code(&self) -> ProtosocketControlCode { + ProtosocketControlCode::from_u8(self.code as u8) + } + + fn cancelled(request_id: u64) -> Self { + Request { + request_id, + code: ProtosocketControlCode::Cancel as u32, + body: None, + } + } + + fn set_message_id(&mut self, message_id: u64) { + self.request_id = message_id; + } + + fn ended(request_id: u64) -> Self { + Self { + request_id, + code: ProtosocketControlCode::End as u32, + body: None, + } + } +} + +impl protosocket_rpc::Message for Response { + fn message_id(&self) -> u64 { + self.request_id + } + + fn control_code(&self) -> ProtosocketControlCode { + ProtosocketControlCode::from_u8(self.code as u8) + } + + fn cancelled(request_id: u64) -> Self { + Response { + request_id, + code: ProtosocketControlCode::Cancel as u32, + body: None, + } + } + + fn set_message_id(&mut self, message_id: u64) { + self.request_id = message_id + } + + fn ended(request_id: u64) -> Self { + Self { + request_id, + code: ProtosocketControlCode::End as u32, + body: None, + } + } +} diff --git a/example-proto-spawn-stream/src/server.rs b/example-proto-spawn-stream/src/server.rs new file mode 100644 index 0000000..679fbb3 --- /dev/null +++ b/example-proto-spawn-stream/src/server.rs @@ -0,0 +1,141 @@ +//! A protosocket-rpc server that spawns its streaming rpc producers onto the runtime. +//! +//! Rpcs are normally polled by the connection, which stops polling them when the peer +//! can't receive. Spawning decouples the producer from that backpressure: the task runs +//! regardless. Bound it with a channel - its capacity is how far the producer can run +//! ahead of the peer. + +use std::pin::pin; +use std::sync::atomic::AtomicUsize; + +use futures::{future::BoxFuture, stream::BoxStream, Stream, StreamExt}; +use messages::{EchoRequest, EchoStream, Request, Response}; +use protosocket::{PooledEncoder, StreamWithAddress, TcpSocketListener}; +use protosocket_prost::{ProstDecoder, ProstSerializer}; +use protosocket_rpc::{ + server::{ConnectionService, RpcKind, SocketService}, + ProtosocketControlCode, +}; +use tokio::net::TcpStream; + +mod messages; + +fn main() -> Result<(), Box> { + static I: AtomicUsize = AtomicUsize::new(0); + let runtime = tokio::runtime::Builder::new_multi_thread() + .thread_name_fn(|| { + format!( + "app-{}", + I.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + ) + }) + .worker_threads(2) + .event_interval(7) + .enable_all() + .build()?; + + runtime.block_on(run_main()) +} + +#[allow(clippy::expect_used)] +async fn run_main() -> Result<(), Box> { + env_logger::init(); + let mut server = protosocket_rpc::server::SocketRpcServer::new( + TcpSocketListener::listen( + std::env::var("HOST") + .unwrap_or_else(|_| "0.0.0.0:9000".to_string()) + .parse()?, + 1024, + None, + )?, + DemoRpcSocketService, + 4 << 20, + 1 << 20, + 128, + )?; + server.set_max_queued_outbound_messages(512); + + tokio::spawn(server).await??; + Ok(()) +} + +struct DemoRpcSocketService; +impl SocketService for DemoRpcSocketService { + type ConnectionService = DemoRpcConnectionServer; + type SocketListener = TcpSocketListener; + + fn codec(&mut self) -> ::Codec { + Default::default() + } + + fn new_stream_service( + &mut self, + stream: &StreamWithAddress, + ) -> Self::ConnectionService { + log::info!("new connection server {}", stream.address()); + DemoRpcConnectionServer { + address: stream.address(), + } + } +} + +struct DemoRpcConnectionServer { + address: std::net::SocketAddr, +} +impl ConnectionService for DemoRpcConnectionServer { + type Codec = ( + PooledEncoder>, + ProstDecoder, + ); + type Request = Request; + type Response = Response; + type UnaryFutureType = BoxFuture<'static, Response>; + type StreamType = BoxStream<'static, Response>; + + fn new_rpc( + &mut self, + initiating_message: Self::Request, + ) -> RpcKind { + log::debug!("{} new rpc: {initiating_message:?}", self.address); + let request_id = initiating_message.request_id; + match initiating_message.body { + Some(echo) => { + // Spawn the producer behind a bounded channel. send().await is the + // backpressure: the producer waits while the channel is full, and quits + // when the rpc is cancelled or the connection closes (the receiver + // drops). + let (sender, mut receiver) = tokio::sync::mpsc::channel(16); + tokio::spawn(async move { + let mut stream = pin!(echo_stream(request_id, echo)); + while let Some(response) = stream.next().await { + if sender.send(response).await.is_err() { + break; + } + } + }); + RpcKind::Streaming( + futures::stream::poll_fn(move |context| receiver.poll_recv(context)).boxed(), + ) + } + None => { + log::warn!("received empty echo request id {request_id}"); + RpcKind::Cancelled + } + } + } +} + +fn echo_stream(request_id: u64, echo: EchoRequest) -> impl Stream { + let nanotime = echo.nanotime; + futures::stream::iter(echo.message.into_bytes().into_iter().enumerate().map( + move |(sequence, c)| Response { + request_id, + code: ProtosocketControlCode::Normal as u32, + body: Some(EchoStream { + message: (c as char).to_string(), + nanotime, + sequence: sequence as u64, + }), + }, + )) +} diff --git a/example-proto-spawn/Cargo.toml b/example-proto-spawn/Cargo.toml new file mode 100644 index 0000000..066dbc7 --- /dev/null +++ b/example-proto-spawn/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "example-proto-spawn" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "example-proto-spawn-server" +path = "src/server.rs" + +[[bin]] +name = "example-proto-spawn-client" +path = "src/client.rs" + +[dependencies] +protosocket = { workspace = true } +protosocket-prost = { workspace = true } +protosocket-rpc = { workspace = true } + +env_logger = { workspace = true } +futures = { workspace = true } +log = { workspace = true } +prost = { workspace = true } +tokio = { workspace = true, features = ["full"] } + +histogram = { version = "0.11" } diff --git a/example-proto-spawn/src/client.rs b/example-proto-spawn/src/client.rs new file mode 100644 index 0000000..7a40385 --- /dev/null +++ b/example-proto-spawn/src/client.rs @@ -0,0 +1,208 @@ +use std::{ + sync::{ + atomic::{AtomicU64, AtomicUsize, Ordering}, + Arc, + }, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use messages::{EchoRequest, Request, Response}; +use protosocket::PooledEncoder; +use protosocket_prost::{ProstDecoder, ProstSerializer}; +use protosocket_rpc::{ + client::{Configuration, RpcClient, TcpStreamConnector}, + ProtosocketControlCode, +}; + +mod messages; + +fn main() -> Result<(), Box> { + static I: AtomicUsize = AtomicUsize::new(0); + let runtime = tokio::runtime::Builder::new_multi_thread() + .thread_name_fn(|| { + format!( + "app-{}", + I.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + ) + }) + .worker_threads(4) + .event_interval(7) + .enable_all() + .build()?; + + runtime.block_on(run_main()) +} + +async fn run_main() -> Result<(), Box> { + env_logger::init(); + + let response_count = Arc::new(AtomicUsize::new(0)); + let shed_count = Arc::new(AtomicUsize::new(0)); + let latency = Arc::new(histogram::AtomicHistogram::new(7, 52).expect("histogram works")); + + let concurrency = 256; + let connections = 4; + + let concurrent_count = Arc::new(AtomicUsize::new(0)); + let request_ids = Arc::new(AtomicU64::new(1)); + let mut configuration = Configuration::new(TcpStreamConnector); + configuration.max_queued_outbound_messages(64); + for _i in 0..connections { + let (client, connection) = protosocket_rpc::client::connect::< + ( + PooledEncoder>, + ProstDecoder, + ), + _, + >( + std::env::var("ENDPOINT") + .unwrap_or_else(|_| "127.0.0.1:9000".to_string()) + .parse() + .expect("must use a valid socket address"), + &configuration, + ) + .await?; + let _connection_handle = tokio::spawn(connection); + let tasks = concurrency / connections; + for _ in 0..tasks { + let _client_handle = tokio::spawn(generate_traffic( + concurrent_count.clone(), + request_ids.clone(), + client.clone(), + response_count.clone(), + shed_count.clone(), + latency.clone(), + )); + } + } + + let metrics = tokio::spawn(print_periodic_metrics( + response_count, + shed_count, + latency, + concurrent_count, + )); + + tokio::select!( + // _ = connection_driver => { + // log::warn!("connection driver quit"); + // } + // _ = client_runtime => { + // log::warn!("client runtime quit"); + // } + _ = metrics => { + log::warn!("metrics runtime quit"); + } + ); + + Ok(()) +} + +async fn print_periodic_metrics( + response_count: Arc, + shed_count: Arc, + latency: Arc, + concurrent_count: Arc, +) { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + loop { + let start = Instant::now(); + interval.tick().await; + let total = response_count.swap(0, std::sync::atomic::Ordering::Relaxed); + let hz = (total as f64) / start.elapsed().as_secs_f64().max(0.1); + + let latency = latency.drain(); + let p90 = latency + .percentile(0.9) + .unwrap_or_default() + .map(|b| *b.range().end()) + .unwrap_or_default() as f64 + / 1000.0; + let p999 = latency + .percentile(0.999) + .unwrap_or_default() + .map(|b| *b.range().end()) + .unwrap_or_default() as f64 + / 1000.0; + let p9999 = latency + .percentile(0.9999) + .unwrap_or_default() + .map(|b| *b.range().end()) + .unwrap_or_default() as f64 + / 1000.0; + let shed = shed_count.swap(0, std::sync::atomic::Ordering::Relaxed); + let concurrent = concurrent_count.load(std::sync::atomic::Ordering::Relaxed); + eprintln!( + "Messages: {total:10} rate: {hz:9.1}hz shed: {shed:6} p90: {p90:6.1}µs p999: {p999:6.1}µs p9999: {p9999:6.1}µs concurrency: {concurrent}" + ); + } +} + +async fn generate_traffic( + concurrent_count: Arc, + request_ids: Arc, + client: RpcClient, + metrics_count: Arc, + shed_count: Arc, + metrics_latency: Arc, +) { + log::debug!("running traffic generator"); + loop { + let i = request_ids.fetch_add(1, Ordering::Relaxed); + let request = Request { + request_id: i, + code: ProtosocketControlCode::Normal as u32, + body: Some(EchoRequest { + message: i.to_string(), + nanotime: SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("time works") + .as_nanos() as u64, + }), + }; + concurrent_count.fetch_add(1, Ordering::Relaxed); + match client.send_unary(request) { + Ok(completion) => match completion.await { + Ok(response) => { + handle_unary_response(response, &metrics_count, &metrics_latency); + } + // The server sheds rpcs with a cancellation when its spawn slots are + // all in flight. + Err(_shed) => { + shed_count.fetch_add(1, Ordering::Relaxed); + } + }, + Err(e) => { + log::error!("send should work: {e:?}"); + return; + } + } + concurrent_count.fetch_sub(1, Ordering::Relaxed); + } +} + +fn handle_unary_response( + response: Response, + metrics_count: &AtomicUsize, + metrics_latency: &histogram::AtomicHistogram, +) { + metrics_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + + let request_id = response.request_id; + assert_ne!(response.request_id, 0, "received bad message"); + match response.body { + Some(echo) => { + assert_eq!(request_id, echo.message.parse().unwrap_or_default()); + + let latency = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("time works") + .as_nanos() as u64 + - echo.nanotime; + let _ = metrics_latency.increment(latency); + } + None => { + log::warn!("no response body"); + } + } +} diff --git a/example-proto-spawn/src/messages.rs b/example-proto-spawn/src/messages.rs new file mode 100644 index 0000000..2bc7625 --- /dev/null +++ b/example-proto-spawn/src/messages.rs @@ -0,0 +1,101 @@ +//! If you're only using rust, of course you can hand-write prost structs, but if you +//! want to use a protosocket server with clients in other languages you'll want to +//! generate from protos. + +use protosocket_rpc::ProtosocketControlCode; + +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct Request { + #[prost(uint64, tag = "1")] + pub request_id: u64, + #[prost(uint32, tag = "2")] + pub code: u32, + #[prost(message, tag = "3")] + pub body: Option, +} + +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct EchoRequest { + #[prost(string, tag = "1")] + pub message: String, + #[prost(uint64, tag = "2")] + pub nanotime: u64, +} + +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct Response { + #[prost(uint64, tag = "1")] + pub request_id: u64, + #[prost(uint32, tag = "2")] + pub code: u32, + #[prost(message, tag = "3")] + pub body: Option, +} + +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct EchoResponse { + #[prost(string, tag = "1")] + pub message: String, + #[prost(uint64, tag = "2")] + pub nanotime: u64, +} + +impl protosocket_rpc::Message for Request { + fn message_id(&self) -> u64 { + self.request_id + } + + fn control_code(&self) -> ProtosocketControlCode { + ProtosocketControlCode::from_u8(self.code as u8) + } + + fn cancelled(request_id: u64) -> Self { + Request { + request_id, + code: ProtosocketControlCode::Cancel as u32, + body: None, + } + } + + fn set_message_id(&mut self, message_id: u64) { + self.request_id = message_id; + } + + fn ended(request_id: u64) -> Self { + Self { + request_id, + code: ProtosocketControlCode::End as u32, + body: None, + } + } +} + +impl protosocket_rpc::Message for Response { + fn message_id(&self) -> u64 { + self.request_id + } + + fn control_code(&self) -> ProtosocketControlCode { + ProtosocketControlCode::from_u8(self.code as u8) + } + + fn cancelled(request_id: u64) -> Self { + Response { + request_id, + code: ProtosocketControlCode::Cancel as u32, + body: None, + } + } + + fn set_message_id(&mut self, message_id: u64) { + self.request_id = message_id + } + + fn ended(request_id: u64) -> Self { + Self { + request_id, + code: ProtosocketControlCode::End as u32, + body: None, + } + } +} diff --git a/example-proto-spawn/src/server.rs b/example-proto-spawn/src/server.rs new file mode 100644 index 0000000..a75f4c0 --- /dev/null +++ b/example-proto-spawn/src/server.rs @@ -0,0 +1,154 @@ +//! A protosocket-rpc server that spawns its unary rpcs onto the runtime. +//! +//! Rpcs are normally polled by the connection, which stops polling them when the peer +//! can't receive. Spawning decouples rpc work from that backpressure: the spawned task +//! runs regardless. Bound what you spawn - here, rpcs shed load when the connection's +//! spawn slots are all in flight. + +use std::sync::atomic::AtomicUsize; +use std::sync::Arc; +use std::time::Duration; + +use futures::{future::BoxFuture, stream::BoxStream, FutureExt}; +use messages::{EchoRequest, EchoResponse, Request, Response}; +use protosocket::{PooledEncoder, StreamWithAddress, TcpSocketListener}; +use protosocket_prost::{ProstDecoder, ProstSerializer}; +use protosocket_rpc::{ + server::{ConnectionService, RpcKind, SocketService}, + Message, ProtosocketControlCode, +}; +use tokio::net::TcpStream; +use tokio::sync::Semaphore; + +mod messages; + +fn main() -> Result<(), Box> { + static I: AtomicUsize = AtomicUsize::new(0); + let runtime = tokio::runtime::Builder::new_multi_thread() + .thread_name_fn(|| { + format!( + "app-{}", + I.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + ) + }) + .worker_threads(2) + .event_interval(7) + .enable_all() + .build()?; + + runtime.block_on(run_main()) +} + +#[allow(clippy::expect_used)] +async fn run_main() -> Result<(), Box> { + env_logger::init(); + let mut server = protosocket_rpc::server::SocketRpcServer::new( + TcpSocketListener::listen( + std::env::var("HOST") + .unwrap_or_else(|_| "0.0.0.0:9000".to_string()) + .parse()?, + 1024, + None, + )?, + DemoRpcSocketService, + 4 << 20, + 1 << 20, + 128, + )?; + server.set_max_queued_outbound_messages(512); + + tokio::spawn(server).await??; + Ok(()) +} + +struct DemoRpcSocketService; +impl SocketService for DemoRpcSocketService { + type ConnectionService = DemoRpcConnectionServer; + type SocketListener = TcpSocketListener; + + fn codec(&mut self) -> ::Codec { + Default::default() + } + + fn new_stream_service( + &mut self, + stream: &StreamWithAddress, + ) -> Self::ConnectionService { + log::info!("new connection server {}", stream.address()); + DemoRpcConnectionServer { + address: stream.address(), + // How many spawned unary rpcs may be in flight for this connection. Small + // enough that the example client can outrun it and show shedding. + spawn_limit: Arc::new(Semaphore::new(16)), + } + } +} + +struct DemoRpcConnectionServer { + address: std::net::SocketAddr, + spawn_limit: Arc, +} +impl ConnectionService for DemoRpcConnectionServer { + type Codec = ( + PooledEncoder>, + ProstDecoder, + ); + type Request = Request; + type Response = Response; + type UnaryFutureType = BoxFuture<'static, Response>; + type StreamType = BoxStream<'static, Response>; + + fn new_rpc( + &mut self, + initiating_message: Self::Request, + ) -> RpcKind { + log::debug!("{} new rpc: {initiating_message:?}", self.address); + let request_id = initiating_message.request_id; + match initiating_message.body { + Some(echo) => { + // Spawning opts this work out of connection backpressure, so bound it + // by shedding: no slot, no rpc. Slots are held until responses are + // handed off, so a peer that stops receiving sheds instead of + // accumulating spawned work. A real service might prefer an explicit + // "too busy" response over a cancellation. + match Arc::clone(&self.spawn_limit).try_acquire_owned() { + Ok(slot) => { + let work = tokio::spawn(echo_request(request_id, echo)); + RpcKind::Unary( + async move { + let _held_until_handoff = slot; + match work.await { + Ok(response) => response, + Err(join_error) => { + log::error!("rpc task failed: {join_error}"); + Response::cancelled(request_id) + } + } + } + .boxed(), + ) + } + Err(_no_slot) => RpcKind::Cancelled, + } + } + None => { + log::warn!("received empty echo request id {request_id}"); + RpcKind::Cancelled + } + } + } +} + +async fn echo_request(request_id: u64, echo: EchoRequest) -> Response { + // Pretend this is heavy work that deserves its own task. It is slow enough that + // the example client outruns the spawn slots, so you can watch shedding happen. + tokio::time::sleep(Duration::from_micros(500)).await; + Response { + request_id, + code: ProtosocketControlCode::Normal as u32, + body: Some(EchoResponse { + message: echo.message, + nanotime: echo.nanotime, + }), + } +} diff --git a/example-proto-stream/src/server.rs b/example-proto-stream/src/server.rs index e2c89b3..87e7a1c 100644 --- a/example-proto-stream/src/server.rs +++ b/example-proto-stream/src/server.rs @@ -1,16 +1,10 @@ -use std::pin::pin; - -use futures::{ - FutureExt, Stream, - future::{BoxFuture, join_all}, - stream::FuturesUnordered, -}; +use futures::{Stream, StreamExt, future::join_all, stream::BoxStream}; use messages::{EchoRequest, EchoResponse, EchoStream, Request, Response, ResponseBehavior}; use protosocket::{PooledEncoder, StreamWithAddress, TcpSocketListener}; use protosocket_prost::{ProstDecoder, ProstSerializer}; use protosocket_rpc::{ - Message, ProtosocketControlCode, - server::{ConnectionService, LevelSpawn, RpcResponder, SocketService}, + ProtosocketControlCode, + server::{ConnectionService, LevelSpawn, RpcKind, SocketService}, }; use tokio::net::TcpStream; @@ -67,81 +61,68 @@ async fn run_main() -> Result<(), std::io::Error> { /// ConnectionServices to application-wide state tracking. struct DemoRpcSocketService; impl SocketService for DemoRpcSocketService { - type Codec = ( - // Use a pooled encoder to amortize memory allocation cost. - // Each connection gets its own little memory pool. - PooledEncoder>, - ProstDecoder, - ); type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&self) -> Self::Codec { + fn codec(&mut self) -> ::Codec { ( PooledEncoder::new_with_pool_size(64, Default::default()), ProstDecoder::default(), ) } - fn new_stream_service(&self, stream: &StreamWithAddress) -> Self::ConnectionService { + fn new_stream_service( + &mut self, + stream: &StreamWithAddress, + ) -> Self::ConnectionService { log::info!("new connection server {}", stream.address()); DemoRpcConnectionServer { address: stream.address(), - streams: FuturesUnordered::new(), } } } /// This is the entry point for each Connection. State per-connection is tracked, and you /// get mutable access to the service on each new rpc for state tracking. +/// +/// The connection drives your rpcs: they are only polled when the connection can send, +/// so a slow peer slows its rpcs down instead of buffering responses without bound. struct DemoRpcConnectionServer { address: std::net::SocketAddr, - streams: FuturesUnordered>, } impl ConnectionService for DemoRpcConnectionServer { + type Codec = ( + // Use a pooled encoder to amortize memory allocation cost. + // Each connection gets its own little memory pool. + PooledEncoder>, + ProstDecoder, + ); type Request = Request; type Response = Response; + type UnaryFutureType = futures::future::Ready; + type StreamType = BoxStream<'static, Response>; fn new_rpc( &mut self, initiating_message: Self::Request, - responder: RpcResponder<'_, Self::Response>, - ) { + ) -> RpcKind { log::debug!("{} new rpc: {initiating_message:?}", self.address); let request_id = initiating_message.request_id; let behavior = initiating_message.response_behavior(); match initiating_message.body { Some(echo) => match behavior { - ResponseBehavior::Unary => { - responder.immediate(immediate_echo_response(request_id, echo)) - } + ResponseBehavior::Unary => RpcKind::Unary(futures::future::ready( + immediate_echo_response(request_id, echo), + )), ResponseBehavior::Stream => { - self.streams - .push(responder.stream(echo_stream(request_id, echo)).boxed()); + RpcKind::Streaming(echo_stream(request_id, echo).boxed()) } }, None => { - log::warn!("received empty echo request: {initiating_message:?}"); - responder.immediate(Response::cancelled(request_id)); - } - } - } - - fn poll( - mut self: std::pin::Pin<&mut Self>, - context: &mut std::task::Context<'_>, - ) -> std::ops::ControlFlow<()> { - while !self.streams.is_empty() - && let std::task::Poll::Ready(next) = pin!(&mut self.streams).poll_next(context) - { - if next.is_none() { - log::error!( - "the stream of futures should never return None. We don't poll it while empty" - ); - return std::ops::ControlFlow::Break(()); + log::warn!("received empty echo request id {request_id}"); + RpcKind::Cancelled } } - std::ops::ControlFlow::Continue(()) } } diff --git a/example-proto-tls/src/server.rs b/example-proto-tls/src/server.rs index e382916..a1b139c 100644 --- a/example-proto-tls/src/server.rs +++ b/example-proto-tls/src/server.rs @@ -1,12 +1,12 @@ use std::sync::atomic::AtomicUsize; -use futures::Stream; +use futures::{future::BoxFuture, stream::BoxStream, FutureExt, Stream, StreamExt}; use messages::{EchoRequest, EchoResponse, EchoStream, Request, Response, ResponseBehavior}; use protosocket::{PooledEncoder, StreamWithAddress, TcpSocketListener, TlsSocketListener}; use protosocket_prost::{ProstDecoder, ProstSerializer}; use protosocket_rpc::{ - server::{ConnectionService, RpcResponder, SocketService}, - Message, ProtosocketControlCode, + server::{ConnectionService, RpcKind, SocketService}, + ProtosocketControlCode, }; use rustls_pemfile::Item; @@ -88,19 +88,15 @@ async fn run_main() -> Result<(), Box> { /// ConnectionServices to application-wide state tracking. struct DemoRpcSocketService {} impl SocketService for DemoRpcSocketService { - type Codec = ( - PooledEncoder>, - ProstDecoder, - ); type ConnectionService = DemoRpcConnectionServer; type SocketListener = TlsSocketListener; - fn codec(&self) -> Self::Codec { + fn codec(&mut self) -> ::Codec { Default::default() } fn new_stream_service( - &self, + &mut self, stream: &StreamWithAddress>, ) -> Self::ConnectionService { DemoRpcConnectionServer { @@ -115,33 +111,33 @@ struct DemoRpcConnectionServer { address: std::net::SocketAddr, } impl ConnectionService for DemoRpcConnectionServer { + type Codec = ( + PooledEncoder>, + ProstDecoder, + ); type Request = Request; type Response = Response; + type UnaryFutureType = BoxFuture<'static, Response>; + type StreamType = BoxStream<'static, Response>; fn new_rpc( &mut self, initiating_message: Self::Request, - responder: RpcResponder<'_, Self::Response>, - ) { + ) -> RpcKind { log::debug!("{} new rpc: {initiating_message:?}", self.address); let request_id = initiating_message.request_id; let behavior = initiating_message.response_behavior(); match initiating_message.body { Some(echo) => match behavior { - ResponseBehavior::Unary => { - // See you can spawn the responder futures however you need to - tokio::spawn(responder.unary(echo_request(request_id, echo))); - } + // An async fn works fine as a unary rpc - box it if it isn't Unpin. + ResponseBehavior::Unary => RpcKind::Unary(echo_request(request_id, echo).boxed()), ResponseBehavior::Stream => { - tokio::spawn(responder.stream(echo_stream(request_id, echo))); + RpcKind::Streaming(echo_stream(request_id, echo).boxed()) } }, None => { - // No completion messages will be sent for this message - log::warn!( - "{request_id} no request in rpc body. This may cause a client memory leak." - ); - responder.immediate(Response::cancelled(request_id)); + log::warn!("{request_id} no request in rpc body"); + RpcKind::Cancelled } } } diff --git a/example-proto/src/server.rs b/example-proto/src/server.rs index a6b3eb2..3adcab4 100644 --- a/example-proto/src/server.rs +++ b/example-proto/src/server.rs @@ -1,10 +1,10 @@ -use futures::{future::join_all, Stream}; +use futures::{future::join_all, stream::BoxStream, Stream, StreamExt}; use messages::{EchoRequest, EchoResponse, EchoStream, Request, Response, ResponseBehavior}; use protosocket::{PooledEncoder, StreamWithAddress, TcpSocketListener}; use protosocket_prost::{ProstDecoder, ProstSerializer}; use protosocket_rpc::{ - server::{ConnectionService, LevelSpawn, RpcResponder, SocketService}, - Message, ProtosocketControlCode, + server::{ConnectionService, LevelSpawn, RpcKind, SocketService}, + ProtosocketControlCode, }; use tokio::net::TcpStream; @@ -61,23 +61,20 @@ async fn run_main() -> Result<(), std::io::Error> { /// ConnectionServices to application-wide state tracking. struct DemoRpcSocketService; impl SocketService for DemoRpcSocketService { - type Codec = ( - // Use a pooled encoder to amortize memory allocation cost. - // Each connection gets its own little memory pool. - PooledEncoder>, - ProstDecoder, - ); type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&self) -> Self::Codec { + fn codec(&mut self) -> ::Codec { ( PooledEncoder::new_with_pool_size(64, Default::default()), ProstDecoder::default(), ) } - fn new_stream_service(&self, stream: &StreamWithAddress) -> Self::ConnectionService { + fn new_stream_service( + &mut self, + stream: &StreamWithAddress, + ) -> Self::ConnectionService { log::info!("new connection server {}", stream.address()); DemoRpcConnectionServer { address: stream.address(), @@ -91,29 +88,36 @@ struct DemoRpcConnectionServer { address: std::net::SocketAddr, } impl ConnectionService for DemoRpcConnectionServer { + type Codec = ( + // Use a pooled encoder to amortize memory allocation cost. + // Each connection gets its own little memory pool. + PooledEncoder>, + ProstDecoder, + ); type Request = Request; type Response = Response; + type UnaryFutureType = futures::future::Ready; + type StreamType = BoxStream<'static, Response>; fn new_rpc( &mut self, initiating_message: Self::Request, - responder: RpcResponder<'_, Self::Response>, - ) { + ) -> RpcKind { log::debug!("{} new rpc: {initiating_message:?}", self.address); let request_id = initiating_message.request_id; let behavior = initiating_message.response_behavior(); match initiating_message.body { Some(echo) => match behavior { - ResponseBehavior::Unary => { - responder.immediate(immediate_echo_response(request_id, echo)) - } + ResponseBehavior::Unary => RpcKind::Unary(futures::future::ready( + immediate_echo_response(request_id, echo), + )), ResponseBehavior::Stream => { - tokio::spawn(responder.stream(echo_stream(request_id, echo))); + RpcKind::Streaming(echo_stream(request_id, echo).boxed()) } }, None => { - log::warn!("received empty echo request: {initiating_message:?}"); - responder.immediate(Response::cancelled(request_id)); + log::warn!("received empty echo request id {request_id}"); + RpcKind::Cancelled } } } diff --git a/example-telnet/src/main.rs b/example-telnet/src/main.rs index 2658389..253a1c8 100644 --- a/example-telnet/src/main.rs +++ b/example-telnet/src/main.rs @@ -35,12 +35,12 @@ impl ServerConnector for ServerContext { type Codec = StringCodec; type Reactor = StringReactor; - fn codec(&self) -> Self::Codec { + fn codec(&mut self) -> Self::Codec { StringCodec } fn new_reactor( - &self, + &mut self, optional_outbound: spillway::Sender<::Message>, _address: &StreamWithAddress, ) -> Self::Reactor { @@ -53,7 +53,6 @@ impl ServerConnector for ServerContext { &self, connection: protosocket::Connection< ::Stream, - Self::Codec, Self::Reactor, >, ) { @@ -65,6 +64,7 @@ struct StringReactor { outbound: spillway::Sender, } impl MessageReactor for StringReactor { + type Codec = StringCodec; type Inbound = String; type Outbound = String; type LogicalOutbound = String; diff --git a/protosocket-connection/examples/custom_allocator.rs b/protosocket-connection/examples/custom_allocator.rs index cd261ea..2202aee 100644 --- a/protosocket-connection/examples/custom_allocator.rs +++ b/protosocket-connection/examples/custom_allocator.rs @@ -27,6 +27,7 @@ struct EchoReactor { outbound_messages: spillway::Sender, } impl MessageReactor for EchoReactor { + type Codec = (BulkEncoder, BulkDecoder); type Inbound = Bytes; type Outbound = Bytes; type LogicalOutbound = Bytes; diff --git a/protosocket-connection/src/connection.rs b/protosocket-connection/src/connection.rs index d6f76b2..f3697a9 100644 --- a/protosocket-connection/src/connection.rs +++ b/protosocket-connection/src/connection.rs @@ -10,9 +10,9 @@ use bytes::Buf; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use crate::{ - encoding::Codec, interrupted, message_reactor::{MessageReactor, ReactorStatus}, + send_budget::SendBudget, would_block, Decoder, DeserializeError, Encoder, }; @@ -30,60 +30,45 @@ use crate::{ pub struct Connection< // Bidirectional Stream type to use for this connection. Like `tokio::net::TcpStream`. TStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + 'static, - // The wire / message codec for this connection. - TCodec: Codec, - // The message reactor for this connection. - TReactor: MessageReactor::Message, Outbound = ::Message>, + // The message reactor for this connection. It brings the wire codec with it. + TReactor: MessageReactor, > { stream: TStream, outbound_messages: spillway::Receiver, - send_buffer: VecDeque<::Serialized>, + send_buffer: VecDeque<::Serialized>, receive_buffer_unread_index: usize, receive_buffer: Vec, max_buffer_length: usize, max_queued_send_messages: usize, buffer_allocation_increment: usize, - codec: TCodec, + codec: TReactor::Codec, reactor: TReactor, } -impl< - TStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + 'static, - TCodec: Codec, - TReactor: MessageReactor< - Inbound = ::Message, - Outbound = ::Message, - >, - > std::fmt::Display for Connection +impl + std::fmt::Display for Connection { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let read_end = self.receive_buffer_unread_index; let read_capacity = self.receive_buffer.len(); let write_queue = self.send_buffer.len(); let write_length: usize = self.send_buffer.iter().map(|b| b.remaining()).sum(); - write!(f, "Connection: {{read{{end: {read_end}, capacity: {read_capacity}}}, write{{queue: {write_queue}, length: {write_length}}} }}") + write!( + f, + "Connection: {{read{{end: {read_end}, capacity: {read_capacity}}}, write{{queue: {write_queue}, length: {write_length}}} }}" + ) } } -impl< - TStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + 'static, - TCodec: Codec, - TReactor: MessageReactor< - Inbound = ::Message, - Outbound = ::Message, - >, - > Unpin for Connection +impl + Unpin for Connection { } impl< TStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + 'static, - TCodec: Codec, - TReactor: MessageReactor< - Inbound = ::Message, - Outbound = ::Message, - >, - > Future for Connection + TReactor: MessageReactor, + > Future for Connection { type Output = (); @@ -128,14 +113,8 @@ impl< } } -impl< - TStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + 'static, - TCodec: Codec, - TReactor: MessageReactor< - Inbound = ::Message, - Outbound = ::Message, - >, - > Drop for Connection +impl Drop + for Connection { fn drop(&mut self) { log::debug!("connection dropped") @@ -156,12 +135,8 @@ enum ReadBufferState { impl< TStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + 'static, - TCodec: Codec, - TReactor: MessageReactor< - Inbound = ::Message, - Outbound = ::Message, - >, - > Connection + TReactor: MessageReactor, + > Connection { /// Create a new protosocket Connection with the given stream and reactor. /// @@ -169,7 +144,7 @@ impl< #[allow(clippy::type_complexity, clippy::too_many_arguments)] pub fn new( stream: TStream, - codec: TCodec, + codec: TReactor::Codec, max_buffer_length: usize, buffer_allocation_increment: usize, max_queued_send_messages: usize, @@ -240,7 +215,11 @@ impl< Err(e) => match e { DeserializeError::IncompleteBuffer { next_message_size } => { if self.max_buffer_length < next_message_size { - log::error!("tried to receive message that is too long. Resetting connection - max: {}, requested: {}", self.max_buffer_length, next_message_size); + log::error!( + "tried to receive message that is too long. Resetting connection - max: {}, requested: {}", + self.max_buffer_length, + next_message_size + ); return ReadBufferState::Disconnected; } log::debug!("waiting for the next message of length {next_message_size}"); @@ -252,7 +231,10 @@ impl< } DeserializeError::SkipMessage { distance } => { if self.receive_buffer_unread_index - buffer_cursor < distance { - log::trace!("cannot skip yet, need to read more. Skipping: {distance}, remaining:{}", self.receive_buffer_unread_index - buffer_cursor); + log::trace!( + "cannot skip yet, need to read more. Skipping: {distance}, remaining:{}", + self.receive_buffer_unread_index - buffer_cursor + ); break ReadBufferState::Pending; } log::debug!("skipping message of length {distance}"); @@ -313,7 +295,10 @@ impl< } } - /// This serializes work-in-progress messages and moves them over into the write queue + /// This serializes work-in-progress messages and moves them over into the write queue. + /// + /// Outbound messages come from the queue and from the reactor, in that order, within + /// the send budget. The connection closes when both sources are exhausted. #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))] fn poll_serialize_outbound_messages(&mut self, context: &mut Context<'_>) -> Poll<()> { let max_outbound = self.max_queued_send_messages - self.send_buffer.len(); @@ -324,26 +309,43 @@ impl< } let start_len = self.send_buffer.len(); - for _ in 0..max_outbound { - let message = match self.outbound_messages.poll_next(context) { - Poll::Pending => { - log::debug!("no more messages to serialize, and we are pending for more"); - break; + let mut queue_closed = false; + while self.send_buffer.len() - start_len < max_outbound { + match self.outbound_messages.poll_next(context) { + Poll::Ready(Some(message)) => { + let message = self.reactor.on_outbound_message(message); + let buffer = self.codec.encode(message); + log::trace!( + "serialized message and enqueueing outbound buffer: {}b", + buffer.remaining() + ); + // queue up a writev + self.send_buffer.push_back(buffer); } Poll::Ready(None) => { - log::info!("outbound message channel was closed"); - return Poll::Ready(()); + queue_closed = true; + break; } - Poll::Ready(Some(next)) => next, - }; - let message = self.reactor.on_outbound_message(message); - let buffer = self.codec.encode(message); - log::trace!( - "serialized message and enqueueing outbound buffer: {}b", - buffer.remaining() - ); - // queue up a writev - self.send_buffer.push_back(buffer); + Poll::Pending => break, + } + } + + let remaining = max_outbound - (self.send_buffer.len() - start_len); + if 0 < remaining { + let Self { + send_buffer, + reactor, + codec, + .. + } = self; + let mut outbound = SendBudget::new(codec, send_buffer, remaining); + // SAFETY: This is a structural pin. If I'm not moved then neither is this reactor. + let reactor_state = + unsafe { Pin::new_unchecked(reactor) }.poll_outbound_many(context, &mut outbound); + if queue_closed && matches!(reactor_state, Poll::Ready(None)) { + log::info!("all outbound message sources are exhausted"); + return Poll::Ready(()); + } } let new_len = self.send_buffer.len(); if start_len != new_len { @@ -450,7 +452,9 @@ impl< } else { // Walk the buffer forward. It needs to be the next bytes on the wire, so we'll put it back in front. // Partial buffer consumption is relatively uncommon, but it definitely happens. - log::debug!("after writing {total_written}b, advancing partially written buffer of {remaining}b by {written}b"); + log::debug!( + "after writing {total_written}b, advancing partially written buffer of {remaining}b by {written}b" + ); front.advance(written); self.send_buffer.push_front(front); break; diff --git a/protosocket-connection/src/lib.rs b/protosocket-connection/src/lib.rs index a71dc9e..b8c1c28 100644 --- a/protosocket-connection/src/lib.rs +++ b/protosocket-connection/src/lib.rs @@ -13,6 +13,7 @@ mod encoding; mod error; mod message_reactor; mod pooled_encoder; +mod send_budget; mod socket_listener; pub use connection::Connection; @@ -26,6 +27,8 @@ pub use message_reactor::ReactorStatus; pub use pooled_encoder::PooledEncoder; pub use pooled_encoder::Reusable; pub use pooled_encoder::Serialize; +pub use send_budget::SendBudget; +pub use send_budget::SendPermit; pub use socket_listener::SocketListener; pub use socket_listener::SocketResult; pub use socket_listener::StreamWithAddress; diff --git a/protosocket-connection/src/message_reactor.rs b/protosocket-connection/src/message_reactor.rs index c300470..255b950 100644 --- a/protosocket-connection/src/message_reactor.rs +++ b/protosocket-connection/src/message_reactor.rs @@ -1,3 +1,5 @@ +use crate::{encoding::Codec, send_budget::SendBudget, Decoder, Encoder}; + /// A message reactor is a stateful object that processes inbound messages. /// You receive &mut self, and you receive your messages by value. /// @@ -7,6 +9,8 @@ /// Your message reactor and your tcp connection share their fate - when one drops or /// disconnects, the other does too. pub trait MessageReactor: 'static { + /// The wire codec for this reactor's connection. + type Codec: Codec + Decoder + Encoder; /// Messages inbound from the remote. type Inbound; /// Messages outbound to a remote. @@ -26,6 +30,9 @@ pub trait MessageReactor: 'static { /// You can use this to drive connection state machines (e.g., `FuturesUnordered`), /// or whatever else you need to do with your reactor between reading from the network and /// writing to it. + /// + /// This is polled unconditionaly. For producing outbound wire messages, you should implement + /// [poll_outbound_many], which exerts backpressure when outbound writes are blocked. fn poll( self: std::pin::Pin<&mut Self>, _context: &mut std::task::Context<'_>, @@ -33,10 +40,34 @@ pub trait MessageReactor: 'static { std::ops::ControlFlow::Continue(()) } - /// Called from the connection's driver task when messages are sent. + /// Called from the connection's driver task when messages from the outbound queue + /// are sent. /// /// You can use this to track outbound messages, or for logging, or metrics, or whatever. + /// + /// Not invoked for messages produced by `poll_outbound_many`. Do any required bookkeeping + /// within that method. fn on_outbound_message(&mut self, message: Self::LogicalOutbound) -> Self::Outbound; + + /// Produce outbound wire messages. + /// + /// This is only called when the connection has room to send, and `outbound` refuses + /// messages beyond that room. When the connection cannot write, work that produces + /// messages is not advanced. + /// + /// Messages emitted here do not trigger `on_outbound_message`: do any required bookkeeping + /// within this method. + /// + /// Return `Ready(Some(()))` after emitting, `Pending` when nothing is available + /// right now, and `Ready(None)` if you will never emit. The connection + /// closes when both its outbound queue is closed and this method returns `Ready(None)`. + fn poll_outbound_many( + self: std::pin::Pin<&mut Self>, + _context: &mut std::task::Context<'_>, + _outbound: &mut SendBudget<'_, Self::Codec>, + ) -> std::task::Poll> { + std::task::Poll::Ready(None) + } } /// What the connection should do after processing a batch of inbound messages. diff --git a/protosocket-connection/src/send_budget.rs b/protosocket-connection/src/send_budget.rs new file mode 100644 index 0000000..5d779e3 --- /dev/null +++ b/protosocket-connection/src/send_budget.rs @@ -0,0 +1,58 @@ +use std::collections::VecDeque; + +use bytes::Buf; + +use crate::Encoder; + +/// A bounded lease on a connection's send capacity. Sends encode directly into +/// the connection's send buffer. +pub struct SendBudget<'a, TEncoder: Encoder> { + encoder: &'a mut TEncoder, + send_buffer: &'a mut VecDeque, + remaining: usize, +} + +impl<'s, TEncoder: Encoder> SendBudget<'s, TEncoder> { + pub(crate) fn new( + encoder: &'s mut TEncoder, + send_buffer: &'s mut VecDeque, + remaining: usize, + ) -> Self { + Self { + encoder, + send_buffer, + remaining, + } + } + + /// Take a permit to send one message. None when the budget is spent. + /// Dropping a permit unused returns its capacity. + /// The permit hold &mut on the budget - you must spend or drop it before + /// you can reserve another. + /// + /// You can reserve ahead of doing work to exert backpressure if outbound writes + /// would be blocked. + pub fn reserve(&mut self) -> Option> { + if self.remaining == 0 { + None + } else { + Some(SendPermit { budget: self }) + } + } +} + +/// A reserved slot in a connection's send queue. A permit is one send: spending it +/// consumes it. +pub struct SendPermit<'a, 'b, TEncoder: Encoder> { + budget: &'a mut SendBudget<'b, TEncoder>, +} + +impl SendPermit<'_, '_, TEncoder> { + /// Spend the permit to send a message. + pub fn send(self, message: TEncoder::Message) { + self.budget.remaining -= 1; + let buffer = self.budget.encoder.encode(message); + log::trace!("serialized reactor message: {}b", buffer.remaining()); + self.budget.send_buffer.push_back(buffer); + } +} diff --git a/protosocket-connection/src/socket_listener.rs b/protosocket-connection/src/socket_listener.rs index a785d73..7c2dcaf 100644 --- a/protosocket-connection/src/socket_listener.rs +++ b/protosocket-connection/src/socket_listener.rs @@ -187,7 +187,7 @@ impl SocketListener for TlsSocketListener { continue; } Poll::Ready(SocketResult::Disconnect) => { - return Poll::Ready(SocketResult::Disconnect) + return Poll::Ready(SocketResult::Disconnect); } Poll::Pending => break, } diff --git a/protosocket-rpc/Cargo.toml b/protosocket-rpc/Cargo.toml index 5d4937d..c488a3e 100644 --- a/protosocket-rpc/Cargo.toml +++ b/protosocket-rpc/Cargo.toml @@ -28,5 +28,6 @@ thiserror = { workspace = true } webpki-roots = { workspace = true } [dev-dependencies] +bytes = { workspace = true } prost = { workspace = true } tokio = { workspace = true, features = ["full"] } diff --git a/protosocket-rpc/README.md b/protosocket-rpc/README.md index f64b7d5..eebf8fa 100644 --- a/protosocket-rpc/README.md +++ b/protosocket-rpc/README.md @@ -49,30 +49,35 @@ bytes you encode are the bytes which are sent. ## A diagram Each RPC is its own interaction: You can have thousands of concurrent RPCs on a single connection, all streaming, completing, or cancelling as each wants to. -``` +Servers return their rpc completions from `new_rpc` - a future for unary, a stream for +streaming - and the connection drives them. Completions are only polled when the +connection has room to send, so a peer that stops receiving stops its rpcs instead of +buffering responses without bound. + +```text Client wire Server RpcClient SocketRpcServer │ │ │ UNARY │ │ │ - ├── send_unary(req) ──── request (id=N) ────────────▶│ new_rpc(req, responder) + ├── send_unary(req) ──── request (id=N) ────────────▶│ new_rpc(req) -> RpcKind::Unary(future) │ │ │ - │ UnaryCompletion ◀─── response (id=N) ────────────┤ responder.unary(future) - │ .await │ (you get the response) + │ UnaryCompletion ◀─── response (id=N) ────────────┤ (the connection polls your future) + │ .await │ │ │ │ STREAMING │ │ │ - ├── send_streaming(req) ─ request (id=M) ───────────▶│ new_rpc(req, responder) + ├── send_streaming(req) ─ request (id=M) ───────────▶│ new_rpc(req) -> RpcKind::Streaming(stream) │ │ │ - │ StreamingCompletion ◀ response (id=M) ───────────┤ responder.stream(stream) + │ StreamingCompletion ◀ response (id=M) ───────────┤ (the connection polls your stream) │ .next().await ◀── response (id=M) ───────────┤ (you get the next message for response M) │ .next().await ◀── End (id=M) ────────────────┤ (stream ended) │ │ │ CANCELLATION │ │ │ - ├── send_unary(req) ──── request (id=L) ────────────▶│ new_rpc(req, responder) + ├── send_unary(req) ──── request (id=L) ────────────▶│ new_rpc(req) -> RpcKind::Unary(future) │ │ │ - ├── drop(completion) ──── Cancel (id) ──────────────▶│ (task aborted) + ├── drop(completion) ──── Cancel (id) ──────────────▶│ (rpc aborted) ``` diff --git a/protosocket-rpc/src/client/configuration.rs b/protosocket-rpc/src/client/configuration.rs index 9a793fc..aba1899 100644 --- a/protosocket-rpc/src/client/configuration.rs +++ b/protosocket-rpc/src/client/configuration.rs @@ -80,10 +80,8 @@ pub async fn connect( >, protosocket::Connection< TStreamConnector::Stream, - Codec, RpcCompletionReactor< - ::Message, - ::Message, + Codec, DoNothingMessageHandler<::Message>, >, >, @@ -113,8 +111,7 @@ where socket.set_reuse_address(true)?; let message_reactor: RpcCompletionReactor< - ::Message, - ::Message, + Codec, DoNothingMessageHandler<::Message>, > = RpcCompletionReactor::new(Default::default()); let (outbound, outbound_messages) = spillway::channel(); @@ -127,10 +124,8 @@ where // Tie outbound_messages to message_reactor via a protosocket::Connection let connection = Connection::< TStreamConnector::Stream, - Codec, RpcCompletionReactor< - ::Message, - ::Message, + Codec, DoNothingMessageHandler<::Message>, >, >::new( diff --git a/protosocket-rpc/src/client/reactor/completion_reactor.rs b/protosocket-rpc/src/client/reactor/completion_reactor.rs index 9fded35..7ec9338 100644 --- a/protosocket-rpc/src/client/reactor/completion_reactor.rs +++ b/protosocket-rpc/src/client/reactor/completion_reactor.rs @@ -5,30 +5,29 @@ use std::{ sync::{atomic::AtomicBool, Arc}, }; -use protosocket::{MessageReactor, ReactorStatus}; +use protosocket::{Codec, Decoder, Encoder, MessageReactor, ReactorStatus}; use crate::{message::ProtosocketControlCode, Message}; use super::completion_registry::{Completion, CompletionRegistry}; #[derive(Debug)] -pub struct RpcCompletionReactor +pub struct RpcCompletionReactor where - Inbound: Message, - Outbound: Message, - TUnregisteredMessageHandler: UnregisteredMessageHandler, + TCodec: Codec, + ::Message: Message, + TUnregisteredMessageHandler: UnregisteredMessageHandler::Message>, { - rpc_registry: CompletionRegistry, + rpc_registry: CompletionRegistry<::Message>, is_alive: Arc, unregistered_message_handler: TUnregisteredMessageHandler, - _phantom: PhantomData, + _phantom: PhantomData, } -impl - RpcCompletionReactor +impl RpcCompletionReactor where - Inbound: Message, - Outbound: Message, - TUnregisteredMessageHandler: UnregisteredMessageHandler, + TCodec: Codec, + ::Message: Message, + TUnregisteredMessageHandler: UnregisteredMessageHandler::Message>, { #[allow(clippy::new_without_default)] pub fn new(unregistered_message_handler: TUnregisteredMessageHandler) -> Self { @@ -45,12 +44,12 @@ where } } -impl Drop - for RpcCompletionReactor +impl Drop + for RpcCompletionReactor where - Inbound: Message, - Outbound: Message, - TUnregisteredMessageHandler: UnregisteredMessageHandler, + TCodec: Codec, + ::Message: Message, + TUnregisteredMessageHandler: UnregisteredMessageHandler::Message>, { fn drop(&mut self) { self.is_alive @@ -71,15 +70,17 @@ pub enum RpcNotification { Cancel(u64), } -impl MessageReactor - for RpcCompletionReactor +impl MessageReactor + for RpcCompletionReactor where - Inbound: Message, - Outbound: Message, - TUnregisteredMessageHandler: UnregisteredMessageHandler, + TCodec: Codec + 'static, + ::Message: Message, + ::Message: Message, + TUnregisteredMessageHandler: UnregisteredMessageHandler::Message>, { - type Inbound = Inbound; - type Outbound = Outbound; + type Codec = TCodec; + type Inbound = ::Message; + type Outbound = ::Message; type LogicalOutbound = RpcNotification; fn on_inbound_message(&mut self, message: Self::Inbound) -> ReactorStatus { @@ -101,15 +102,21 @@ where Entry::Occupied(mut registered_rpc) => { if let Completion::RemoteStreaming(stream) = registered_rpc.get_mut() { if let Err(e) = stream.send(message) { - log::debug!("{message_id} completion channel closed - did the client lose interest in this request? {e:?}"); + log::debug!( + "{message_id} completion channel closed - did the client lose interest in this request? {e:?}" + ); registered_rpc.remove(); } } else if let Completion::Unary(completion) = registered_rpc.remove() { if let Err(e) = completion.send(Ok(message)) { - log::debug!("{message_id} completion channel closed - did the client lose interest in this request? {e:?}"); + log::debug!( + "{message_id} completion channel closed - did the client lose interest in this request? {e:?}" + ); } } else { - panic!("{message_id} unexpected command response type. Sorry, I wanted to borrow for streaming and remove by value for unary without doing 2 map lookups, so I couldn't match"); + panic!( + "{message_id} unexpected command response type. Sorry, I wanted to borrow for streaming and remove by value for unary without doing 2 map lookups, so I couldn't match" + ); } } Entry::Vacant(_vacant_entry) => { @@ -137,7 +144,7 @@ where RpcNotification::Cancel(message_id) => { log::trace!("{} cancelling rpc in completion reactor", message_id); self.rpc_registry.deregister(message_id); - Outbound::cancelled(message_id) + ::cancelled(message_id) } } } diff --git a/protosocket-rpc/src/client/rpc_client.rs b/protosocket-rpc/src/client/rpc_client.rs index 602e7ce..e2fbb1c 100644 --- a/protosocket-rpc/src/client/rpc_client.rs +++ b/protosocket-rpc/src/client/rpc_client.rs @@ -43,14 +43,15 @@ where Request: Message, Response: Message, { - pub(crate) fn new( + pub(crate) fn new( submission_queue: spillway::Sender>, - message_reactor: &RpcCompletionReactor< - Response, - Request, - DoNothingMessageHandler, - >, - ) -> Self { + message_reactor: &RpcCompletionReactor>, + ) -> Self + where + TCodec: protosocket::Codec + + protosocket::Decoder + + protosocket::Encoder, + { Self { submission_queue, is_alive: message_reactor.alive_handle(), @@ -168,6 +169,34 @@ mod test { } } + /// A codec for the u64 test messages. These tests never move real bytes; this + /// only gives the reactor its codec type. + struct TestCodec; + impl protosocket::Encoder for TestCodec { + type Message = u64; + type Serialized = bytes::Bytes; + + fn encode(&mut self, message: Self::Message) -> Self::Serialized { + bytes::Bytes::copy_from_slice(&message.to_le_bytes()) + } + } + impl protosocket::Decoder for TestCodec { + type Message = u64; + + fn decode( + &mut self, + mut buffer: impl bytes::Buf, + ) -> Result<(usize, Self::Message), protosocket::DeserializeError> { + if buffer.remaining() < 8 { + return Err(protosocket::DeserializeError::IncompleteBuffer { + next_message_size: 8, + }); + } + Ok((8, buffer.get_u64_le())) + } + } + impl protosocket::Codec for TestCodec {} + fn drive_future(f: F) -> F::Output { let mut f = pin!(f); loop { @@ -182,11 +211,11 @@ mod test { fn get_client() -> ( spillway::Receiver>, RpcClient, - RpcCompletionReactor>, + RpcCompletionReactor>, ) { let (sender, remote_end) = spillway::channel(); let rpc_reactor = - RpcCompletionReactor::::new(DoNothingMessageHandler::default()); + RpcCompletionReactor::::new(DoNothingMessageHandler::default()); let client = RpcClient::new(sender, &rpc_reactor); (remote_end, client, rpc_reactor) } @@ -234,7 +263,7 @@ mod test { Vec<( spillway::Receiver>, RpcClient, - RpcCompletionReactor>, + RpcCompletionReactor>, )>, >, fail_connections: AtomicBool, diff --git a/protosocket-rpc/src/server/abortable.rs b/protosocket-rpc/src/server/abortable.rs deleted file mode 100644 index 143f756..0000000 --- a/protosocket-rpc/src/server/abortable.rs +++ /dev/null @@ -1,118 +0,0 @@ -use std::{ - future::Future, - pin::Pin, - sync::{ - atomic::{AtomicUsize, Ordering}, - Arc, - }, - task::{Context, Poll}, -}; - -use futures::{task::AtomicWaker, Stream}; - -#[derive(Debug)] -pub struct IdentifiableAbortable { - f: F, - aborted: Arc, - waker: Arc, -} - -impl IdentifiableAbortable { - pub fn new(f: F) -> (Self, IdentifiableAbortHandle) { - let aborted = Arc::new(AtomicUsize::new(0)); - let waker = Arc::new(AtomicWaker::new()); - ( - Self { - f, - aborted: aborted.clone(), - waker: waker.clone(), - }, - IdentifiableAbortHandle { aborted, waker }, - ) - } -} - -#[derive(Debug)] -pub enum AbortableState { - Abort, - Aborted, - Ready(T), -} - -impl Future for IdentifiableAbortable -where - F: Future, -{ - type Output = AbortableState>; - - fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { - let state = self.aborted.load(Ordering::Relaxed); - if 1 == state { - self.aborted.store(2, Ordering::Relaxed); - return Poll::Ready(AbortableState::Abort); - } - if 2 == state { - return Poll::Ready(AbortableState::Aborted); - } - self.waker.register(context.waker()); - // SAFETY: This is a structural pin. If I'm not moved then neither is this future. - let structurally_pinned_future = unsafe { self.as_mut().map_unchecked_mut(|me| &mut me.f) }; - structurally_pinned_future - .poll(context) - .map(|output| AbortableState::Ready(Ok(output))) - } -} - -impl Stream for IdentifiableAbortable -where - S: Stream, -{ - type Item = AbortableState>; - - fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { - self.waker.register(context.waker()); - match self.aborted.load(Ordering::Relaxed) { - 0 => { - // SAFETY: This is a structural pin. If I'm not moved then neither is this stream. - match unsafe { self.as_mut().map_unchecked_mut(|me| &mut me.f) }.poll_next(context) - { - Poll::Ready(next) => { - match next { - Some(next) => Poll::Ready(Some(AbortableState::Ready(Ok(next)))), - None => { - // stream is done - self.aborted.store(3, Ordering::Relaxed); - Poll::Ready(Some(AbortableState::Ready(Err( - crate::Error::Finished, - )))) - } - } - } - Poll::Pending => Poll::Pending, - } - } - 1 => { - self.aborted.store(2, Ordering::Relaxed); - Poll::Ready(Some(AbortableState::Abort)) - } - 2 => { - self.aborted.store(3, Ordering::Relaxed); - Poll::Ready(Some(AbortableState::Aborted)) - } - _ => Poll::Ready(None), - } - } -} - -#[derive(Debug)] -pub struct IdentifiableAbortHandle { - aborted: Arc, - waker: Arc, -} -impl IdentifiableAbortHandle { - /// Mark the future or stream as externally cancelled - don't send a cancellation - pub fn mark_aborted(self) { - self.aborted.store(2, Ordering::Relaxed); - self.waker.wake(); - } -} diff --git a/protosocket-rpc/src/server/abortion_tracker.rs b/protosocket-rpc/src/server/abortion_tracker.rs index 0267c18..cc7695e 100644 --- a/protosocket-rpc/src/server/abortion_tracker.rs +++ b/protosocket-rpc/src/server/abortion_tracker.rs @@ -1,10 +1,10 @@ use std::collections::HashMap; -use crate::server::abortable::IdentifiableAbortHandle; +use crate::server::rpc_stream::RpcAbortHandle; #[derive(Default)] pub struct AbortionTracker { - aborts: HashMap, + aborts: HashMap, } impl std::fmt::Debug for AbortionTracker { @@ -16,15 +16,11 @@ impl std::fmt::Debug for AbortionTracker { } impl AbortionTracker { - pub fn register( - &mut self, - id: u64, - handle: IdentifiableAbortHandle, - ) -> Option { + pub fn register(&mut self, id: u64, handle: RpcAbortHandle) -> Option { self.aborts.insert(id, handle) } - pub fn take_abort(&mut self, id: u64) -> Option { + pub fn take_abort(&mut self, id: u64) -> Option { self.aborts.remove(&id) } } diff --git a/protosocket-rpc/src/server/connection_server.rs b/protosocket-rpc/src/server/connection_server.rs deleted file mode 100644 index 322a08e..0000000 --- a/protosocket-rpc/src/server/connection_server.rs +++ /dev/null @@ -1,653 +0,0 @@ -use std::{ - collections::HashMap, - future::Future, - pin::{pin, Pin}, - task::{Context, Poll}, -}; - -use futures::{ - stream::{FuturesUnordered, SelectAll}, - Stream, -}; -use tokio::sync::mpsc; -use tokio_util::sync::PollSender; - -use crate::{server::RpcKind, Error, Message, ProtosocketControlCode}; - -use super::{ - abortable::{AbortableState, IdentifiableAbortHandle, IdentifiableAbortable}, - ConnectionService, -}; - -#[derive(Debug)] -pub struct RpcConnectionServer -where - TConnectionServer: ConnectionService, -{ - connection_server: TConnectionServer, - // inbound: mpsc::UnboundedReceiver<::Request>, - outbound: PollSender<::Response>, - // next_messages_buffer: Vec<::Request>, - // outstanding_unary_rpcs: - // FuturesUnordered>, - // outstanding_streaming_rpcs: SelectAll>, - aborts: HashMap, -} - -impl Future for RpcConnectionServer -where - TConnectionServer: ConnectionService, -{ - type Output = Result<(), crate::Error>; - - fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { - // receive new messages - if let Some(early_out) = self.as_mut().poll_receive_buffer(context) { - return early_out; - } - // either we're pending on inbound or we're awake - self.as_mut().handle_message_buffer(); - - // retire and advance outstanding rpcs - if let Some(early_out) = self.as_mut().poll_advance_unary_rpcs(context) { - return early_out; - } - if let Some(early_out) = self.poll_advance_streaming_rpcs(context) { - return early_out; - } - - Poll::Pending - } -} - -impl RpcConnectionServer -where - TConnectionServer: ConnectionService, -{ - pub fn new( - connection_server: TConnectionServer, - inbound: mpsc::UnboundedReceiver<::Request>, - outbound: mpsc::Sender<::Response>, - ) -> Self { - Self { - connection_server, - inbound, - outbound: PollSender::new(outbound), - next_messages_buffer: Default::default(), - outstanding_unary_rpcs: Default::default(), - outstanding_streaming_rpcs: Default::default(), - aborts: Default::default(), - } - } - - -} - -#[cfg(test)] -mod test { - use std::{ - future::Future, - pin::pin, - ptr, - task::{Context, Poll, RawWaker, RawWakerVTable, Waker}, - }; - - use futures::{FutureExt, StreamExt}; - use tokio::sync::mpsc; - - use crate::{ - server::{ConnectionService, RpcKind}, - ProtosocketControlCode, - }; - - use super::RpcConnectionServer; - - #[derive(Clone, PartialEq, Eq, prost::Message, PartialOrd, Ord)] - pub struct Message { - #[prost(uint64, tag = "1")] - pub id: u64, - #[prost(uint32, tag = "2")] - pub code: u32, - #[prost(uint64, tag = "3")] - pub n: u64, - } - - impl crate::Message for Message { - fn message_id(&self) -> u64 { - self.id - } - - fn control_code(&self) -> crate::ProtosocketControlCode { - crate::ProtosocketControlCode::from_u8(self.code as u8) - } - - fn set_message_id(&mut self, message_id: u64) { - self.id = message_id; - } - - fn cancelled(message_id: u64) -> Self { - Self { - id: message_id, - n: 0, - code: ProtosocketControlCode::Cancel.as_u8() as u32, - } - } - - fn ended(message_id: u64) -> Self { - Self { - id: message_id, - n: 0, - code: ProtosocketControlCode::End.as_u8() as u32, - } - } - } - - const HANGING_UNARY_MESSAGE: u64 = 2000; - const HANGING_STREAMING_MESSAGE: u64 = 3000; - struct TestConnectionService; - impl std::fmt::Debug for TestConnectionService { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TestConnectionService").finish() - } - } - - impl ConnectionService for TestConnectionService { - type Request = Message; - type Response = Message; - // Boxing is used for convenience in tests. You should try to use a static type in your real code. - type UnaryFutureType = futures::future::BoxFuture<'static, Message>; - type StreamType = futures::stream::BoxStream<'static, Message>; - - fn new_rpc( - &mut self, - request: Self::Request, - ) -> crate::server::RpcKind { - if request.id == HANGING_UNARY_MESSAGE { - RpcKind::Unary(futures::future::pending().boxed()) - } else if request.id == HANGING_STREAMING_MESSAGE { - RpcKind::Streaming(futures::stream::pending().boxed()) - } else if request.id < 1000 { - RpcKind::Unary( - futures::future::ready(Message { - id: request.id, - code: ProtosocketControlCode::Normal.as_u8() as u32, - n: request.n + 1, - }) - .boxed(), - ) - } else { - RpcKind::Streaming( - futures::stream::iter((0..request.n).map(move |n| Message { - id: request.id, - code: ProtosocketControlCode::Normal.as_u8() as u32, - n, - })) - .boxed(), - ) - } - } - } - - pub fn noop_waker() -> Waker { - const NOOP_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new( - |_| RawWaker::new(ptr::null(), &NOOP_WAKER_VTABLE), - |_| {}, - |_| {}, - |_| {}, - ); - let raw = RawWaker::new(ptr::null(), &NOOP_WAKER_VTABLE); - // SAFETY: the contracts for RawWaker and RawWakerVTable are trivially upheld by always making new wakers - unsafe { Waker::from_raw(raw) } - } - - fn test_server( - outbound_buffer: usize, - ) -> ( - mpsc::UnboundedSender, - mpsc::Receiver, - RpcConnectionServer, - ) { - let (inbound_sender, inbound) = mpsc::unbounded_channel(); - let (outbound, outbound_receiver) = mpsc::channel(outbound_buffer); - let server = RpcConnectionServer::new(TestConnectionService, inbound, outbound); - (inbound_sender, outbound_receiver, server) - } - - #[track_caller] - fn assert_next( - message: Message, - outbound_receiver: &mut mpsc::Receiver, - context: &mut Context<'_>, - ) { - assert_eq!( - Poll::Ready(Some(message)), - outbound_receiver.poll_recv(context) - ); - } - - #[track_caller] - fn poll_next( - outbound_receiver: &mut mpsc::Receiver, - context: &mut Context<'_>, - ) -> Message { - match outbound_receiver.poll_recv(context) { - Poll::Ready(Some(message)) => message, - got => panic!("expected message, got {got:?}"), - } - } - - #[test] - fn unary() { - let waker = noop_waker(); - let mut context = Context::from_waker(&waker); - - let (inbound_sender, mut outbound_receiver, mut server) = test_server(3); - - // test messages below 1000 are unary. Response is n + 1 - let _ = inbound_sender.send(Message { - id: 1, - code: 0, - n: 1, - }); - - assert_eq!( - Poll::Pending, - outbound_receiver.poll_recv(&mut context), - "nothing should be sent until the server advances to accept the message" - ); - - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - assert_eq!( - 0, - server.outstanding_unary_rpcs.len(), - "it completed in one poll" - ); - - assert_next( - Message { - id: 1, - code: 0, - n: 2, - }, - &mut outbound_receiver, - &mut context, - ); - } - - #[test] - fn concurrent_unary() { - let waker = noop_waker(); - let mut context = Context::from_waker(&waker); - - let (inbound_sender, mut outbound_receiver, mut server) = test_server(3); - - let _ = inbound_sender.send(Message { - id: 1, - code: 0, - n: 1, - }); - let _ = inbound_sender.send(Message { - id: 2, - code: 0, - n: 3, - }); - let _ = inbound_sender.send(Message { - id: 3, - code: 0, - n: 5, - }); - - // the server takes up to MAXIMUM_MESSAGES_PER_POLL per poll. I only submitted 3, so they should - // all get processed in the a single round of poll. - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - assert_eq!( - 0, - server.outstanding_unary_rpcs.len(), - "it completed in one poll" - ); - - let mut concurrent_completions = vec![ - poll_next(&mut outbound_receiver, &mut context), - poll_next(&mut outbound_receiver, &mut context), - poll_next(&mut outbound_receiver, &mut context), - ]; - // they are allowed to complete in any order but I'd like a deterministic order for the assertion - concurrent_completions.sort(); - - assert_eq!( - vec![ - Message { - id: 1, - code: 0, - n: 2 - }, - Message { - id: 2, - code: 0, - n: 4 - }, - Message { - id: 3, - code: 0, - n: 6 - }, - ], - concurrent_completions, - ); - assert_eq!( - Poll::Pending, - outbound_receiver.poll_recv(&mut context), - "no made up messages" - ); - } - - #[test] - fn streaming() { - let waker = noop_waker(); - let mut context = Context::from_waker(&waker); - - let (inbound_sender, mut outbound_receiver, mut server) = test_server(3); - // "test" messages at and above 1000 are streaming. Stream has responses n=0..n - let _ = inbound_sender.send(Message { - id: 1000, - code: 0, - n: 2, - }); - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - - let first_message = poll_next(&mut outbound_receiver, &mut context); - assert_eq!( - 1, - server.outstanding_streaming_rpcs.len(), - "there should still be an outstanding rpc because the stream is not done" - ); - let messages = vec![ - first_message, - poll_next(&mut outbound_receiver, &mut context), - poll_next(&mut outbound_receiver, &mut context), - ]; - // these must come in the correct order. - - assert_eq!( - vec![ - Message { - id: 1000, - code: 0, - n: 0 - }, - Message { - id: 1000, - code: 0, - n: 1 - }, - Message { - id: 1000, - code: ProtosocketControlCode::End.as_u8() as u32, - n: 0 - }, - ], - messages, - ); - - assert_eq!(1, server.outstanding_streaming_rpcs.len(), "server has not yet discovered that this rpc is complete. This might change if the poll batch process is changed"); - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - assert_eq!( - 0, - server.outstanding_streaming_rpcs.len(), - "all rpcs should be completed" - ); - assert_eq!( - Poll::Pending, - outbound_receiver.poll_recv(&mut context), - "no made up messages" - ); - } - - #[test] - fn streaming_concurrent() { - let waker = noop_waker(); - let mut context = Context::from_waker(&waker); - - let (inbound_sender, mut outbound_receiver, mut server) = test_server(3); - // "test" messages at and above 1000 are streaming. Stream has responses n=0..n - let _ = inbound_sender.send(Message { - id: 1000, - code: 0, - n: 2, - }); - let _ = inbound_sender.send(Message { - id: 1001, - code: 0, - n: 2, - }); - let _ = inbound_sender.send(Message { - id: 1002, - code: 0, - n: 2, - }); - - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - assert_eq!(3, server.outstanding_streaming_rpcs.len()); - - let mut messages = vec![ - poll_next(&mut outbound_receiver, &mut context), - poll_next(&mut outbound_receiver, &mut context), - poll_next(&mut outbound_receiver, &mut context), - ]; - assert_eq!( - Poll::Pending, - outbound_receiver.poll_recv(&mut context), - "outbound buffer is only 3. It is unknown if any of the rpcs are complete" - ); - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - messages.push(poll_next(&mut outbound_receiver, &mut context)); - messages.push(poll_next(&mut outbound_receiver, &mut context)); - messages.push(poll_next(&mut outbound_receiver, &mut context)); - assert_eq!(Poll::Pending, outbound_receiver.poll_recv(&mut context), "though we only defined 6 messages, the server sends an End message for each gracefully ended stream"); - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - messages.push(poll_next(&mut outbound_receiver, &mut context)); - messages.push(poll_next(&mut outbound_receiver, &mut context)); - messages.push(poll_next(&mut outbound_receiver, &mut context)); - - // The messages may be intermixed per-rpc, but they must be mutually in order per-rpc. - // It is a weak assertion to sort these, because that would allow _reordered streams_ to pass the test. - let first_rpc: Vec<_> = messages - .iter() - .filter(|message| message.id == 1000) - .cloned() - .collect(); - let second_rpc: Vec<_> = messages - .iter() - .filter(|message| message.id == 1001) - .cloned() - .collect(); - let third_rpc: Vec<_> = messages - .iter() - .filter(|message| message.id == 1002) - .cloned() - .collect(); - - assert_eq!( - vec![ - Message { - id: 1000, - code: 0, - n: 0 - }, - Message { - id: 1000, - code: 0, - n: 1 - }, - Message { - id: 1000, - code: ProtosocketControlCode::End.as_u8() as u32, - n: 0 - }, - ], - first_rpc, - ); - assert_eq!( - vec![ - Message { - id: 1001, - code: 0, - n: 0 - }, - Message { - id: 1001, - code: 0, - n: 1 - }, - Message { - id: 1001, - code: ProtosocketControlCode::End.as_u8() as u32, - n: 0 - }, - ], - second_rpc, - ); - assert_eq!( - vec![ - Message { - id: 1002, - code: 0, - n: 0 - }, - Message { - id: 1002, - code: 0, - n: 1 - }, - Message { - id: 1002, - code: ProtosocketControlCode::End.as_u8() as u32, - n: 0 - }, - ], - third_rpc, - ); - // server may have 0-3 pending rpcs, but they should all complete with the next poll. - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - assert_eq!( - 0, - server.outstanding_streaming_rpcs.len(), - "all rpcs should be completed" - ); - assert_eq!( - Poll::Pending, - outbound_receiver.poll_recv(&mut context), - "no made up messages" - ); - } - - // This test makes sure that the server drops a unary rpc when it asked to do so. - #[test] - fn unary_client_cancellation() { - let waker = noop_waker(); - let mut context = Context::from_waker(&waker); - - let (inbound_sender, mut outbound_receiver, mut server) = test_server(3); - - let _ = inbound_sender.send(Message { - id: HANGING_UNARY_MESSAGE, - code: 0, - n: 1, - }); - assert!(pin!(&mut server).poll(&mut context).is_pending()); - - assert_eq!( - 1, - server.outstanding_unary_rpcs.len(), - "it will never complete" - ); - - let _ = inbound_sender.send(Message { - id: HANGING_UNARY_MESSAGE, - code: ProtosocketControlCode::Cancel.as_u8() as u32, - n: 0, - }); - - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - assert_eq!( - 0, - server.outstanding_unary_rpcs.len(), - "all rpcs should be completed" - ); - assert_eq!( - Poll::Pending, - outbound_receiver.poll_recv(&mut context), - "no made up messages" - ); - } - - // This test makes sure that the server drops a streaming rpc when it asked to do so. - #[test] - fn streaming_client_cancellation() { - let waker = noop_waker(); - let mut context = Context::from_waker(&waker); - - let (inbound_sender, mut outbound_receiver, mut server) = test_server(3); - - let _ = inbound_sender.send(Message { - id: HANGING_STREAMING_MESSAGE, - code: 0, - n: 1, - }); - assert!(pin!(&mut server).poll(&mut context).is_pending()); - - assert_eq!( - 1, - server.outstanding_streaming_rpcs.len(), - "it will never complete" - ); - - let _ = inbound_sender.send(Message { - id: HANGING_STREAMING_MESSAGE, - code: ProtosocketControlCode::Cancel.as_u8() as u32, - n: 0, - }); - - assert!( - pin!(&mut server).poll(&mut context).is_pending(), - "server should be pending forever" - ); - assert_eq!( - 0, - server.outstanding_streaming_rpcs.len(), - "all rpcs should be completed" - ); - assert_eq!( - Poll::Pending, - outbound_receiver.poll_recv(&mut context), - "no made up messages" - ); - } -} diff --git a/protosocket-rpc/src/server/forward_streaming.rs b/protosocket-rpc/src/server/forward_streaming.rs deleted file mode 100644 index 062b03d..0000000 --- a/protosocket-rpc/src/server/forward_streaming.rs +++ /dev/null @@ -1,139 +0,0 @@ -use std::{ - future::Future, - pin::Pin, - task::{Context, Poll}, -}; - -use futures::Stream; - -use crate::{ - server::{abortable::AbortableState, rpc_submitter::RpcResponse}, - Message, -}; - -pub struct ForwardAbortableStreamingRpc -where - S: Stream>>, - T: Message, -{ - stream: S, - id: u64, - forward: spillway::Sender>, - completed_for_drop: bool, -} -impl Drop for ForwardAbortableStreamingRpc -where - S: Stream>>, - T: Message, -{ - fn drop(&mut self) { - if !self.completed_for_drop { - log::debug!("dropping unary rpc before completion: {}", self.id); - let _ = self.forward.send(RpcResponse::Final(T::cancelled(self.id))); - } - } -} -impl ForwardAbortableStreamingRpc -where - S: Stream>>, - T: Message, -{ - pub fn new(stream: S, id: u64, forward: spillway::Sender>) -> Self { - Self { - stream, - id, - forward, - completed_for_drop: false, - } - } - - fn complete_for_drop(self: Pin<&mut Self>) { - // SAFETY: This is a structural pin. If I'm not moved then neither is this boolean (it was an invariant for the future anyway). - unsafe { - self.get_unchecked_mut().completed_for_drop = true; - } - } -} -impl Future for ForwardAbortableStreamingRpc -where - S: Stream>>, - T: Message, -{ - type Output = (); - - fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { - loop { - // SAFETY: This is a structural pin. If I'm not moved then neither is this future. - break match unsafe { self.as_mut().map_unchecked_mut(|me| &mut me.stream) } - .poll_next(context) - { - Poll::Ready(state) => { - match state { - Some(AbortableState::Ready(Ok(response))) => { - log::trace!("{} unary rpc response", self.id); - if let Err(_e) = self.forward.send(RpcResponse::Partial(response)) { - log::debug!("outbound connection is closed"); - } - continue; - } - Some(AbortableState::Ready(Err(e))) => { - match e { - crate::Error::IoFailure(error) => { - log::warn!( - "{} io failure while servicing rpc: {error:?}", - self.id - ); - let _ = self - .forward - .send(RpcResponse::Final(T::cancelled(self.id))); - } - crate::Error::CancelledRemotely => { - log::debug!("{} rpc cancelled remotely", self.id); - let _ = self - .forward - .send(RpcResponse::Final(T::cancelled(self.id))); - } - crate::Error::ConnectionIsClosed => { - log::debug!("{} rpc cancelled remotely", self.id); - let _ = self - .forward - .send(RpcResponse::Final(T::cancelled(self.id))); - } - crate::Error::Finished => { - log::debug!("{} unary rpc ended", self.id); - if let Err(_e) = - self.forward.send(RpcResponse::Final(T::ended(self.id))) - { - log::debug!("outbound connection is closed"); - } - } - } - self.complete_for_drop(); - Poll::Ready(()) - } - Some(AbortableState::Abort) => { - // This happens when the upstream stuff is dropped and there are no messages that can be produced. We'll send a cancellation. - log::debug!("{} unary rpc abort", self.id); - let _ = self.forward.send(RpcResponse::Final(T::cancelled(self.id))); - self.complete_for_drop(); - Poll::Ready(()) - } - Some(AbortableState::Aborted) => { - log::debug!("{} unary rpc was cancelled by the client.", self.id); - let _ = self.forward.send(RpcResponse::Final(T::cancelled(self.id))); - self.complete_for_drop(); - Poll::Ready(()) - } - None => { - log::debug!("{} streaming rpc reached the end", self.id); - let _ = self.forward.send(RpcResponse::Final(T::ended(self.id))); - self.complete_for_drop(); - Poll::Ready(()) - } - } - } - Poll::Pending => Poll::Pending, - }; - } - } -} diff --git a/protosocket-rpc/src/server/forward_unary.rs b/protosocket-rpc/src/server/forward_unary.rs deleted file mode 100644 index f261221..0000000 --- a/protosocket-rpc/src/server/forward_unary.rs +++ /dev/null @@ -1,119 +0,0 @@ -use std::{ - future::Future, - pin::Pin, - task::{Context, Poll}, -}; - -use crate::{ - server::{abortable::AbortableState, rpc_submitter::RpcResponse}, - Message, -}; - -pub struct ForwardAbortableUnaryRpc -where - F: Future>>, - T: Message, -{ - future: F, - id: u64, - forward: spillway::Sender>, - completed_for_drop: bool, -} -impl Drop for ForwardAbortableUnaryRpc -where - F: Future>>, - T: Message, -{ - fn drop(&mut self) { - if !self.completed_for_drop { - log::debug!("dropping unary rpc before completion: {}", self.id); - let _ = self.forward.send(RpcResponse::Final(T::cancelled(self.id))); - } - } -} -impl ForwardAbortableUnaryRpc -where - F: Future>>, - T: Message, -{ - pub fn new(future: F, id: u64, forward: spillway::Sender>) -> Self { - Self { - future, - id, - forward, - completed_for_drop: false, - } - } -} -impl Future for ForwardAbortableUnaryRpc -where - F: Future>>, - T: Message, -{ - type Output = (); - - fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { - // SAFETY: This is a structural pin. If I'm not moved then neither is this future. - let structurally_pinned_future = - unsafe { self.as_mut().map_unchecked_mut(|me| &mut me.future) }; - match structurally_pinned_future.poll(context) { - Poll::Ready(state) => { - // SAFETY: This is a structural pin. If I'm not moved then neither is this boolean (it was an invariant for the future anyway). - unsafe { - self.as_mut().get_unchecked_mut().completed_for_drop = true; - } - - match state { - AbortableState::Ready(Ok(response)) => { - log::trace!("{} unary rpc response", self.id); - if let Err(_e) = self.forward.send(RpcResponse::Final(response)) { - log::debug!("outbound connection is closed"); - } - Poll::Ready(()) - } - AbortableState::Ready(Err(e)) => { - match e { - crate::Error::IoFailure(error) => { - log::warn!("{} io failure while servicing rpc: {error:?}", self.id); - let _ = - self.forward.send(RpcResponse::Final(T::cancelled(self.id))); - } - crate::Error::CancelledRemotely => { - log::debug!("{} rpc cancelled remotely", self.id); - let _ = - self.forward.send(RpcResponse::Final(T::cancelled(self.id))); - } - crate::Error::ConnectionIsClosed => { - log::debug!("{} rpc cancelled remotely", self.id); - let _ = - self.forward.send(RpcResponse::Final(T::cancelled(self.id))); - } - crate::Error::Finished => { - log::debug!("{} unary rpc ended", self.id); - if let Err(_e) = - self.forward.send(RpcResponse::Final(T::ended(self.id))) - { - log::debug!("outbound connection is closed"); - } - } - } - Poll::Ready(()) - // cancelled - } - AbortableState::Abort => { - // This happens when the upstream stuff is dropped and there are no messages that can be produced. We'll send a cancellation. - log::debug!("{} unary rpc abort", self.id); - let _ = self.forward.send(RpcResponse::Final(T::cancelled(self.id))); - Poll::Ready(()) - } - AbortableState::Aborted => { - log::debug!("{} unary rpc was cancelled by the client.", self.id); - let _ = self.forward.send(RpcResponse::Final(T::cancelled(self.id))); - Poll::Ready(()) - } - } - } - Poll::Pending => Poll::Pending, - } - } -} diff --git a/protosocket-rpc/src/server/mod.rs b/protosocket-rpc/src/server/mod.rs index 3886d58..758f034 100644 --- a/protosocket-rpc/src/server/mod.rs +++ b/protosocket-rpc/src/server/mod.rs @@ -2,17 +2,13 @@ //! //! See example-proto or example-messagepack for how to make servers. -mod abortable; mod abortion_tracker; -mod forward_streaming; -mod forward_unary; -mod rpc_responder; +mod rpc_stream; mod rpc_submitter; mod server_traits; mod socket_server; mod spawn; -pub use rpc_responder::RpcResponder; -pub use server_traits::{ConnectionService, SocketService}; +pub use server_traits::{ConnectionService, RpcKind, SocketService}; pub use socket_server::SocketRpcServer; pub use spawn::{LevelSpawn, Spawn, TokioSpawn}; diff --git a/protosocket-rpc/src/server/rpc_responder.rs b/protosocket-rpc/src/server/rpc_responder.rs deleted file mode 100644 index 53196c6..0000000 --- a/protosocket-rpc/src/server/rpc_responder.rs +++ /dev/null @@ -1,66 +0,0 @@ -use std::future::Future; - -use crate::{ - server::{ - abortable::IdentifiableAbortable, abortion_tracker::AbortionTracker, - forward_streaming::ForwardAbortableStreamingRpc, forward_unary::ForwardAbortableUnaryRpc, - rpc_submitter::RpcResponse, - }, - Message, -}; - -/// A request context's temporary lease to an RPC Reactor's state. -/// You want to consume your RpcResponder as quickly as possible. -#[must_use] -pub struct RpcResponder<'a, Response> { - outbound: &'a spillway::Sender>, - aborts: &'a mut AbortionTracker, - message_id: u64, -} -impl<'a, Response> RpcResponder<'a, Response> -where - Response: Message, -{ - pub(crate) fn new_responder_reference( - outbound: &'a spillway::Sender>, - aborts: &'a mut AbortionTracker, - message_id: u64, - ) -> Self { - Self { - outbound, - aborts, - message_id, - } - } - - /// Consume the responder by providing a future that will materialize the response to this request. - pub fn unary(self, unary_rpc: impl Future) -> impl Future { - let (abortable, abort) = IdentifiableAbortable::new(unary_rpc); - self.aborts.register(self.message_id, abort); - ForwardAbortableUnaryRpc::new(abortable, self.message_id, self.outbound.clone()) - } - - /// Consume the responder by providing a stream that will materialize the response to this request. - pub fn stream( - self, - streaming_rpc: impl futures::Stream, - ) -> impl Future { - let (abortable_stream, abort) = IdentifiableAbortable::new(streaming_rpc); - self.aborts.register(self.message_id, abort); - ForwardAbortableStreamingRpc::new(abortable_stream, self.message_id, self.outbound.clone()) - } - - /// Consume the responder by providing an immediate response. - /// - /// This is the cheapest, fastest way to respond, but you must only use it when you can get a response - /// without blocking! - pub fn immediate(self, response: Response) { - if self - .outbound - .send(RpcResponse::Untracked(response)) - .is_err() - { - log::debug!("outbound channel closed while sending response"); - } - } -} diff --git a/protosocket-rpc/src/server/rpc_stream.rs b/protosocket-rpc/src/server/rpc_stream.rs new file mode 100644 index 0000000..2fc0b49 --- /dev/null +++ b/protosocket-rpc/src/server/rpc_stream.rs @@ -0,0 +1,137 @@ +use std::{ + future::Future, + pin::{pin, Pin}, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + task::{Context, Poll}, +}; + +use futures::{task::AtomicWaker, Stream}; + +/// A pooled rpc completion: unary and streaming rpcs presented as one shape. +/// +/// A unary rpc is a stream that yields exactly one terminal `Complete` event. A streaming +/// rpc yields `Item` events until the underlying stream finishes with a terminal `Finished`. +/// Cancellation yields a terminal `Cancelled` from either kind. After a terminal event the +/// stream yields `None`. +#[derive(Debug)] +pub struct RpcStream { + id: u64, + aborted: Arc, + waker: Arc, + kind: RpcStreamKind, +} + +#[derive(Debug)] +enum RpcStreamKind { + Unary(F), + Streaming(S), + Done, +} + +/// An event from a pooled rpc. +#[derive(Debug)] +pub enum RpcStreamEvent { + /// A streaming rpc produced a response item. Stream continues. + Item(T), + /// A unary rpc completed with its single response. Terminal. + Complete(T), + /// A streaming rpc's stream finished. Terminal. + Finished, + /// The rpc was cancelled. Terminal. + Cancelled, +} + +impl RpcStreamEvent { + pub fn is_terminal(&self) -> bool { + match self { + RpcStreamEvent::Item(_) => false, + RpcStreamEvent::Complete(_) | RpcStreamEvent::Finished | RpcStreamEvent::Cancelled => { + true + } + } + } +} + +impl RpcStream { + pub fn new_unary(id: u64, completion: F) -> (Self, RpcAbortHandle) { + Self::new(id, RpcStreamKind::Unary(completion)) + } + + pub fn new_streaming(id: u64, stream: S) -> (Self, RpcAbortHandle) { + Self::new(id, RpcStreamKind::Streaming(stream)) + } + + fn new(id: u64, kind: RpcStreamKind) -> (Self, RpcAbortHandle) { + let aborted = Arc::new(AtomicBool::new(false)); + let waker = Arc::new(AtomicWaker::new()); + ( + Self { + id, + aborted: aborted.clone(), + waker: waker.clone(), + kind, + }, + RpcAbortHandle { aborted, waker }, + ) + } +} + +impl Stream for RpcStream +where + F: Future + Unpin, + S: Stream + Unpin, +{ + type Item = (u64, RpcStreamEvent); + + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + let id = self.id; + self.waker.register(context.waker()); + if self.aborted.load(Ordering::Relaxed) { + return if matches!(self.kind, RpcStreamKind::Done) { + Poll::Ready(None) + } else { + self.kind = RpcStreamKind::Done; + Poll::Ready(Some((id, RpcStreamEvent::Cancelled))) + }; + } + match &mut self.kind { + RpcStreamKind::Unary(completion) => match pin!(completion).poll(context) { + Poll::Ready(response) => { + self.kind = RpcStreamKind::Done; + Poll::Ready(Some((id, RpcStreamEvent::Complete(response)))) + } + Poll::Pending => Poll::Pending, + }, + RpcStreamKind::Streaming(stream) => match pin!(stream).poll_next(context) { + Poll::Ready(Some(item)) => Poll::Ready(Some((id, RpcStreamEvent::Item(item)))), + Poll::Ready(None) => { + self.kind = RpcStreamKind::Done; + Poll::Ready(Some((id, RpcStreamEvent::Finished))) + } + Poll::Pending => Poll::Pending, + }, + RpcStreamKind::Done => Poll::Ready(None), + } + } +} + +/// Cancels a pooled rpc. The rpc yields a terminal `Cancelled` event on its next poll. +#[derive(Debug)] +pub struct RpcAbortHandle { + aborted: Arc, + waker: Arc, +} + +impl RpcAbortHandle { + /// Mark the rpc aborted and wake it so the cancellation is observed. + /// + /// This races with normal rpc completion. If an rpc has already completed + /// by the time it observes cancel, it doesn't also cancel itself. + pub fn mark_aborted(self) { + self.aborted.store(true, Ordering::Relaxed); + self.waker.wake(); + } +} diff --git a/protosocket-rpc/src/server/rpc_submitter.rs b/protosocket-rpc/src/server/rpc_submitter.rs index afcc2b3..ae21ded 100644 --- a/protosocket-rpc/src/server/rpc_submitter.rs +++ b/protosocket-rpc/src/server/rpc_submitter.rs @@ -1,63 +1,123 @@ -use protosocket::MessageReactor; +use std::{ + pin::{pin, Pin}, + task::{Context, Poll}, +}; + +use futures::stream::{FuturesUnordered, StreamFuture}; +use futures::{Stream, StreamExt}; +use protosocket::{MessageReactor, SendBudget}; use crate::{ - server::{abortion_tracker::AbortionTracker, rpc_responder::RpcResponder, ConnectionService}, + server::{abortion_tracker::AbortionTracker, ConnectionService, RpcKind}, Message, ProtosocketControlCode, }; -/// A MessageReactor that sends RPCs along to a sink -#[derive(Debug)] -pub struct RpcSubmitter +use super::rpc_stream::{RpcStream, RpcStreamEvent}; + +type PooledRpc = RpcStream< + ::UnaryFutureType, + ::StreamType, +>; + +/// A MessageReactor that hosts a ConnectionService's rpcs. +/// +/// New rpcs are registered from inbound messages. Their completions live in a pool and +/// are advanced by `poll_outbound_many`. The connection only polls outbound when it +/// has room to send. If you spawn your rpc's, you may want to limit concurrent RPCs to +/// avoid spawning unbounded work against a write-blocked connection. +/// +/// Rpc messages are sent in readiness order. +pub struct RpcSubmitter where - TConnectionServer: ConnectionService, + TConnectionService: ConnectionService, { - connection_server: TConnectionServer, - outbound: spillway::Sender::Response>>, + connection_server: TConnectionService, aborts: AbortionTracker, + rejections: Vec, + rpcs: FuturesUnordered>>, } + +impl std::fmt::Debug for RpcSubmitter +where + TConnectionService: ConnectionService, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RpcSubmitter") + .field("aborts", &self.aborts) + .field("rejections", &self.rejections.len()) + .field("rpcs", &self.rpcs.len()) + .finish() + } +} + impl RpcSubmitter where TConnectionService: ConnectionService, { - pub fn new( - connection_server: TConnectionService, - outbound: spillway::Sender>, - ) -> Self { + pub fn new(connection_server: TConnectionService) -> Self { Self { connection_server, - outbound, aborts: Default::default(), + rejections: Default::default(), + rpcs: Default::default(), } } -} -pub enum RpcResponse { - Partial(T), - Final(T), - Untracked(T), + fn message_for_event( + &mut self, + id: u64, + event: RpcStreamEvent, + ) -> TConnectionService::Response { + match event { + RpcStreamEvent::Item(mut item) => { + item.set_message_id(id); + item + } + RpcStreamEvent::Complete(mut response) => { + let _ = self.aborts.take_abort(id); + response.set_message_id(id); + response + } + RpcStreamEvent::Finished => { + let _ = self.aborts.take_abort(id); + ::ended(id) + } + RpcStreamEvent::Cancelled => { + let _ = self.aborts.take_abort(id); + ::cancelled(id) + } + } + } } impl MessageReactor for RpcSubmitter where TConnectionService: ConnectionService, { + type Codec = TConnectionService::Codec; type Inbound = TConnectionService::Request; type Outbound = TConnectionService::Response; - type LogicalOutbound = RpcResponse; + type LogicalOutbound = TConnectionService::Response; fn on_inbound_message(&mut self, message: Self::Inbound) -> protosocket::ReactorStatus { let message_id = message.message_id(); match message.control_code() { - ProtosocketControlCode::Normal => { - self.connection_server.new_rpc( - message, - RpcResponder::new_responder_reference( - &self.outbound, - &mut self.aborts, - message_id, - ), - ); - } + ProtosocketControlCode::Normal => match self.connection_server.new_rpc(message) { + RpcKind::Unary(completion) => { + let (rpc, handle) = RpcStream::new_unary(message_id, completion); + self.aborts.register(message_id, handle); + self.rpcs.push(rpc.into_future()); + } + RpcKind::Streaming(stream) => { + let (rpc, handle) = RpcStream::new_streaming(message_id, stream); + self.aborts.register(message_id, handle); + self.rpcs.push(rpc.into_future()); + } + RpcKind::Cancelled => { + log::debug!("rejecting rpc {message_id}"); + self.rejections.push(message_id); + } + }, ProtosocketControlCode::Cancel => { if let Some(abort) = self.aborts.take_abort(message_id) { log::debug!("cancelling message {message_id}"); @@ -74,114 +134,85 @@ where } fn on_outbound_message(&mut self, response: Self::LogicalOutbound) -> Self::Outbound { - match response { - RpcResponse::Partial(message) => message, - RpcResponse::Untracked(message) => message, - RpcResponse::Final(message) => { - if self.aborts.take_abort(message.message_id()).is_none() { - log::debug!( - "final response for untracked message {}", - message.message_id() - ); - } - message - } - } + response } - fn poll( - mut self: std::pin::Pin<&mut Self>, - context: &mut std::task::Context<'_>, - ) -> std::ops::ControlFlow<()> { - // SAFETY: This is a structural pin. If I'm not moved then neither is this future. + fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> std::ops::ControlFlow<()> { + // SAFETY: This is a structural pin. If I'm not moved then neither is this service. let structurally_pinned_connection_server = unsafe { self.as_mut() .map_unchecked_mut(|me| &mut me.connection_server) }; structurally_pinned_connection_server.poll(context) } -} -impl RpcSubmitter -where - TConnectionService: ConnectionService, -{ - // fn poll_advance_streaming_rpcs( - // mut self: Pin<&mut Self>, - // context: &mut Context<'_>, - // ) -> Option>> { - // if self.outstanding_streaming_rpcs.is_empty() { - // log::trace!("no outstanding streaming rpcs to advance"); - // return None; - // } - // while let Poll::Ready(streaming_next) = - // futures::Stream::poll_next(pin!(&mut self.outstanding_streaming_rpcs), context) - // { - // match streaming_next { - // Some((id, AbortableState::Ready(Ok(next)))) => { - // log::debug!("{id} streaming rpc next {next:?}"); - // if let Err(_e) = self.outbound.send(next) { - // log::debug!("outbound connection is closed"); - // return Some(Poll::Ready(Err(crate::Error::ConnectionIsClosed))); - // } - // } - // Some((id, AbortableState::Ready(Err(e)))) => { - // let abort = self.aborts.remove(&id); - // match e { - // crate::Error::IoFailure(error) => { - // log::warn!("{id} io failure while servicing rpc: {error:?}"); - // if let Some(abort) = abort { - // abort.abort(); - // } - // } - // crate::Error::CancelledRemotely => { - // log::debug!("{id} rpc cancelled remotely"); - // if let Some(abort) = abort { - // abort.abort(); - // } - // } - // crate::Error::ConnectionIsClosed => { - // log::debug!("{id} rpc cancelled remotely"); - // if let Some(abort) = abort { - // abort.abort(); - // } - // } - // crate::Error::Finished => { - // log::debug!("{id} streaming rpc ended"); - // if let Some(abort) = abort { - // if let Err(_e) = self - // .outbound - // .send(::ended(id)) - // { - // log::debug!("outbound connection is closed"); - // return Some(Poll::Ready(Err( - // crate::Error::ConnectionIsClosed, - // ))); - // } - // abort.mark_aborted(); - // } - // } - // } - // } - // Some((id, AbortableState::Abort)) => { - // // This happens when the upstream stuff is dropped and there are no messages that can be produced. We'll send a cancellation. - // log::debug!("{id} streaming rpc abort"); - // if let Some(abort) = self.aborts.remove(&id) { - // abort.abort(); - // } - // } - // Some((id, AbortableState::Aborted)) => { - // log::debug!("{id} streaming rpc done"); - // if let Some(abort) = self.aborts.remove(&id) { - // abort.mark_aborted(); - // } - // } - // None => { - // // nothing to wait for - // break; - // } - // } - // } - // None - // } + fn poll_outbound_many( + self: Pin<&mut Self>, + context: &mut Context<'_>, + outbound: &mut SendBudget<'_, TConnectionService::Codec>, + ) -> Poll> { + let me = self.get_mut(); + let mut produced = 0usize; + // First, yield any rejected RPC responses. + while !me.rejections.is_empty() { + let Some(permit) = outbound.reserve() else { + break; + }; + let Some(message_id) = me.rejections.pop() else { + break; + }; + permit.send(::cancelled(message_id)); + produced += 1; + } + // Then poll ready RPCs until we either run out of room to send responses or run out of rpcs. + while let Some(permit) = outbound.reserve() { + match pin!(&mut me.rpcs).poll_next(context) { + Poll::Ready(Some((first, mut rpc))) => { + let Some((id, event)) = first else { + // StreamFuture yielded None: this rpc's stream was already + // exhausted, and there's no more work to do for it. + continue; + }; + let terminal = event.is_terminal(); + permit.send(me.message_for_event(id, event)); + produced += 1; + if terminal { + continue; + } + // If an RPC is ready, continue to poll it while it has ready responses, so that + // streaming RPCs with many ready responses don't have to retransit the rpc pool + // for each response. Track whether the RPC has terminated - if it's exhausted, we + // don't need to put it back. + let retired = loop { + let Some(permit) = outbound.reserve() else { + break false; + }; + match pin!(&mut rpc).poll_next(context) { + Poll::Ready(Some((id, event))) => { + let terminal = event.is_terminal(); + permit.send(me.message_for_event(id, event)); + produced += 1; + if terminal { + break true; + } + } + Poll::Ready(None) => break true, + Poll::Pending => break false, + } + }; + if !retired { + me.rpcs.push(rpc.into_future()); + } + } + // An empty pool is not a finished reactor: new rpcs arrive from inbound + // processing, which wakes this connection. + Poll::Ready(None) | Poll::Pending => break, + } + } + if 0 < produced { + Poll::Ready(Some(())) + } else { + Poll::Pending + } + } } diff --git a/protosocket-rpc/src/server/server_traits.rs b/protosocket-rpc/src/server/server_traits.rs index c811a46..f037977 100644 --- a/protosocket-rpc/src/server/server_traits.rs +++ b/protosocket-rpc/src/server/server_traits.rs @@ -1,6 +1,9 @@ +use std::future::Future; + +use futures::Stream; use protosocket::{Codec, Decoder, Encoder, SocketListener}; -use crate::{server::rpc_responder::RpcResponder, Message}; +use crate::Message; /// SocketService receives connections and produces ConnectionServices. /// @@ -8,70 +11,75 @@ use crate::{server::rpc_responder::RpcResponder, Message}; /// remote peer and it returns a ConnectionService for that connection. You can think of this as the /// "connection factory" for your server. It is the "top" of your service stack. pub trait SocketService: 'static { - /// Message encoding scheme - /// - /// Consider pooling your allocations, like with `protosocket::PooledEncoder`. - /// The write out to the network uses the raw `Encoder::Serialized` type, so you - /// can make outbound messages low-allocation via simple pooling. - type Codec: Codec + Decoder + Encoder; - /// The type of connection service that will be created for each connection. - type ConnectionService: ConnectionService< - Request = ::Message, - Response = ::Message, - >; + type ConnectionService: ConnectionService; /// The listener type for this service. E.g., `TcpSocketListener` type SocketListener: SocketListener; /// Create a new message codec for a connection. - fn codec(&self) -> Self::Codec; + fn codec(&mut self) -> ::Codec; /// Create a new ConnectionService for your new connection. /// The Stream will be wired into a `protosocket::Connection`. You can look at it in here /// if it has data you want (like a SocketAddr). fn new_stream_service( - &self, + &mut self, _stream: &::Stream, ) -> Self::ConnectionService; } -/// A connection service receives rpcs from clients and sends responses. +/// A connection service receives rpcs from clients and returns the work that produces +/// responses. +/// +/// Each client connection gets a ConnectionService; put your per-connection state here. +/// You are called with each initiating message, and you return the rpc that completes it: +/// a future for a unary rpc, or a stream for a streaming rpc. /// -/// Each client connection gets a ConnectionService. You put your per-connection state in your -/// ConnectionService implementation. +/// The connection drives your rpcs, and only polls them when it has room to send. A peer +/// that stops receiving stops its rpcs. If your stream models lag (like +/// `tokio::sync::broadcast`), a slow peer engages that model instead of buffering +/// responses without bound. /// -/// Every interaction with a client is done via an RPC. You are called with the initiating message -/// from the client, and you return the kind of response future that is used to complete the RPC. +/// Rpcs are polled in the connection's task. If an rpc is compute-heavy, you can spawn +/// it and return a completion future or a channel-backed stream. /// -/// A ConnectionService is executed in the context of an RPC connection server, which is a future. -/// This means you get `&mut self` when you are called with a new rpc. You can use simple mutable -/// state per-connection; but if you need to share state between connections or elsewhere in your -/// application, you will need to use an appropriate state sharing mechanism. +/// Response message ids are stamped by the connection. pub trait ConnectionService: Unpin + 'static { + /// The wire codec for connections hosting this service. + /// + /// Consider pooling your allocations, like with `protosocket::PooledEncoder`. + /// The write out to the network uses the raw `Encoder::Serialized` type, so you + /// can make outbound messages low-allocation via simple pooling. + type Codec: Codec + + Decoder + + Encoder + + 'static; /// The type of request message, These messages initiate rpcs. type Request: Message; /// The type of response message, These messages complete rpcs, or are streamed from them. type Response: Message; + /// The type of future that completes a unary rpc. Use a `BoxFuture` if your future + /// is not `Unpin`. + type UnaryFutureType: Future + Unpin; + /// The type of stream that produces the responses for a streaming rpc. Use a + /// `BoxStream` if your stream is not `Unpin`. + type StreamType: Stream + Unpin; - /// Called with an initiating message from the client, the Reponder is how you send your response. - /// - /// If you drop the responder, the rpc is cancelled. If you send a response, the rpc is completed. - /// - /// If you want to, you can `Send` the responder into a task, e.g., via `tokio::spawn`, and - /// complete the rpc from there. The client will wait until the responder is dropped or a response - /// is sent, in theory. + /// Called with an initiating message from the client. Return the rpc completion for the + /// message: `RpcKind::Unary` for a single response, `RpcKind::Streaming` for a stream of + /// responses, or `RpcKind::Cancelled` to reject the message (the client receives a + /// cancellation for it). fn new_rpc( &mut self, initiating_message: Self::Request, - responder: RpcResponder<'_, Self::Response>, - ); + ) -> RpcKind; /// Optional poll to allow the connection to push work forward internally. /// - /// You can use this to drive connection state machines (e.g., `FuturesUnordered`), - /// or whatever else you need to do with your connection between reading from the network and - /// writing to it. + /// This is called unconditionally on every connection poll - it is NOT subject to + /// send-capacity backpressure. Do not produce outbound messages from here; use it for + /// bookkeeping, timers, and other connection state machines. fn poll( self: std::pin::Pin<&mut Self>, _context: &mut std::task::Context<'_>, @@ -79,3 +87,13 @@ pub trait ConnectionService: Unpin + 'static { std::ops::ControlFlow::Continue(()) } } + +/// Type of rpc to be completed +pub enum RpcKind { + /// A rpc that produces a single response. + Unary(Unary), + /// A rpc that produces a stream of responses. + Streaming(Streaming), + /// An immediate cancellation of an rpc. + Cancelled, +} diff --git a/protosocket-rpc/src/server/socket_server.rs b/protosocket-rpc/src/server/socket_server.rs index 7f148f1..893183f 100644 --- a/protosocket-rpc/src/server/socket_server.rs +++ b/protosocket-rpc/src/server/socket_server.rs @@ -11,7 +11,10 @@ use std::task::Poll; use crate::server::Spawn; use super::rpc_submitter::RpcSubmitter; -use super::server_traits::SocketService; +use super::server_traits::{ConnectionService, SocketService}; + +type ServiceCodec = + <::ConnectionService as ConnectionService>::Codec; /// A `SocketRpcServer` is a server future. It listens on a socket and spawns new connections, /// with a ConnectionService to handle each connection. @@ -41,17 +44,18 @@ impl super::TokioSpawn< Connection< ::Stream, - TSocketService::Codec, RpcSubmitter, >, >, > where TSocketService: SocketService, - TSocketService::Codec: Send, - ::Serialized: Send, + ServiceCodec: Send, + as Encoder>::Serialized: Send, ::Stream: Send, TSocketService::ConnectionService: Send, + ::UnaryFutureType: Send, + ::StreamType: Send, { /// Construct a new `SocketRpcServer` with a listener. /// @@ -81,7 +85,6 @@ where TSpawnConnection: Spawn< Connection< ::Stream, - TSocketService::Codec, RpcSubmitter, >, >, @@ -127,7 +130,6 @@ where TSpawnConnection: Spawn< Connection< ::Stream, - TSocketService::Codec, RpcSubmitter, >, >, @@ -140,12 +142,14 @@ where Poll::Ready(result) => match result { SocketResult::Stream(stream) => { let connection_service = self.socket_server.new_stream_service(&stream); - let (outbound_messages, outbound_messages_receiver) = spillway::channel(); - let submitter = RpcSubmitter::new(connection_service, outbound_messages); + // The rpc server produces all responses through the submitter, so + // the queue's senders are dropped immediately: the connection's + // lifetime is governed by the socket and the submitter alone. + let (_no_senders, outbound_messages_receiver) = spillway::channel(); + let submitter = RpcSubmitter::new(connection_service); #[allow(clippy::type_complexity)] let connection: Connection< ::Stream, - TSocketService::Codec, RpcSubmitter, > = Connection::new( stream, diff --git a/protosocket-server/src/connection_server.rs b/protosocket-server/src/connection_server.rs index b940a38..f58c966 100644 --- a/protosocket-server/src/connection_server.rs +++ b/protosocket-server/src/connection_server.rs @@ -1,7 +1,5 @@ use protosocket::Codec; use protosocket::Connection; -use protosocket::Decoder; -use protosocket::Encoder; use protosocket::MessageReactor; use protosocket::SocketListener; use protosocket::SocketResult; @@ -16,21 +14,19 @@ use std::task::Poll; /// The ServerConnector listens to a socket and spawns a Reactor for each new connection. pub trait ServerConnector: Unpin { /// Message encoding - type Codec: Codec - + Decoder::Inbound> - + Encoder::Outbound>; + type Codec: Codec; /// Per-connection message handler - type Reactor: MessageReactor; + type Reactor: MessageReactor; /// The listener type for this service. E.g., `TcpSocketListener` type SocketListener: SocketListener; /// Create a new message codec for a connection - fn codec(&self) -> Self::Codec; + fn codec(&mut self) -> Self::Codec; /// Create a per-connection message Reactor. /// You can look at the connection in here if you need some data, like a SocketAddr fn new_reactor( - &self, + &mut self, optional_outbound: spillway::Sender<::LogicalOutbound>, _connection: &::Stream, ) -> Self::Reactor; @@ -38,11 +34,7 @@ pub trait ServerConnector: Unpin { /// Spawn a connection - probably you just want tokio::spawn, but you might have other needs. fn spawn_connection( &self, - connection: Connection< - ::Stream, - Self::Codec, - Self::Reactor, - >, + connection: Connection<::Stream, Self::Reactor>, ); }