From 6be8a9de13e41aa8592b56feac4dc0a22fe7b539 Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Tue, 14 Jul 2026 08:19:15 -0700 Subject: [PATCH 01/10] fix: poll server rpcs only when the connection can send --- example-messagepack/src/server.rs | 22 +- example-proto-stream/src/server.rs | 52 ++-- example-proto-tls/src/server.rs | 26 +- example-proto/src/server.rs | 23 +- protosocket-connection/src/connection.rs | 41 +++- protosocket-connection/src/message_reactor.rs | 20 ++ protosocket-rpc/README.md | 21 +- .../src/server/abortion_tracker.rs | 12 +- protosocket-rpc/src/server/mod.rs | 8 +- protosocket-rpc/src/server/rpc_stream.rs | 125 ++++++++++ protosocket-rpc/src/server/rpc_submitter.rs | 230 ++++++++---------- protosocket-rpc/src/server/server_traits.rs | 72 ++++-- protosocket-rpc/src/server/socket_server.rs | 4 +- 13 files changed, 403 insertions(+), 253 deletions(-) create mode 100644 protosocket-rpc/src/server/rpc_stream.rs diff --git a/example-messagepack/src/server.rs b/example-messagepack/src/server.rs index 99f5821..551332c 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; @@ -84,30 +84,28 @@ struct DemoRpcConnectionServer { impl ConnectionService for DemoRpcConnectionServer { 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-proto-stream/src/server.rs b/example-proto-stream/src/server.rs index e2c89b3..0cd28e5 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; @@ -87,61 +81,45 @@ impl SocketService for DemoRpcSocketService { 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 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..a8180cb 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; @@ -117,31 +117,27 @@ struct DemoRpcConnectionServer { impl ConnectionService for DemoRpcConnectionServer { 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..6cf11cc 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; @@ -93,27 +93,28 @@ struct DemoRpcConnectionServer { impl ConnectionService for DemoRpcConnectionServer { 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/protosocket-connection/src/connection.rs b/protosocket-connection/src/connection.rs index d6f76b2..69a1339 100644 --- a/protosocket-connection/src/connection.rs +++ b/protosocket-connection/src/connection.rs @@ -61,7 +61,10 @@ impl< 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}}} }}" + ) } } @@ -240,7 +243,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 +259,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}"); @@ -327,8 +337,25 @@ impl< 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; + // The queue is drained; the reactor may drive its own message sources + // (e.g., streaming rpcs). It is only polled here, within the send + // budget, so a connection that cannot write does not advance them. + // SAFETY: This is a structural pin. If I'm not moved then neither is this reactor. + match unsafe { Pin::new_unchecked(&mut self.reactor) } + .poll_next_outbound(context) + { + Poll::Pending => { + log::debug!( + "no more messages to serialize, and we are pending for more" + ); + break; + } + Poll::Ready(None) => { + log::info!("reactor is finished producing messages"); + return Poll::Ready(()); + } + Poll::Ready(Some(next)) => next, + } } Poll::Ready(None) => { log::info!("outbound message channel was closed"); @@ -450,7 +477,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/message_reactor.rs b/protosocket-connection/src/message_reactor.rs index c300470..14159dd 100644 --- a/protosocket-connection/src/message_reactor.rs +++ b/protosocket-connection/src/message_reactor.rs @@ -37,6 +37,26 @@ pub trait MessageReactor: 'static { /// /// You can use this to track outbound messages, or for logging, or metrics, or whatever. fn on_outbound_message(&mut self, message: Self::LogicalOutbound) -> Self::Outbound; + + /// Poll for the next outbound message produced by the reactor itself. + /// + /// This is only called when the connection has room in its send queue. This is how + /// backpressure is applied to reactor-driven work: when the connection cannot write, + /// the reactor is not polled for outbound messages, and any streams or futures the + /// reactor drives to produce them are not advanced. If your source of messages models + /// lag or load-shedding (like `tokio::sync::broadcast`), a slow or stalled peer causes + /// that model to engage instead of buffering without bound. + /// + /// Return `Poll::Ready(None)` to tell the connection the reactor is finished producing + /// messages and the connection should close. If your reactor does not produce messages + /// on its own (for example, a client whose outbound messages are all submitted through + /// the connection's outbound queue), leave the default implementation. + fn poll_next_outbound( + self: std::pin::Pin<&mut Self>, + _context: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Pending + } } /// What the connection should do after processing a batch of inbound messages. 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/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/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_stream.rs b/protosocket-rpc/src/server/rpc_stream.rs new file mode 100644 index 0000000..637bdf2 --- /dev/null +++ b/protosocket-rpc/src/server/rpc_stream.rs @@ -0,0 +1,125 @@ +use std::{ + future::Future, + 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 returns `None`, so a `SelectAll` retires it. +#[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, tagged with terminal-ness so the caller knows what, if +/// anything, to put on the wire after it. +#[derive(Debug)] +pub enum RpcStreamEvent { + /// A streaming rpc produced a response item. + Item(T), + /// A unary rpc completed with its single response. Terminal; the response is the + /// completion, no trailing control message belongs on the wire. + Complete(T), + /// A streaming rpc's stream finished. Terminal; the peer should be told the rpc ended. + Finished, + /// The rpc was cancelled. Terminal. + Cancelled, +} + +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::new(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::new(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. + 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..5a04fb6 100644 --- a/protosocket-rpc/src/server/rpc_submitter.rs +++ b/protosocket-rpc/src/server/rpc_submitter.rs @@ -1,63 +1,101 @@ +use std::{ + collections::VecDeque, + pin::Pin, + task::{Context, Poll}, +}; + +use futures::stream::SelectAll; +use futures::Stream; use protosocket::MessageReactor; 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}; + +/// A MessageReactor that hosts a ConnectionService's rpcs and drives them within the +/// connection's send budget. +/// +/// New rpcs are registered from inbound messages. Their completions - unary and streaming +/// alike - live in one pool and are only advanced by `poll_next_outbound`, which the +/// connection calls only when it has room to send. This is the backpressure contract: a +/// connection that cannot write does not advance the work that produces responses. The +/// pool yields in readiness order; no priority between unary and streaming rpcs is imposed. +pub struct RpcSubmitter where - TConnectionServer: ConnectionService, + TConnectionService: ConnectionService, { - connection_server: TConnectionServer, - outbound: spillway::Sender::Response>>, + connection_server: TConnectionService, + /// Holds the connection's outbound queue open. Rpc responses do not flow through the + /// queue; they are produced on demand by poll_next_outbound. + _outbound: spillway::Sender, aborts: AbortionTracker, + /// Message ids of rpcs to reject. Small and self-limiting: bounded by the inbound + /// messages processed between sends. + rejections: VecDeque, + rpcs: SelectAll>, +} + +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>, + outbound: spillway::Sender, ) -> Self { Self { connection_server, - outbound, + _outbound: outbound, aborts: Default::default(), + rejections: Default::default(), + rpcs: Default::default(), } } } -pub enum RpcResponse { - Partial(T), - Final(T), - Untracked(T), -} - impl MessageReactor for RpcSubmitter where TConnectionService: ConnectionService, { 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); + } + RpcKind::Streaming(stream) => { + let (rpc, handle) = RpcStream::new_streaming(message_id, stream); + self.aborts.register(message_id, handle); + self.rpcs.push(rpc); + } + RpcKind::Cancelled => { + log::debug!("rejecting unknown rpc {message_id}"); + self.rejections.push_back(message_id); + } + }, ProtosocketControlCode::Cancel => { if let Some(abort) = self.aborts.take_abort(message_id) { log::debug!("cancelling message {message_id}"); @@ -74,114 +112,52 @@ 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 - // } + /// Produce the next response, in rpc readiness order. This is only called when the + /// connection can accept a message for serialization, so rpc work only advances when + /// its output has somewhere to go. + fn poll_next_outbound( + self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll> { + let me = self.get_mut(); + if let Some(message_id) = me.rejections.pop_front() { + return Poll::Ready(Some(::cancelled(message_id))); + } + match Pin::new(&mut me.rpcs).poll_next(context) { + Poll::Ready(Some((id, event))) => match event { + RpcStreamEvent::Item(mut item) => { + item.set_message_id(id); + Poll::Ready(Some(item)) + } + RpcStreamEvent::Complete(mut response) => { + let _ = me.aborts.take_abort(id); + response.set_message_id(id); + Poll::Ready(Some(response)) + } + RpcStreamEvent::Finished => { + let _ = me.aborts.take_abort(id); + Poll::Ready(Some(::ended(id))) + } + RpcStreamEvent::Cancelled => { + let _ = me.aborts.take_abort(id); + Poll::Ready(Some(::cancelled(id))) + } + }, + // An empty pool is not a closed reactor: new rpcs arrive from inbound + // processing, which wakes this connection on its own. + Poll::Ready(None) | Poll::Pending => Poll::Pending, + } + } } diff --git a/protosocket-rpc/src/server/server_traits.rs b/protosocket-rpc/src/server/server_traits.rs index c811a46..2cc72ff 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. /// @@ -17,9 +20,9 @@ pub trait SocketService: 'static { /// The type of connection service that will be created for each connection. type ConnectionService: ConnectionService< - Request = ::Message, - Response = ::Message, - >; + Request = ::Message, + Response = ::Message, + >; /// The listener type for this service. E.g., `TcpSocketListener` type SocketListener: SocketListener; @@ -36,42 +39,57 @@ pub trait SocketService: 'static { ) -> 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. You put your per-connection state in your /// ConnectionService implementation. /// /// 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. +/// from the client, and you return the kind of rpc that completes it: a future for a unary rpc, +/// or a stream for a streaming rpc. +/// +/// Your futures and streams are driven by the connection itself, and they are only polled when +/// the connection has room to send responses. This is how backpressure works: when the peer +/// stops receiving, your rpcs stop being polled. If your stream models lag or load-shedding +/// (like `tokio::sync::broadcast`), a slow peer causes that model to engage instead of buffering +/// responses without bound. /// -/// 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. +/// Because rpcs are polled in the context of the connection, a connection and all of its +/// outstanding rpcs advance on 1 cpu at a time. That might be good for your use case, or it +/// might be suboptimal. If an rpc is compute-heavy, you can of course spawn a task and return +/// a future that completes when the task completes, e.g., with a `tokio::sync::oneshot`. In +/// general, try to do as little as possible: Return a future and let the connection poll it. +/// This keeps your task count low and your wakes more tightly related to the cooperating +/// tasks that need to be woken. +/// +/// Response message ids are stamped by the connection: every response or stream item is sent +/// with the message id of the rpc that initiated it. pub trait ConnectionService: Unpin + '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 +97,13 @@ pub trait ConnectionService: Unpin + 'static { std::ops::ControlFlow::Continue(()) } } + +/// Type of rpc to be completed +pub enum RpcKind { + /// This is a unary rpc. It will complete with a single response. + Unary(Unary), + /// This is a streaming rpc. It will complete with a stream of responses. + Streaming(Streaming), + /// Do not process this rpc. The initiating message is answered with a cancellation. + Cancelled, +} diff --git a/protosocket-rpc/src/server/socket_server.rs b/protosocket-rpc/src/server/socket_server.rs index 7f148f1..25ca913 100644 --- a/protosocket-rpc/src/server/socket_server.rs +++ b/protosocket-rpc/src/server/socket_server.rs @@ -11,7 +11,7 @@ use std::task::Poll; use crate::server::Spawn; use super::rpc_submitter::RpcSubmitter; -use super::server_traits::SocketService; +use super::server_traits::{ConnectionService, SocketService}; /// A `SocketRpcServer` is a server future. It listens on a socket and spawns new connections, /// with a ConnectionService to handle each connection. @@ -52,6 +52,8 @@ where ::Serialized: Send, ::Stream: Send, TSocketService::ConnectionService: Send, + ::UnaryFutureType: Send, + ::StreamType: Send, { /// Construct a new `SocketRpcServer` with a listener. /// From 2ef2465b9dd883a6f48427b794c9905c969105ba Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Tue, 14 Jul 2026 08:20:12 -0700 Subject: [PATCH 02/10] chore: remove server files replaced by connection-driven rpcs --- protosocket-rpc/src/server/abortable.rs | 118 ---- .../src/server/connection_server.rs | 653 ------------------ .../src/server/forward_streaming.rs | 139 ---- protosocket-rpc/src/server/forward_unary.rs | 119 ---- protosocket-rpc/src/server/rpc_responder.rs | 66 -- 5 files changed, 1095 deletions(-) delete mode 100644 protosocket-rpc/src/server/abortable.rs delete mode 100644 protosocket-rpc/src/server/connection_server.rs delete mode 100644 protosocket-rpc/src/server/forward_streaming.rs delete mode 100644 protosocket-rpc/src/server/forward_unary.rs delete mode 100644 protosocket-rpc/src/server/rpc_responder.rs 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/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/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"); - } - } -} From fee0ecc5521bbd433b028f97ac42fee49f50d5dd Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Tue, 14 Jul 2026 09:08:55 -0700 Subject: [PATCH 03/10] refactor: close connections when both outbound sources are exhausted --- protosocket-connection/src/connection.rs | 31 +++++++++++-------- protosocket-connection/src/message_reactor.rs | 12 ++++--- protosocket-rpc/src/server/rpc_submitter.rs | 9 +----- protosocket-rpc/src/server/server_traits.rs | 6 ++-- protosocket-rpc/src/server/socket_server.rs | 7 +++-- 5 files changed, 34 insertions(+), 31 deletions(-) diff --git a/protosocket-connection/src/connection.rs b/protosocket-connection/src/connection.rs index 69a1339..ab6366e 100644 --- a/protosocket-connection/src/connection.rs +++ b/protosocket-connection/src/connection.rs @@ -335,15 +335,21 @@ impl< let start_len = self.send_buffer.len(); for _ in 0..max_outbound { - let message = match self.outbound_messages.poll_next(context) { - Poll::Pending => { - // The queue is drained; the reactor may drive its own message sources - // (e.g., streaming rpcs). It is only polled here, within the send - // budget, so a connection that cannot write does not advance them. + // The connection has two outbound sources: the queued outbound messages and + // the reactor's own message production (e.g., streaming rpcs). The reactor is + // only polled here, within the send budget, so a connection that cannot write + // does not advance the work that produces messages. The connection closes + // when both sources are exhausted: the queue's senders are all dropped and + // the reactor is finished. + let queue = self.outbound_messages.poll_next(context); + let message = match queue { + Poll::Ready(Some(next)) => next, + Poll::Pending | Poll::Ready(None) => { // SAFETY: This is a structural pin. If I'm not moved then neither is this reactor. match unsafe { Pin::new_unchecked(&mut self.reactor) } .poll_next_outbound(context) { + Poll::Ready(Some(next)) => next, Poll::Pending => { log::debug!( "no more messages to serialize, and we are pending for more" @@ -351,17 +357,16 @@ impl< break; } Poll::Ready(None) => { - log::info!("reactor is finished producing messages"); - return Poll::Ready(()); + if matches!(queue, Poll::Ready(None)) { + log::info!("all outbound message sources are exhausted"); + return Poll::Ready(()); + } + // The reactor never produces messages; the live outbound + // queue governs this connection's lifetime. + break; } - Poll::Ready(Some(next)) => next, } } - Poll::Ready(None) => { - log::info!("outbound message channel was closed"); - return Poll::Ready(()); - } - Poll::Ready(Some(next)) => next, }; let message = self.reactor.on_outbound_message(message); let buffer = self.codec.encode(message); diff --git a/protosocket-connection/src/message_reactor.rs b/protosocket-connection/src/message_reactor.rs index 14159dd..d8e9f70 100644 --- a/protosocket-connection/src/message_reactor.rs +++ b/protosocket-connection/src/message_reactor.rs @@ -47,15 +47,17 @@ pub trait MessageReactor: 'static { /// lag or load-shedding (like `tokio::sync::broadcast`), a slow or stalled peer causes /// that model to engage instead of buffering without bound. /// - /// Return `Poll::Ready(None)` to tell the connection the reactor is finished producing - /// messages and the connection should close. If your reactor does not produce messages - /// on its own (for example, a client whose outbound messages are all submitted through - /// the connection's outbound queue), leave the default implementation. + /// A connection has two outbound sources: its outbound message queue and its reactor. + /// It closes when both are exhausted - the queue's senders are all dropped and this + /// returns `Poll::Ready(None)`. The default implementation says this reactor never + /// produces messages of its own, which leaves the connection's lifetime governed by + /// the outbound queue, as before this method existed. Return `Poll::Pending` while + /// you might produce messages later. fn poll_next_outbound( self: std::pin::Pin<&mut Self>, _context: &mut std::task::Context<'_>, ) -> std::task::Poll> { - std::task::Poll::Pending + std::task::Poll::Ready(None) } } diff --git a/protosocket-rpc/src/server/rpc_submitter.rs b/protosocket-rpc/src/server/rpc_submitter.rs index 5a04fb6..b8f7438 100644 --- a/protosocket-rpc/src/server/rpc_submitter.rs +++ b/protosocket-rpc/src/server/rpc_submitter.rs @@ -28,9 +28,6 @@ where TConnectionService: ConnectionService, { connection_server: TConnectionService, - /// Holds the connection's outbound queue open. Rpc responses do not flow through the - /// queue; they are produced on demand by poll_next_outbound. - _outbound: spillway::Sender, aborts: AbortionTracker, /// Message ids of rpcs to reject. Small and self-limiting: bounded by the inbound /// messages processed between sends. @@ -55,13 +52,9 @@ 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: outbound, aborts: Default::default(), rejections: Default::default(), rpcs: Default::default(), diff --git a/protosocket-rpc/src/server/server_traits.rs b/protosocket-rpc/src/server/server_traits.rs index 2cc72ff..45be820 100644 --- a/protosocket-rpc/src/server/server_traits.rs +++ b/protosocket-rpc/src/server/server_traits.rs @@ -20,9 +20,9 @@ pub trait SocketService: 'static { /// The type of connection service that will be created for each connection. type ConnectionService: ConnectionService< - Request = ::Message, - Response = ::Message, - >; + Request = ::Message, + Response = ::Message, + >; /// The listener type for this service. E.g., `TcpSocketListener` type SocketListener: SocketListener; diff --git a/protosocket-rpc/src/server/socket_server.rs b/protosocket-rpc/src/server/socket_server.rs index 25ca913..cb5ca13 100644 --- a/protosocket-rpc/src/server/socket_server.rs +++ b/protosocket-rpc/src/server/socket_server.rs @@ -142,8 +142,11 @@ 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, From b930234b4dee3a83ccded087de50e3a0cd577b81 Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Tue, 14 Jul 2026 09:40:53 -0700 Subject: [PATCH 04/10] perf: drain hot rpcs in batches to amortize pool re-insertion --- protosocket-rpc/src/server/rpc_submitter.rs | 128 +++++++++++++++----- 1 file changed, 98 insertions(+), 30 deletions(-) diff --git a/protosocket-rpc/src/server/rpc_submitter.rs b/protosocket-rpc/src/server/rpc_submitter.rs index b8f7438..7dc3ba6 100644 --- a/protosocket-rpc/src/server/rpc_submitter.rs +++ b/protosocket-rpc/src/server/rpc_submitter.rs @@ -4,17 +4,28 @@ use std::{ task::{Context, Poll}, }; -use futures::stream::SelectAll; -use futures::Stream; +use futures::stream::{FuturesUnordered, StreamFuture}; +use futures::{Stream, StreamExt}; use protosocket::MessageReactor; use crate::{ - server::{abortion_tracker::AbortionTracker, ConnectionService, RpcKind}, Message, ProtosocketControlCode, + server::{ConnectionService, RpcKind, abortion_tracker::AbortionTracker}, }; use super::rpc_stream::{RpcStream, RpcStreamEvent}; +/// How many messages to take from one rpc before returning it to the pool. Re-inserting +/// a stream into the pool costs a node cycle per visit, so hot streams are drained in +/// small batches to amortize it. This bounds both the staging queue and how far one rpc +/// can burst ahead of its peers within a send budget. +const PER_RPC_BATCH: usize = 64; + +type PooledRpc = RpcStream< + ::UnaryFutureType, + ::StreamType, +>; + /// A MessageReactor that hosts a ConnectionService's rpcs and drives them within the /// connection's send budget. /// @@ -22,7 +33,8 @@ use super::rpc_stream::{RpcStream, RpcStreamEvent}; /// alike - live in one pool and are only advanced by `poll_next_outbound`, which the /// connection calls only when it has room to send. This is the backpressure contract: a /// connection that cannot write does not advance the work that produces responses. The -/// pool yields in readiness order; no priority between unary and streaming rpcs is imposed. +/// pool yields in readiness order; no priority between unary and streaming rpcs is +/// imposed, and a ready rpc yields at most `PER_RPC_BATCH` messages per pool visit. pub struct RpcSubmitter where TConnectionService: ConnectionService, @@ -32,7 +44,10 @@ where /// Message ids of rpcs to reject. Small and self-limiting: bounded by the inbound /// messages processed between sends. rejections: VecDeque, - rpcs: SelectAll>, + rpcs: FuturesUnordered>>, + /// Messages already claimed from a batched rpc, awaiting hand-off to the connection. + /// Bounded by `PER_RPC_BATCH`. + staged: VecDeque, } impl std::fmt::Debug for RpcSubmitter @@ -44,6 +59,7 @@ where .field("aborts", &self.aborts) .field("rejections", &self.rejections.len()) .field("rpcs", &self.rpcs.len()) + .field("staged", &self.staged.len()) .finish() } } @@ -58,6 +74,35 @@ where aborts: Default::default(), rejections: Default::default(), rpcs: Default::default(), + staged: Default::default(), + } + } + + /// Turn an rpc event into the message that goes on the wire, retiring the rpc's + /// cancellation bookkeeping on terminal events. + 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) + } } } } @@ -77,15 +122,15 @@ where RpcKind::Unary(completion) => { let (rpc, handle) = RpcStream::new_unary(message_id, completion); self.aborts.register(message_id, handle); - self.rpcs.push(rpc); + 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); + self.rpcs.push(rpc.into_future()); } RpcKind::Cancelled => { - log::debug!("rejecting unknown rpc {message_id}"); + log::debug!("rejecting rpc {message_id}"); self.rejections.push_back(message_id); } }, @@ -125,32 +170,55 @@ where context: &mut Context<'_>, ) -> Poll> { let me = self.get_mut(); + if let Some(staged) = me.staged.pop_front() { + return Poll::Ready(Some(staged)); + } if let Some(message_id) = me.rejections.pop_front() { return Poll::Ready(Some(::cancelled(message_id))); } - match Pin::new(&mut me.rpcs).poll_next(context) { - Poll::Ready(Some((id, event))) => match event { - RpcStreamEvent::Item(mut item) => { - item.set_message_id(id); - Poll::Ready(Some(item)) - } - RpcStreamEvent::Complete(mut response) => { - let _ = me.aborts.take_abort(id); - response.set_message_id(id); - Poll::Ready(Some(response)) - } - RpcStreamEvent::Finished => { - let _ = me.aborts.take_abort(id); - Poll::Ready(Some(::ended(id))) + loop { + match Pin::new(&mut me.rpcs).poll_next(context) { + Poll::Ready(Some((first, mut rpc))) => { + let Some((id, event)) = first else { + // The rpc's terminal event was already delivered; retire it and + // look at the next ready rpc. + continue; + }; + let terminal = !matches!(event, RpcStreamEvent::Item(_)); + let first = me.message_for_event(id, event); + if terminal { + return Poll::Ready(Some(first)); + } + // Drain a small batch from this rpc while it is hot, so a busy stream + // doesn't pay a pool re-insertion per message. + let mut exhausted = false; + while me.staged.len() < PER_RPC_BATCH - 1 { + match Pin::new(&mut rpc).poll_next(context) { + Poll::Ready(Some((id, event))) => { + let terminal = !matches!(event, RpcStreamEvent::Item(_)); + let message = me.message_for_event(id, event); + me.staged.push_back(message); + if terminal { + exhausted = true; + break; + } + } + Poll::Ready(None) => { + exhausted = true; + break; + } + Poll::Pending => break, + } + } + if !exhausted { + me.rpcs.push(rpc.into_future()); + } + return Poll::Ready(Some(first)); } - RpcStreamEvent::Cancelled => { - let _ = me.aborts.take_abort(id); - Poll::Ready(Some(::cancelled(id))) - } - }, - // An empty pool is not a closed reactor: new rpcs arrive from inbound - // processing, which wakes this connection on its own. - Poll::Ready(None) | Poll::Pending => Poll::Pending, + // An empty pool is not a closed reactor: new rpcs arrive from inbound + // processing, which wakes this connection on its own. + Poll::Ready(None) | Poll::Pending => return Poll::Pending, + } } } } From efedeaa2731dd57e256f36c5685632ea68a7aa5b Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Tue, 14 Jul 2026 10:12:16 -0700 Subject: [PATCH 05/10] perf: produce reactor outbound messages in send-budget batches --- protosocket-connection/src/connection.rs | 95 +++++++++++-------- protosocket-connection/src/message_reactor.rs | 31 +++--- protosocket-rpc/src/server/rpc_submitter.rs | 76 ++++++++------- 3 files changed, 112 insertions(+), 90 deletions(-) diff --git a/protosocket-connection/src/connection.rs b/protosocket-connection/src/connection.rs index ab6366e..7fb7da0 100644 --- a/protosocket-connection/src/connection.rs +++ b/protosocket-connection/src/connection.rs @@ -37,6 +37,9 @@ pub struct Connection< > { stream: TStream, outbound_messages: spillway::Receiver, + /// Reusable landing area for reactor-produced messages, drained through + /// `on_outbound_message` and the codec each poll. Bounded by the send budget. + reactor_outbound_scratch: VecDeque, send_buffer: VecDeque<::Serialized>, receive_buffer_unread_index: usize, receive_buffer: Vec, @@ -183,6 +186,7 @@ impl< Self { stream, outbound_messages, + reactor_outbound_scratch: Default::default(), send_buffer: Default::default(), receive_buffer: Vec::new(), max_buffer_length, @@ -333,49 +337,60 @@ impl< return Poll::Pending; } + // The connection has two outbound sources: the queued outbound messages and the + // reactor's own message production (e.g., streaming rpcs). Both are only polled + // here, within the send budget, so a connection that cannot write does not + // advance the work that produces messages. The connection closes when both + // sources are exhausted: the queue's senders are all dropped and the reactor is + // finished. let start_len = self.send_buffer.len(); - for _ in 0..max_outbound { - // The connection has two outbound sources: the queued outbound messages and - // the reactor's own message production (e.g., streaming rpcs). The reactor is - // only polled here, within the send budget, so a connection that cannot write - // does not advance the work that produces messages. The connection closes - // when both sources are exhausted: the queue's senders are all dropped and - // the reactor is finished. - let queue = self.outbound_messages.poll_next(context); - let message = match queue { - Poll::Ready(Some(next)) => next, - Poll::Pending | Poll::Ready(None) => { - // SAFETY: This is a structural pin. If I'm not moved then neither is this reactor. - match unsafe { Pin::new_unchecked(&mut self.reactor) } - .poll_next_outbound(context) - { - Poll::Ready(Some(next)) => next, - Poll::Pending => { - log::debug!( - "no more messages to serialize, and we are pending for more" - ); - break; - } - Poll::Ready(None) => { - if matches!(queue, Poll::Ready(None)) { - log::info!("all outbound message sources are exhausted"); - return Poll::Ready(()); - } - // The reactor never produces messages; the live outbound - // queue governs this connection's lifetime. - 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) => { + queue_closed = true; + break; } + Poll::Pending => break, + } + } + + let remaining = max_outbound - (self.send_buffer.len() - start_len); + if 0 < remaining { + let scratch = &mut self.reactor_outbound_scratch; + debug_assert!(scratch.is_empty(), "scratch is always drained after use"); + let mut sink = |message: TReactor::LogicalOutbound| { + scratch.push_back(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); + // SAFETY: This is a structural pin. If I'm not moved then neither is this reactor. + let reactor_state = unsafe { Pin::new_unchecked(&mut self.reactor) } + .poll_outbound_many(context, &mut sink, remaining); + // Reactor-produced messages get the same treatment as queued messages: + // through on_outbound_message, then the codec. + while let Some(message) = self.reactor_outbound_scratch.pop_front() { + let message = self.reactor.on_outbound_message(message); + let buffer = self.codec.encode(message); + log::trace!("serialized reactor message: {}b", buffer.remaining()); + self.send_buffer.push_back(buffer); + } + if let Poll::Ready(None) = reactor_state { + if queue_closed { + log::info!("all outbound message sources are exhausted"); + return Poll::Ready(()); + } + // The reactor never produces messages; the live outbound queue + // governs this connection's lifetime. + } } let new_len = self.send_buffer.len(); if start_len != new_len { diff --git a/protosocket-connection/src/message_reactor.rs b/protosocket-connection/src/message_reactor.rs index d8e9f70..981fe01 100644 --- a/protosocket-connection/src/message_reactor.rs +++ b/protosocket-connection/src/message_reactor.rs @@ -38,25 +38,34 @@ pub trait MessageReactor: 'static { /// You can use this to track outbound messages, or for logging, or metrics, or whatever. fn on_outbound_message(&mut self, message: Self::LogicalOutbound) -> Self::Outbound; - /// Poll for the next outbound message produced by the reactor itself. + /// Produce outbound messages from the reactor itself, up to `budget` of them, + /// delivered through `sink`. /// - /// This is only called when the connection has room in its send queue. This is how - /// backpressure is applied to reactor-driven work: when the connection cannot write, - /// the reactor is not polled for outbound messages, and any streams or futures the - /// reactor drives to produce them are not advanced. If your source of messages models - /// lag or load-shedding (like `tokio::sync::broadcast`), a slow or stalled peer causes - /// that model to engage instead of buffering without bound. + /// This is only called when the connection has room in its send queue, and `budget` + /// is exactly that room. This is how backpressure is applied to reactor-driven work: + /// when the connection cannot write, the reactor is not polled for outbound messages, + /// and any streams or futures the reactor drives to produce them are not advanced. + /// If your source of messages models lag or load-shedding (like + /// `tokio::sync::broadcast`), a slow or stalled peer causes that model to engage + /// instead of buffering without bound. + /// + /// Messages delivered through `sink` pass through `on_outbound_message` before + /// serialization, exactly like messages from the outbound queue: every outbound + /// message is observed by that hook, whichever source produced it. /// /// A connection has two outbound sources: its outbound message queue and its reactor. /// It closes when both are exhausted - the queue's senders are all dropped and this /// returns `Poll::Ready(None)`. The default implementation says this reactor never /// produces messages of its own, which leaves the connection's lifetime governed by - /// the outbound queue, as before this method existed. Return `Poll::Pending` while - /// you might produce messages later. - fn poll_next_outbound( + /// the outbound queue, as before this method existed. Return `Poll::Pending` when + /// nothing is available right now (registering wakers), and `Poll::Ready(Some(()))` + /// after emitting one or more messages. + fn poll_outbound_many( self: std::pin::Pin<&mut Self>, _context: &mut std::task::Context<'_>, - ) -> std::task::Poll> { + _sink: &mut impl FnMut(Self::LogicalOutbound), + _budget: usize, + ) -> std::task::Poll> { std::task::Poll::Ready(None) } } diff --git a/protosocket-rpc/src/server/rpc_submitter.rs b/protosocket-rpc/src/server/rpc_submitter.rs index 7dc3ba6..9d93811 100644 --- a/protosocket-rpc/src/server/rpc_submitter.rs +++ b/protosocket-rpc/src/server/rpc_submitter.rs @@ -15,12 +15,6 @@ use crate::{ use super::rpc_stream::{RpcStream, RpcStreamEvent}; -/// How many messages to take from one rpc before returning it to the pool. Re-inserting -/// a stream into the pool costs a node cycle per visit, so hot streams are drained in -/// small batches to amortize it. This bounds both the staging queue and how far one rpc -/// can burst ahead of its peers within a send budget. -const PER_RPC_BATCH: usize = 64; - type PooledRpc = RpcStream< ::UnaryFutureType, ::StreamType, @@ -30,11 +24,11 @@ type PooledRpc = RpcStream< /// connection's send budget. /// /// New rpcs are registered from inbound messages. Their completions - unary and streaming -/// alike - live in one pool and are only advanced by `poll_next_outbound`, which the +/// alike - live in one pool and are only advanced by `poll_outbound_many`, which the /// connection calls only when it has room to send. This is the backpressure contract: a /// connection that cannot write does not advance the work that produces responses. The /// pool yields in readiness order; no priority between unary and streaming rpcs is -/// imposed, and a ready rpc yields at most `PER_RPC_BATCH` messages per pool visit. +/// imposed, and a ready rpc is drained while it is hot, up to the send budget. pub struct RpcSubmitter where TConnectionService: ConnectionService, @@ -45,9 +39,6 @@ where /// messages processed between sends. rejections: VecDeque, rpcs: FuturesUnordered>>, - /// Messages already claimed from a batched rpc, awaiting hand-off to the connection. - /// Bounded by `PER_RPC_BATCH`. - staged: VecDeque, } impl std::fmt::Debug for RpcSubmitter @@ -59,7 +50,6 @@ where .field("aborts", &self.aborts) .field("rejections", &self.rejections.len()) .field("rpcs", &self.rpcs.len()) - .field("staged", &self.staged.len()) .finish() } } @@ -74,7 +64,6 @@ where aborts: Default::default(), rejections: Default::default(), rpcs: Default::default(), - staged: Default::default(), } } @@ -162,21 +151,24 @@ where structurally_pinned_connection_server.poll(context) } - /// Produce the next response, in rpc readiness order. This is only called when the - /// connection can accept a message for serialization, so rpc work only advances when - /// its output has somewhere to go. - fn poll_next_outbound( + /// Produce responses in rpc readiness order, draining each ready rpc while it is + /// hot to amortize pool re-insertion. This is only called when the connection can + /// accept messages for serialization, and `budget` is exactly the room it has, so + /// rpc work only advances when its output has somewhere to go. + fn poll_outbound_many( self: Pin<&mut Self>, context: &mut Context<'_>, - ) -> Poll> { + sink: &mut impl FnMut(Self::LogicalOutbound), + budget: usize, + ) -> Poll> { let me = self.get_mut(); - if let Some(staged) = me.staged.pop_front() { - return Poll::Ready(Some(staged)); - } - if let Some(message_id) = me.rejections.pop_front() { - return Poll::Ready(Some(::cancelled(message_id))); - } - loop { + let mut produced = 0; + while produced < budget { + if let Some(message_id) = me.rejections.pop_front() { + sink(::cancelled(message_id)); + produced += 1; + continue; + } match Pin::new(&mut me.rpcs).poll_next(context) { Poll::Ready(Some((first, mut rpc))) => { let Some((id, event)) = first else { @@ -185,40 +177,46 @@ where continue; }; let terminal = !matches!(event, RpcStreamEvent::Item(_)); - let first = me.message_for_event(id, event); + sink(me.message_for_event(id, event)); + produced += 1; if terminal { - return Poll::Ready(Some(first)); + continue; } - // Drain a small batch from this rpc while it is hot, so a busy stream - // doesn't pay a pool re-insertion per message. - let mut exhausted = false; - while me.staged.len() < PER_RPC_BATCH - 1 { + // Drain this rpc while it is hot, so a busy stream doesn't pay a + // pool re-insertion per message. Fairness across rpcs is round-robin + // per pool visit: an rpc that exhausts the budget goes to the back. + let mut retired = false; + while produced < budget { match Pin::new(&mut rpc).poll_next(context) { Poll::Ready(Some((id, event))) => { let terminal = !matches!(event, RpcStreamEvent::Item(_)); - let message = me.message_for_event(id, event); - me.staged.push_back(message); + sink(me.message_for_event(id, event)); + produced += 1; if terminal { - exhausted = true; + retired = true; break; } } Poll::Ready(None) => { - exhausted = true; + retired = true; break; } Poll::Pending => break, } } - if !exhausted { + if !retired { me.rpcs.push(rpc.into_future()); } - return Poll::Ready(Some(first)); } - // An empty pool is not a closed reactor: new rpcs arrive from inbound + // An empty pool is not a finished reactor: new rpcs arrive from inbound // processing, which wakes this connection on its own. - Poll::Ready(None) | Poll::Pending => return Poll::Pending, + Poll::Ready(None) | Poll::Pending => break, } } + if 0 < produced { + Poll::Ready(Some(())) + } else { + Poll::Pending + } } } From efea5ed5edd7f133921339023588db678be7821e Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Tue, 14 Jul 2026 15:51:57 -0700 Subject: [PATCH 06/10] docs: tighten server docs and add a spawning example --- Cargo.toml | 1 + example-proto-spawn/Cargo.toml | 26 +++ example-proto-spawn/src/client.rs | 246 ++++++++++++++++++++ example-proto-spawn/src/messages.rs | 130 +++++++++++ example-proto-spawn/src/server.rs | 172 ++++++++++++++ protosocket-rpc/src/server/rpc_stream.rs | 11 +- protosocket-rpc/src/server/server_traits.rs | 31 +-- 7 files changed, 593 insertions(+), 24 deletions(-) create mode 100644 example-proto-spawn/Cargo.toml create mode 100644 example-proto-spawn/src/client.rs create mode 100644 example-proto-spawn/src/messages.rs create mode 100644 example-proto-spawn/src/server.rs diff --git a/Cargo.toml b/Cargo.toml index 45d6398..d2741c3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ members = [ "example-messagepack", "example-pool", "example-proto", + "example-proto-spawn", "example-proto-stream", "example-proto-tls", "example-telnet", 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..8539d74 --- /dev/null +++ b/example-proto-spawn/src/client.rs @@ -0,0 +1,246 @@ +use std::{ + sync::{ + atomic::{AtomicU64, AtomicUsize, Ordering}, + Arc, + }, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use futures::StreamExt; +use messages::{EchoRequest, EchoResponseKind, Request, Response, ResponseBehavior}; +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!( + // _ = 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, + 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); + // Alternate unary and streaming so both spawned server paths get exercised. + let behavior = if i.is_multiple_of(2) { + ResponseBehavior::Unary + } else { + ResponseBehavior::Stream + }; + 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, + }), + response_behavior: behavior as i32, + }; + concurrent_count.fetch_add(1, Ordering::Relaxed); + match behavior { + ResponseBehavior::Unary => match client.send_unary(request) { + Ok(completion) => { + let response = completion.await.expect("response must be successful"); + handle_unary_response(response, &metrics_count, &metrics_latency); + } + Err(e) => { + log::error!("send should work: {e:?}"); + return; + } + }, + ResponseBehavior::Stream => 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_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.kind { + Some(EchoResponseKind::Echo(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); + } + Some(EchoResponseKind::Stream(_char_response)) => { + log::error!("got a stream response for a unary request"); + } + None => { + log::warn!("no response body"); + } + } +} + +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.kind { + Some(EchoResponseKind::Echo(_echo)) => { + log::error!("got a unary response for a stream request"); + } + Some(EchoResponseKind::Stream(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/src/messages.rs b/example-proto-spawn/src/messages.rs new file mode 100644 index 0000000..eecf5de --- /dev/null +++ b/example-proto-spawn/src/messages.rs @@ -0,0 +1,130 @@ +//! 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, + #[prost(enumeration = "ResponseBehavior", tag = "4")] + pub response_behavior: i32, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, prost::Enumeration)] +#[repr(i32)] +pub enum ResponseBehavior { + Unary = 0, + Stream = 1, +} + +#[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(oneof = "EchoResponseKind", tags = "3, 4")] + pub kind: Option, +} + +#[derive(Clone, PartialEq, Eq, prost::Oneof)] +pub enum EchoResponseKind { + #[prost(message, tag = "3")] + Echo(EchoResponse), + #[prost(message, tag = "4")] + Stream(EchoStream), +} + +#[derive(Clone, PartialEq, Eq, prost::Message)] +pub struct EchoResponse { + #[prost(string, tag = "1")] + pub message: String, + #[prost(uint64, tag = "2")] + pub nanotime: u64, +} + +#[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, + response_behavior: ResponseBehavior::Unary as i32, + } + } + + 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, + response_behavior: ResponseBehavior::Unary as i32, + } + } +} + +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, + kind: 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, + kind: None, + } + } +} diff --git a/example-proto-spawn/src/server.rs b/example-proto-spawn/src/server.rs new file mode 100644 index 0000000..9e41a14 --- /dev/null +++ b/example-proto-spawn/src/server.rs @@ -0,0 +1,172 @@ +//! A protosocket-rpc server that spawns its 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. Use a completion future for unary work, and a producer task behind +//! a bounded channel for streams - the channel capacity is how far the producer can run +//! ahead of the peer. + +use std::pin::pin; +use std::sync::atomic::AtomicUsize; +use std::time::Duration; + +use futures::{future::BoxFuture, stream::BoxStream, FutureExt, 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, RpcKind, SocketService}, + Message, 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 Codec = ( + PooledEncoder>, + ProstDecoder, + ); + type ConnectionService = DemoRpcConnectionServer; + type SocketListener = TcpSocketListener; + + fn codec(&self) -> Self::Codec { + Default::default() + } + + fn new_stream_service(&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 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; + let behavior = initiating_message.response_behavior(); + match initiating_message.body { + Some(echo) => match behavior { + ResponseBehavior::Unary => { + // Spawn the work and return a future that completes with the task. + let work = tokio::spawn(echo_request(request_id, echo)); + RpcKind::Unary( + async move { + match work.await { + Ok(response) => response, + Err(join_error) => { + log::error!("rpc task failed: {join_error}"); + Response::cancelled(request_id) + } + } + } + .boxed(), + ) + } + ResponseBehavior::Stream => { + // Spawn the producer behind a bounded channel. The producer waits + // when 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 + } + } + } +} + +async fn echo_request(request_id: u64, echo: EchoRequest) -> Response { + // Pretend this is compute-heavy work that deserves its own task. + tokio::time::sleep(Duration::from_micros(1)).await; + Response { + request_id, + code: ProtosocketControlCode::Normal as u32, + kind: Some(messages::EchoResponseKind::Echo(EchoResponse { + message: echo.message, + nanotime: echo.nanotime, + })), + } +} + +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, + kind: Some(messages::EchoResponseKind::Stream(EchoStream { + message: (c as char).to_string(), + nanotime, + sequence: sequence as u64, + })), + }, + )) +} diff --git a/protosocket-rpc/src/server/rpc_stream.rs b/protosocket-rpc/src/server/rpc_stream.rs index 637bdf2..95bfa82 100644 --- a/protosocket-rpc/src/server/rpc_stream.rs +++ b/protosocket-rpc/src/server/rpc_stream.rs @@ -1,6 +1,6 @@ use std::{ future::Future, - pin::Pin, + pin::{pin, Pin}, sync::{ atomic::{AtomicBool, Ordering}, Arc, @@ -15,7 +15,7 @@ use futures::{task::AtomicWaker, Stream}; /// 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 returns `None`, so a `SelectAll` retires it. +/// stream yields `None`. #[derive(Debug)] pub struct RpcStream { id: u64, @@ -89,14 +89,14 @@ where }; } match &mut self.kind { - RpcStreamKind::Unary(completion) => match Pin::new(completion).poll(context) { + 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::new(stream).poll_next(context) { + 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; @@ -118,6 +118,9 @@ pub struct RpcAbortHandle { 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/server_traits.rs b/protosocket-rpc/src/server/server_traits.rs index 45be820..1ab1a3e 100644 --- a/protosocket-rpc/src/server/server_traits.rs +++ b/protosocket-rpc/src/server/server_traits.rs @@ -39,31 +39,22 @@ pub trait SocketService: 'static { ) -> Self::ConnectionService; } -/// A connection service receives rpcs from clients and returns the work that produces responses. +/// A connection service receives rpcs from clients and returns the work that produces +/// responses. /// -/// Each client connection gets a ConnectionService. You put your per-connection state in your -/// ConnectionService implementation. +/// 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. /// -/// 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 rpc that completes it: a future for a unary rpc, -/// or a stream for a streaming rpc. -/// -/// Your futures and streams are driven by the connection itself, and they are only polled when -/// the connection has room to send responses. This is how backpressure works: when the peer -/// stops receiving, your rpcs stop being polled. If your stream models lag or load-shedding -/// (like `tokio::sync::broadcast`), a slow peer causes that model to engage instead of buffering +/// 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. /// -/// Because rpcs are polled in the context of the connection, a connection and all of its -/// outstanding rpcs advance on 1 cpu at a time. That might be good for your use case, or it -/// might be suboptimal. If an rpc is compute-heavy, you can of course spawn a task and return -/// a future that completes when the task completes, e.g., with a `tokio::sync::oneshot`. In -/// general, try to do as little as possible: Return a future and let the connection poll it. -/// This keeps your task count low and your wakes more tightly related to the cooperating -/// tasks that need to be woken. +/// 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. /// -/// Response message ids are stamped by the connection: every response or stream item is sent -/// with the message id of the rpc that initiated it. +/// Response message ids are stamped by the connection. pub trait ConnectionService: Unpin + 'static { /// The type of request message, These messages initiate rpcs. type Request: Message; From a4e75c5e3893869072a0d1b06e209db514b1e71a Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Tue, 14 Jul 2026 15:51:57 -0700 Subject: [PATCH 07/10] feat: reserve send permits to enforce the outbound send budget --- Cargo.lock | 15 +++ protosocket-connection/src/connection.rs | 123 +++++++++++++----- protosocket-connection/src/lib.rs | 3 + protosocket-connection/src/message_reactor.rs | 65 +++++---- protosocket-connection/src/socket_listener.rs | 2 +- protosocket-rpc/src/server/rpc_submitter.rs | 86 ++++++------ 6 files changed, 193 insertions(+), 101 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index eb81a62..bdfeeec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -493,6 +493,21 @@ 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-stream" version = "0.1.0" diff --git a/protosocket-connection/src/connection.rs b/protosocket-connection/src/connection.rs index 7fb7da0..e760ff8 100644 --- a/protosocket-connection/src/connection.rs +++ b/protosocket-connection/src/connection.rs @@ -12,10 +12,80 @@ use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use crate::{ encoding::Codec, interrupted, - message_reactor::{MessageReactor, ReactorStatus}, + message_reactor::{MessageReactor, ReactorStatus, SendBudget, SendManyPermit, SendPermit}, would_block, Decoder, DeserializeError, Encoder, }; +/// The connection's send capacity, spent by encoding directly into the send buffer. +struct EncodingSendBudget<'a, TCodec: Codec> { + codec: &'a mut TCodec, + send_buffer: &'a mut VecDeque<::Serialized>, + remaining: usize, +} + +impl SendBudget<::Message> for EncodingSendBudget<'_, TCodec> { + fn reserve(&mut self) -> Option::Message> + '_> { + if self.remaining == 0 { + None + } else { + Some(EncodingSendPermit { budget: self }) + } + } + + fn reserve_many( + &mut self, + count: usize, + ) -> impl SendManyPermit<::Message> + '_ { + let reserved = count.min(self.remaining); + EncodingSendManyPermit { + budget: self, + reserved, + } + } +} + +/// A reserved slot in the send buffer. +struct EncodingSendPermit<'a, 'b, TCodec: Codec> { + budget: &'a mut EncodingSendBudget<'b, TCodec>, +} + +impl SendPermit<::Message> + for EncodingSendPermit<'_, '_, TCodec> +{ + fn send(self, message: ::Message) { + self.budget.remaining -= 1; + let buffer = self.budget.codec.encode(message); + log::trace!("serialized reactor message: {}b", buffer.remaining()); + self.budget.send_buffer.push_back(buffer); + } +} + +/// A reserved run of slots in the send buffer. +struct EncodingSendManyPermit<'a, 'b, TCodec: Codec> { + budget: &'a mut EncodingSendBudget<'b, TCodec>, + reserved: usize, +} + +impl SendManyPermit<::Message> + for EncodingSendManyPermit<'_, '_, TCodec> +{ + fn reserved(&self) -> usize { + self.reserved + } + + fn send(&mut self, message: ::Message) { + assert!( + 0 < self.reserved, + "sent more messages than the permit reserved" + ); + self.reserved -= 1; + self.budget.remaining -= 1; + let buffer = self.budget.codec.encode(message); + log::trace!("serialized reactor message: {}b", buffer.remaining()); + self.budget.send_buffer.push_back(buffer); + } +} + /// A bidirectional, message-oriented AsyncRead/AsyncWrite stream wrapper. /// /// Connections are Futures that you spawn. @@ -37,9 +107,6 @@ pub struct Connection< > { stream: TStream, outbound_messages: spillway::Receiver, - /// Reusable landing area for reactor-produced messages, drained through - /// `on_outbound_message` and the codec each poll. Bounded by the send budget. - reactor_outbound_scratch: VecDeque, send_buffer: VecDeque<::Serialized>, receive_buffer_unread_index: usize, receive_buffer: Vec, @@ -186,7 +253,6 @@ impl< Self { stream, outbound_messages, - reactor_outbound_scratch: Default::default(), send_buffer: Default::default(), receive_buffer: Vec::new(), max_buffer_length, @@ -327,7 +393,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(); @@ -337,12 +406,6 @@ impl< return Poll::Pending; } - // The connection has two outbound sources: the queued outbound messages and the - // reactor's own message production (e.g., streaming rpcs). Both are only polled - // here, within the send budget, so a connection that cannot write does not - // advance the work that produces messages. The connection closes when both - // sources are exhausted: the queue's senders are all dropped and the reactor is - // finished. let start_len = self.send_buffer.len(); let mut queue_closed = false; while self.send_buffer.len() - start_len < max_outbound { @@ -367,29 +430,23 @@ impl< let remaining = max_outbound - (self.send_buffer.len() - start_len); if 0 < remaining { - let scratch = &mut self.reactor_outbound_scratch; - debug_assert!(scratch.is_empty(), "scratch is always drained after use"); - let mut sink = |message: TReactor::LogicalOutbound| { - scratch.push_back(message); + let Self { + send_buffer, + reactor, + codec, + .. + } = self; + let mut outbound = EncodingSendBudget { + 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(&mut self.reactor) } - .poll_outbound_many(context, &mut sink, remaining); - // Reactor-produced messages get the same treatment as queued messages: - // through on_outbound_message, then the codec. - while let Some(message) = self.reactor_outbound_scratch.pop_front() { - let message = self.reactor.on_outbound_message(message); - let buffer = self.codec.encode(message); - log::trace!("serialized reactor message: {}b", buffer.remaining()); - self.send_buffer.push_back(buffer); - } - if let Poll::Ready(None) = reactor_state { - if queue_closed { - log::info!("all outbound message sources are exhausted"); - return Poll::Ready(()); - } - // The reactor never produces messages; the live outbound queue - // governs this connection's lifetime. + 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(); diff --git a/protosocket-connection/src/lib.rs b/protosocket-connection/src/lib.rs index a71dc9e..e359706 100644 --- a/protosocket-connection/src/lib.rs +++ b/protosocket-connection/src/lib.rs @@ -23,6 +23,9 @@ pub use encoding::OwnedBuffer; pub use error::DeserializeError; pub use message_reactor::MessageReactor; pub use message_reactor::ReactorStatus; +pub use message_reactor::SendBudget; +pub use message_reactor::SendManyPermit; +pub use message_reactor::SendPermit; pub use pooled_encoder::PooledEncoder; pub use pooled_encoder::Reusable; pub use pooled_encoder::Serialize; diff --git a/protosocket-connection/src/message_reactor.rs b/protosocket-connection/src/message_reactor.rs index 981fe01..c101b5e 100644 --- a/protosocket-connection/src/message_reactor.rs +++ b/protosocket-connection/src/message_reactor.rs @@ -33,43 +33,64 @@ 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. Messages produced by `poll_outbound_many` do not come through here. /// /// You can use this to track outbound messages, or for logging, or metrics, or whatever. fn on_outbound_message(&mut self, message: Self::LogicalOutbound) -> Self::Outbound; - /// Produce outbound messages from the reactor itself, up to `budget` of them, - /// delivered through `sink`. + /// Produce outbound wire messages into `outbound`. /// - /// This is only called when the connection has room in its send queue, and `budget` - /// is exactly that room. This is how backpressure is applied to reactor-driven work: - /// when the connection cannot write, the reactor is not polled for outbound messages, - /// and any streams or futures the reactor drives to produce them are not advanced. - /// If your source of messages models lag or load-shedding (like - /// `tokio::sync::broadcast`), a slow or stalled peer causes that model to engage - /// instead of buffering without bound. + /// 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. If your message source models lag (like + /// `tokio::sync::broadcast`), a slow peer engages that model instead of buffering + /// without bound. /// - /// Messages delivered through `sink` pass through `on_outbound_message` before - /// serialization, exactly like messages from the outbound queue: every outbound - /// message is observed by that hook, whichever source produced it. + /// Messages you emit here are yours: do any bookkeeping before you send them. /// - /// A connection has two outbound sources: its outbound message queue and its reactor. - /// It closes when both are exhausted - the queue's senders are all dropped and this - /// returns `Poll::Ready(None)`. The default implementation says this reactor never - /// produces messages of its own, which leaves the connection's lifetime governed by - /// the outbound queue, as before this method existed. Return `Poll::Pending` when - /// nothing is available right now (registering wakers), and `Poll::Ready(Some(()))` - /// after emitting one or more messages. + /// Return `Ready(Some(()))` after emitting, `Pending` when nothing is available + /// right now, and `Ready(None)` if you will never emit - the default. The connection + /// closes when its outbound queue is closed and this returns `Ready(None)`. fn poll_outbound_many( self: std::pin::Pin<&mut Self>, _context: &mut std::task::Context<'_>, - _sink: &mut impl FnMut(Self::LogicalOutbound), - _budget: usize, + _outbound: &mut impl SendBudget, ) -> std::task::Poll> { std::task::Poll::Ready(None) } } +/// A bounded lease on a connection's send capacity. +pub trait SendBudget { + /// Take a permit to send one message. None when the budget is spent. + /// Dropping a permit unused returns its capacity. + fn reserve(&mut self) -> Option + '_>; + + /// Reserve up to `count` sends. The permit reports how many it actually holds - + /// possibly zero, possibly fewer than requested. Unspent sends return to the + /// budget when the permit drops. + fn reserve_many(&mut self, count: usize) -> impl SendManyPermit + '_; +} + +/// A reserved slot in a connection's send queue. +pub trait SendPermit { + /// Spend the permit. + fn send(self, message: T); +} + +/// A reserved run of slots in a connection's send queue. +pub trait SendManyPermit { + /// How many sends this permit holds. + fn reserved(&self) -> usize; + + /// Spend one of the reserved sends. + /// + /// Sending more than `reserved` messages panics, like indexing out of bounds: + /// check `reserved` and send at most that many. + fn send(&mut self, message: T); +} + /// What the connection should do after processing a batch of inbound messages. #[derive(Debug, PartialEq, Eq)] pub enum ReactorStatus { 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/src/server/rpc_submitter.rs b/protosocket-rpc/src/server/rpc_submitter.rs index 9d93811..31d5efd 100644 --- a/protosocket-rpc/src/server/rpc_submitter.rs +++ b/protosocket-rpc/src/server/rpc_submitter.rs @@ -1,12 +1,11 @@ use std::{ - collections::VecDeque, - pin::Pin, + pin::{Pin, pin}, task::{Context, Poll}, }; use futures::stream::{FuturesUnordered, StreamFuture}; use futures::{Stream, StreamExt}; -use protosocket::MessageReactor; +use protosocket::{MessageReactor, SendManyPermit, SendPermit}; use crate::{ Message, ProtosocketControlCode, @@ -20,15 +19,17 @@ type PooledRpc = RpcStream< ::StreamType, >; -/// A MessageReactor that hosts a ConnectionService's rpcs and drives them within the -/// connection's send budget. +/// A MessageReactor that hosts a ConnectionService's rpcs. /// -/// New rpcs are registered from inbound messages. Their completions - unary and streaming -/// alike - live in one pool and are only advanced by `poll_outbound_many`, which the -/// connection calls only when it has room to send. This is the backpressure contract: a -/// connection that cannot write does not advance the work that produces responses. The -/// pool yields in readiness order; no priority between unary and streaming rpcs is -/// imposed, and a ready rpc is drained while it is hot, up to the send budget. +/// 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'll probably also want to limit concurrent +/// rpcs, or you'll only get one-sided backpressure. +/// A connection that cannot write does not advance the response futures. +/// +/// Rpc messages are sent in readiness order. There is no ordering across rpcs, however +/// streaming rpcs still get relative ordering for their own messages. You'll receive streams +/// in the order they were yielded by your rpc. pub struct RpcSubmitter where TConnectionService: ConnectionService, @@ -37,7 +38,7 @@ where aborts: AbortionTracker, /// Message ids of rpcs to reject. Small and self-limiting: bounded by the inbound /// messages processed between sends. - rejections: VecDeque, + rejections: Vec, rpcs: FuturesUnordered>>, } @@ -67,8 +68,6 @@ where } } - /// Turn an rpc event into the message that goes on the wire, retiring the rpc's - /// cancellation bookkeeping on terminal events. fn message_for_event( &mut self, id: u64, @@ -120,7 +119,7 @@ where } RpcKind::Cancelled => { log::debug!("rejecting rpc {message_id}"); - self.rejections.push_back(message_id); + self.rejections.push(message_id); } }, ProtosocketControlCode::Cancel => { @@ -151,59 +150,56 @@ where structurally_pinned_connection_server.poll(context) } - /// Produce responses in rpc readiness order, draining each ready rpc while it is - /// hot to amortize pool re-insertion. This is only called when the connection can - /// accept messages for serialization, and `budget` is exactly the room it has, so - /// rpc work only advances when its output has somewhere to go. + /// Produce responses in rpc readiness order. Ready rpcs are drained in runs to + /// amortize pool re-insertion; an rpc that exhausts the budget goes to the back. fn poll_outbound_many( self: Pin<&mut Self>, context: &mut Context<'_>, - sink: &mut impl FnMut(Self::LogicalOutbound), - budget: usize, + outbound: &mut impl protosocket::SendBudget, ) -> Poll> { let me = self.get_mut(); - let mut produced = 0; - while produced < budget { - if let Some(message_id) = me.rejections.pop_front() { - sink(::cancelled(message_id)); + let mut produced = 0usize; + if !me.rejections.is_empty() { + let mut permits = outbound.reserve_many(me.rejections.len()); + for message_id in me.rejections.drain(..permits.reserved()) { + permits.send(::cancelled(message_id)); produced += 1; - continue; } - match Pin::new(&mut me.rpcs).poll_next(context) { + } + loop { + let Some(permit) = outbound.reserve() else { + break; + }; + match pin!(&mut me.rpcs).poll_next(context) { Poll::Ready(Some((first, mut rpc))) => { let Some((id, event)) = first else { - // The rpc's terminal event was already delivered; retire it and - // look at the next ready rpc. + // StreamFuture yielded None: this rpc's stream was already + // exhausted, and there's no more work to do for it. continue; }; let terminal = !matches!(event, RpcStreamEvent::Item(_)); - sink(me.message_for_event(id, event)); + permit.send(me.message_for_event(id, event)); produced += 1; if terminal { continue; } - // Drain this rpc while it is hot, so a busy stream doesn't pay a - // pool re-insertion per message. Fairness across rpcs is round-robin - // per pool visit: an rpc that exhausts the budget goes to the back. - let mut retired = false; - while produced < budget { - match Pin::new(&mut rpc).poll_next(context) { + 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 = !matches!(event, RpcStreamEvent::Item(_)); - sink(me.message_for_event(id, event)); + permit.send(me.message_for_event(id, event)); produced += 1; if terminal { - retired = true; - break; + break true; } } - Poll::Ready(None) => { - retired = true; - break; - } - Poll::Pending => break, + Poll::Ready(None) => break true, + Poll::Pending => break false, } - } + }; if !retired { me.rpcs.push(rpc.into_future()); } From 0a6df5e30fda16f3fb2ae4e01fe2936e09a84283 Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Wed, 15 Jul 2026 08:43:08 -0700 Subject: [PATCH 08/10] docs: split the spawning example into unary and streaming examples --- Cargo.lock | 15 ++ Cargo.toml | 1 + example-proto-spawn-stream/Cargo.toml | 26 +++ example-proto-spawn-stream/src/client.rs | 189 ++++++++++++++++++++ example-proto-spawn-stream/src/messages.rs | 103 +++++++++++ example-proto-spawn-stream/src/server.rs | 138 ++++++++++++++ example-proto-spawn/src/client.rs | 80 +++------ example-proto-spawn/src/messages.rs | 35 +--- example-proto-spawn/src/server.rs | 97 ++++------ protosocket-rpc/src/server/rpc_submitter.rs | 23 +-- 10 files changed, 546 insertions(+), 161 deletions(-) create mode 100644 example-proto-spawn-stream/Cargo.toml create mode 100644 example-proto-spawn-stream/src/client.rs create mode 100644 example-proto-spawn-stream/src/messages.rs create mode 100644 example-proto-spawn-stream/src/server.rs diff --git a/Cargo.lock b/Cargo.lock index bdfeeec..98477f3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -508,6 +508,21 @@ dependencies = [ "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" diff --git a/Cargo.toml b/Cargo.toml index d2741c3..53d0688 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ members = [ "example-pool", "example-proto", "example-proto-spawn", + "example-proto-spawn-stream", "example-proto-stream", "example-proto-tls", "example-telnet", 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..09fe0f0 --- /dev/null +++ b/example-proto-spawn-stream/src/server.rs @@ -0,0 +1,138 @@ +//! 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::{Stream, StreamExt, future::BoxFuture, stream::BoxStream}; +use messages::{EchoRequest, EchoStream, Request, Response}; +use protosocket::{PooledEncoder, StreamWithAddress, TcpSocketListener}; +use protosocket_prost::{ProstDecoder, ProstSerializer}; +use protosocket_rpc::{ + ProtosocketControlCode, + server::{ConnectionService, RpcKind, SocketService}, +}; +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 Codec = ( + PooledEncoder>, + ProstDecoder, + ); + type ConnectionService = DemoRpcConnectionServer; + type SocketListener = TcpSocketListener; + + fn codec(&self) -> Self::Codec { + Default::default() + } + + fn new_stream_service(&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 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/src/client.rs b/example-proto-spawn/src/client.rs index 8539d74..7a40385 100644 --- a/example-proto-spawn/src/client.rs +++ b/example-proto-spawn/src/client.rs @@ -6,8 +6,7 @@ use std::{ time::{Duration, Instant, SystemTime, UNIX_EPOCH}, }; -use futures::StreamExt; -use messages::{EchoRequest, EchoResponseKind, Request, Response, ResponseBehavior}; +use messages::{EchoRequest, Request, Response}; use protosocket::PooledEncoder; use protosocket_prost::{ProstDecoder, ProstSerializer}; use protosocket_rpc::{ @@ -38,6 +37,7 @@ 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; @@ -70,6 +70,7 @@ async fn run_main() -> Result<(), Box> { request_ids.clone(), client.clone(), response_count.clone(), + shed_count.clone(), latency.clone(), )); } @@ -77,6 +78,7 @@ async fn run_main() -> Result<(), Box> { let metrics = tokio::spawn(print_periodic_metrics( response_count, + shed_count, latency, concurrent_count, )); @@ -98,6 +100,7 @@ async fn run_main() -> Result<(), Box> { async fn print_periodic_metrics( response_count: Arc, + shed_count: Arc, latency: Arc, concurrent_count: Arc, ) { @@ -127,9 +130,10 @@ async fn print_periodic_metrics( .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 p90: {p90:6.1}µs p999: {p999:6.1}µs p9999: {p9999:6.1}µs concurrency: {concurrent}" + "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}" ); } } @@ -139,17 +143,12 @@ async fn generate_traffic( 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); - // Alternate unary and streaming so both spawned server paths get exercised. - let behavior = if i.is_multiple_of(2) { - ResponseBehavior::Unary - } else { - ResponseBehavior::Stream - }; let request = Request { request_id: i, code: ProtosocketControlCode::Normal as u32, @@ -160,31 +159,23 @@ async fn generate_traffic( .expect("time works") .as_nanos() as u64, }), - response_behavior: behavior as i32, }; concurrent_count.fetch_add(1, Ordering::Relaxed); - match behavior { - ResponseBehavior::Unary => match client.send_unary(request) { - Ok(completion) => { - let response = completion.await.expect("response must be successful"); + match client.send_unary(request) { + Ok(completion) => match completion.await { + Ok(response) => { handle_unary_response(response, &metrics_count, &metrics_latency); } - Err(e) => { - log::error!("send should work: {e:?}"); - return; - } - }, - ResponseBehavior::Stream => 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; + // 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); } @@ -199,8 +190,8 @@ fn handle_unary_response( let request_id = response.request_id; assert_ne!(response.request_id, 0, "received bad message"); - match response.kind { - Some(EchoResponseKind::Echo(echo)) => { + match response.body { + Some(echo) => { assert_eq!(request_id, echo.message.parse().unwrap_or_default()); let latency = SystemTime::now() @@ -210,35 +201,6 @@ fn handle_unary_response( - echo.nanotime; let _ = metrics_latency.increment(latency); } - Some(EchoResponseKind::Stream(_char_response)) => { - log::error!("got a stream response for a unary request"); - } - None => { - log::warn!("no response body"); - } - } -} - -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.kind { - Some(EchoResponseKind::Echo(_echo)) => { - log::error!("got a unary response for a stream request"); - } - Some(EchoResponseKind::Stream(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/src/messages.rs b/example-proto-spawn/src/messages.rs index eecf5de..2bc7625 100644 --- a/example-proto-spawn/src/messages.rs +++ b/example-proto-spawn/src/messages.rs @@ -12,15 +12,6 @@ pub struct Request { pub code: u32, #[prost(message, tag = "3")] pub body: Option, - #[prost(enumeration = "ResponseBehavior", tag = "4")] - pub response_behavior: i32, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, prost::Enumeration)] -#[repr(i32)] -pub enum ResponseBehavior { - Unary = 0, - Stream = 1, } #[derive(Clone, PartialEq, Eq, prost::Message)] @@ -37,16 +28,8 @@ pub struct Response { pub request_id: u64, #[prost(uint32, tag = "2")] pub code: u32, - #[prost(oneof = "EchoResponseKind", tags = "3, 4")] - pub kind: Option, -} - -#[derive(Clone, PartialEq, Eq, prost::Oneof)] -pub enum EchoResponseKind { #[prost(message, tag = "3")] - Echo(EchoResponse), - #[prost(message, tag = "4")] - Stream(EchoStream), + pub body: Option, } #[derive(Clone, PartialEq, Eq, prost::Message)] @@ -57,16 +40,6 @@ pub struct EchoResponse { pub nanotime: u64, } -#[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 @@ -81,7 +54,6 @@ impl protosocket_rpc::Message for Request { request_id, code: ProtosocketControlCode::Cancel as u32, body: None, - response_behavior: ResponseBehavior::Unary as i32, } } @@ -94,7 +66,6 @@ impl protosocket_rpc::Message for Request { request_id, code: ProtosocketControlCode::End as u32, body: None, - response_behavior: ResponseBehavior::Unary as i32, } } } @@ -112,7 +83,7 @@ impl protosocket_rpc::Message for Response { Response { request_id, code: ProtosocketControlCode::Cancel as u32, - kind: None, + body: None, } } @@ -124,7 +95,7 @@ impl protosocket_rpc::Message for Response { Self { request_id, code: ProtosocketControlCode::End as u32, - kind: None, + body: None, } } } diff --git a/example-proto-spawn/src/server.rs b/example-proto-spawn/src/server.rs index 9e41a14..c98335b 100644 --- a/example-proto-spawn/src/server.rs +++ b/example-proto-spawn/src/server.rs @@ -1,17 +1,16 @@ -//! A protosocket-rpc server that spawns its rpcs onto the runtime. +//! 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. Use a completion future for unary work, and a producer task behind -//! a bounded channel for streams - the channel capacity is how far the producer can run -//! ahead of the peer. +//! runs regardless. Bound what you spawn - here, rpcs shed load when the connection's +//! spawn slots are all in flight. -use std::pin::pin; use std::sync::atomic::AtomicUsize; +use std::sync::Arc; use std::time::Duration; -use futures::{future::BoxFuture, stream::BoxStream, FutureExt, Stream, StreamExt}; -use messages::{EchoRequest, EchoResponse, EchoStream, Request, Response, ResponseBehavior}; +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::{ @@ -19,6 +18,7 @@ use protosocket_rpc::{ Message, ProtosocketControlCode, }; use tokio::net::TcpStream; +use tokio::sync::Semaphore; mod messages; @@ -78,12 +78,16 @@ impl SocketService for DemoRpcSocketService { 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 Request = Request; @@ -97,44 +101,33 @@ impl ConnectionService for DemoRpcConnectionServer { ) -> 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 => { - // Spawn the work and return a future that completes with the task. - let work = tokio::spawn(echo_request(request_id, echo)); - RpcKind::Unary( - async move { - match work.await { - Ok(response) => response, - Err(join_error) => { - log::error!("rpc task failed: {join_error}"); - Response::cancelled(request_id) + 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(), - ) - } - ResponseBehavior::Stream => { - // Spawn the producer behind a bounded channel. The producer waits - // when 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(), - ) + ) + } + Err(_no_slot) => RpcKind::Cancelled, } - }, + } None => { log::warn!("received empty echo request id {request_id}"); RpcKind::Cancelled @@ -144,29 +137,15 @@ impl ConnectionService for DemoRpcConnectionServer { } async fn echo_request(request_id: u64, echo: EchoRequest) -> Response { - // Pretend this is compute-heavy work that deserves its own task. - tokio::time::sleep(Duration::from_micros(1)).await; + // 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, - kind: Some(messages::EchoResponseKind::Echo(EchoResponse { + body: Some(EchoResponse { message: echo.message, nanotime: echo.nanotime, - })), + }), } } - -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, - kind: Some(messages::EchoResponseKind::Stream(EchoStream { - message: (c as char).to_string(), - nanotime, - sequence: sequence as u64, - })), - }, - )) -} diff --git a/protosocket-rpc/src/server/rpc_submitter.rs b/protosocket-rpc/src/server/rpc_submitter.rs index 31d5efd..e60dc43 100644 --- a/protosocket-rpc/src/server/rpc_submitter.rs +++ b/protosocket-rpc/src/server/rpc_submitter.rs @@ -1,5 +1,5 @@ use std::{ - pin::{Pin, pin}, + pin::{pin, Pin}, task::{Context, Poll}, }; @@ -8,8 +8,8 @@ use futures::{Stream, StreamExt}; use protosocket::{MessageReactor, SendManyPermit, SendPermit}; use crate::{ + server::{abortion_tracker::AbortionTracker, ConnectionService, RpcKind}, Message, ProtosocketControlCode, - server::{ConnectionService, RpcKind, abortion_tracker::AbortionTracker}, }; use super::rpc_stream::{RpcStream, RpcStreamEvent}; @@ -23,13 +23,10 @@ type PooledRpc = RpcStream< /// /// 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'll probably also want to limit concurrent -/// rpcs, or you'll only get one-sided backpressure. -/// A connection that cannot write does not advance the response futures. +/// 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. There is no ordering across rpcs, however -/// streaming rpcs still get relative ordering for their own messages. You'll receive streams -/// in the order they were yielded by your rpc. +/// Rpc messages are sent in readiness order. pub struct RpcSubmitter where TConnectionService: ConnectionService, @@ -150,8 +147,6 @@ where structurally_pinned_connection_server.poll(context) } - /// Produce responses in rpc readiness order. Ready rpcs are drained in runs to - /// amortize pool re-insertion; an rpc that exhausts the budget goes to the back. fn poll_outbound_many( self: Pin<&mut Self>, context: &mut Context<'_>, @@ -159,6 +154,7 @@ where ) -> Poll> { let me = self.get_mut(); let mut produced = 0usize; + // First, yield any rejected RPC responses. if !me.rejections.is_empty() { let mut permits = outbound.reserve_many(me.rejections.len()); for message_id in me.rejections.drain(..permits.reserved()) { @@ -166,6 +162,7 @@ where produced += 1; } } + // Then poll ready RPCs until we either run out of room to send responses or run out of rpcs. loop { let Some(permit) = outbound.reserve() else { break; @@ -183,6 +180,10 @@ where 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; @@ -205,7 +206,7 @@ where } } // An empty pool is not a finished reactor: new rpcs arrive from inbound - // processing, which wakes this connection on its own. + // processing, which wakes this connection. Poll::Ready(None) | Poll::Pending => break, } } From 154e1ab63aab147f1a03b36a15822784fda86be0 Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Wed, 15 Jul 2026 11:01:25 -0700 Subject: [PATCH 09/10] refactor: put codecs on reactors and drop the panicking many-permit --- Cargo.lock | 1 + example-messagepack/src/server.rs | 14 +- example-pool/src/server.rs | 2 +- example-proto-spawn-stream/src/server.rs | 14 +- example-proto-spawn/src/server.rs | 10 +- example-proto-stream/src/server.rs | 14 +- example-proto-tls/src/server.rs | 10 +- example-proto/src/server.rs | 14 +- example-telnet/src/main.rs | 2 +- .../examples/custom_allocator.rs | 1 + protosocket-connection/src/connection.rs | 138 +++--------------- protosocket-connection/src/lib.rs | 6 +- protosocket-connection/src/message_reactor.rs | 57 +++----- protosocket-connection/src/send_budget.rs | 58 ++++++++ protosocket-rpc/Cargo.toml | 1 + protosocket-rpc/src/client/configuration.rs | 11 +- .../src/client/reactor/completion_reactor.rs | 63 ++++---- protosocket-rpc/src/client/rpc_client.rs | 49 +++++-- protosocket-rpc/src/server/rpc_stream.rs | 21 ++- protosocket-rpc/src/server/rpc_submitter.rs | 29 ++-- protosocket-rpc/src/server/server_traits.rs | 29 ++-- protosocket-rpc/src/server/socket_server.rs | 11 +- protosocket-server/src/connection_server.rs | 14 +- 23 files changed, 268 insertions(+), 301 deletions(-) create mode 100644 protosocket-connection/src/send_budget.rs diff --git a/Cargo.lock b/Cargo.lock index 98477f3..b59eb98 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1460,6 +1460,7 @@ name = "protosocket-rpc" version = "1.1.0" dependencies = [ "ahash", + "bytes", "futures", "level-runtime", "log", diff --git a/example-messagepack/src/server.rs b/example-messagepack/src/server.rs index 551332c..89b5ae8 100644 --- a/example-messagepack/src/server.rs +++ b/example-messagepack/src/server.rs @@ -55,16 +55,10 @@ 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(&self) -> ::Codec { Default::default() } @@ -82,6 +76,12 @@ 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; diff --git a/example-pool/src/server.rs b/example-pool/src/server.rs index 81c5f75..b974550 100644 --- a/example-pool/src/server.rs +++ b/example-pool/src/server.rs @@ -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/src/server.rs b/example-proto-spawn-stream/src/server.rs index 09fe0f0..8a426b0 100644 --- a/example-proto-spawn-stream/src/server.rs +++ b/example-proto-spawn-stream/src/server.rs @@ -8,13 +8,13 @@ use std::pin::pin; use std::sync::atomic::AtomicUsize; -use futures::{Stream, StreamExt, future::BoxFuture, stream::BoxStream}; +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::{ - ProtosocketControlCode, server::{ConnectionService, RpcKind, SocketService}, + ProtosocketControlCode, }; use tokio::net::TcpStream; @@ -61,14 +61,10 @@ async fn run_main() -> Result<(), Box> { struct DemoRpcSocketService; impl SocketService for DemoRpcSocketService { - type Codec = ( - PooledEncoder>, - ProstDecoder, - ); type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&self) -> Self::Codec { + fn codec(&self) -> ::Codec { Default::default() } @@ -84,6 +80,10 @@ 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>; diff --git a/example-proto-spawn/src/server.rs b/example-proto-spawn/src/server.rs index c98335b..527594d 100644 --- a/example-proto-spawn/src/server.rs +++ b/example-proto-spawn/src/server.rs @@ -63,14 +63,10 @@ async fn run_main() -> Result<(), Box> { struct DemoRpcSocketService; impl SocketService for DemoRpcSocketService { - type Codec = ( - PooledEncoder>, - ProstDecoder, - ); type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&self) -> Self::Codec { + fn codec(&self) -> ::Codec { Default::default() } @@ -90,6 +86,10 @@ struct DemoRpcConnectionServer { spawn_limit: Arc, } impl ConnectionService for DemoRpcConnectionServer { + type Codec = ( + PooledEncoder>, + ProstDecoder, + ); type Request = Request; type Response = Response; type UnaryFutureType = BoxFuture<'static, Response>; diff --git a/example-proto-stream/src/server.rs b/example-proto-stream/src/server.rs index 0cd28e5..4135e02 100644 --- a/example-proto-stream/src/server.rs +++ b/example-proto-stream/src/server.rs @@ -61,16 +61,10 @@ 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(&self) -> ::Codec { ( PooledEncoder::new_with_pool_size(64, Default::default()), ProstDecoder::default(), @@ -94,6 +88,12 @@ 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; diff --git a/example-proto-tls/src/server.rs b/example-proto-tls/src/server.rs index a8180cb..1feae3c 100644 --- a/example-proto-tls/src/server.rs +++ b/example-proto-tls/src/server.rs @@ -88,14 +88,10 @@ 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(&self) -> ::Codec { Default::default() } @@ -115,6 +111,10 @@ 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>; diff --git a/example-proto/src/server.rs b/example-proto/src/server.rs index 6cf11cc..4f9e10f 100644 --- a/example-proto/src/server.rs +++ b/example-proto/src/server.rs @@ -61,16 +61,10 @@ 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(&self) -> ::Codec { ( PooledEncoder::new_with_pool_size(64, Default::default()), ProstDecoder::default(), @@ -91,6 +85,12 @@ 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; diff --git a/example-telnet/src/main.rs b/example-telnet/src/main.rs index 2658389..5056a85 100644 --- a/example-telnet/src/main.rs +++ b/example-telnet/src/main.rs @@ -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 e760ff8..f3697a9 100644 --- a/protosocket-connection/src/connection.rs +++ b/protosocket-connection/src/connection.rs @@ -10,82 +10,12 @@ use bytes::Buf; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use crate::{ - encoding::Codec, interrupted, - message_reactor::{MessageReactor, ReactorStatus, SendBudget, SendManyPermit, SendPermit}, + message_reactor::{MessageReactor, ReactorStatus}, + send_budget::SendBudget, would_block, Decoder, DeserializeError, Encoder, }; -/// The connection's send capacity, spent by encoding directly into the send buffer. -struct EncodingSendBudget<'a, TCodec: Codec> { - codec: &'a mut TCodec, - send_buffer: &'a mut VecDeque<::Serialized>, - remaining: usize, -} - -impl SendBudget<::Message> for EncodingSendBudget<'_, TCodec> { - fn reserve(&mut self) -> Option::Message> + '_> { - if self.remaining == 0 { - None - } else { - Some(EncodingSendPermit { budget: self }) - } - } - - fn reserve_many( - &mut self, - count: usize, - ) -> impl SendManyPermit<::Message> + '_ { - let reserved = count.min(self.remaining); - EncodingSendManyPermit { - budget: self, - reserved, - } - } -} - -/// A reserved slot in the send buffer. -struct EncodingSendPermit<'a, 'b, TCodec: Codec> { - budget: &'a mut EncodingSendBudget<'b, TCodec>, -} - -impl SendPermit<::Message> - for EncodingSendPermit<'_, '_, TCodec> -{ - fn send(self, message: ::Message) { - self.budget.remaining -= 1; - let buffer = self.budget.codec.encode(message); - log::trace!("serialized reactor message: {}b", buffer.remaining()); - self.budget.send_buffer.push_back(buffer); - } -} - -/// A reserved run of slots in the send buffer. -struct EncodingSendManyPermit<'a, 'b, TCodec: Codec> { - budget: &'a mut EncodingSendBudget<'b, TCodec>, - reserved: usize, -} - -impl SendManyPermit<::Message> - for EncodingSendManyPermit<'_, '_, TCodec> -{ - fn reserved(&self) -> usize { - self.reserved - } - - fn send(&mut self, message: ::Message) { - assert!( - 0 < self.reserved, - "sent more messages than the permit reserved" - ); - self.reserved -= 1; - self.budget.remaining -= 1; - let buffer = self.budget.codec.encode(message); - log::trace!("serialized reactor message: {}b", buffer.remaining()); - self.budget.send_buffer.push_back(buffer); - } -} - /// A bidirectional, message-oriented AsyncRead/AsyncWrite stream wrapper. /// /// Connections are Futures that you spawn. @@ -100,31 +30,23 @@ impl SendManyPermit<::Message> 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; @@ -138,25 +60,15 @@ impl< } } -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 = (); @@ -201,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") @@ -229,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. /// @@ -242,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, @@ -436,11 +338,7 @@ impl< codec, .. } = self; - let mut outbound = EncodingSendBudget { - codec, - send_buffer, - remaining, - }; + 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); diff --git a/protosocket-connection/src/lib.rs b/protosocket-connection/src/lib.rs index e359706..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; @@ -23,12 +24,11 @@ pub use encoding::OwnedBuffer; pub use error::DeserializeError; pub use message_reactor::MessageReactor; pub use message_reactor::ReactorStatus; -pub use message_reactor::SendBudget; -pub use message_reactor::SendManyPermit; -pub use message_reactor::SendPermit; 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 c101b5e..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<'_>, @@ -34,63 +41,35 @@ pub trait MessageReactor: 'static { } /// Called from the connection's driver task when messages from the outbound queue - /// are sent. Messages produced by `poll_outbound_many` do not come through here. + /// 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 into `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. If your message source models lag (like - /// `tokio::sync::broadcast`), a slow peer engages that model instead of buffering - /// without bound. + /// messages is not advanced. /// - /// Messages you emit here are yours: do any bookkeeping before you send them. + /// 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 default. The connection - /// closes when its outbound queue is closed and this returns `Ready(None)`. + /// 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 impl SendBudget, + _outbound: &mut SendBudget<'_, Self::Codec>, ) -> std::task::Poll> { std::task::Poll::Ready(None) } } -/// A bounded lease on a connection's send capacity. -pub trait SendBudget { - /// Take a permit to send one message. None when the budget is spent. - /// Dropping a permit unused returns its capacity. - fn reserve(&mut self) -> Option + '_>; - - /// Reserve up to `count` sends. The permit reports how many it actually holds - - /// possibly zero, possibly fewer than requested. Unspent sends return to the - /// budget when the permit drops. - fn reserve_many(&mut self, count: usize) -> impl SendManyPermit + '_; -} - -/// A reserved slot in a connection's send queue. -pub trait SendPermit { - /// Spend the permit. - fn send(self, message: T); -} - -/// A reserved run of slots in a connection's send queue. -pub trait SendManyPermit { - /// How many sends this permit holds. - fn reserved(&self) -> usize; - - /// Spend one of the reserved sends. - /// - /// Sending more than `reserved` messages panics, like indexing out of bounds: - /// check `reserved` and send at most that many. - fn send(&mut self, message: T); -} - /// What the connection should do after processing a batch of inbound messages. #[derive(Debug, PartialEq, Eq)] pub enum ReactorStatus { 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-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/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/rpc_stream.rs b/protosocket-rpc/src/server/rpc_stream.rs index 95bfa82..2fc0b49 100644 --- a/protosocket-rpc/src/server/rpc_stream.rs +++ b/protosocket-rpc/src/server/rpc_stream.rs @@ -31,21 +31,30 @@ enum RpcStreamKind { Done, } -/// An event from a pooled rpc, tagged with terminal-ness so the caller knows what, if -/// anything, to put on the wire after it. +/// An event from a pooled rpc. #[derive(Debug)] pub enum RpcStreamEvent { - /// A streaming rpc produced a response item. + /// A streaming rpc produced a response item. Stream continues. Item(T), - /// A unary rpc completed with its single response. Terminal; the response is the - /// completion, no trailing control message belongs on the wire. + /// A unary rpc completed with its single response. Terminal. Complete(T), - /// A streaming rpc's stream finished. Terminal; the peer should be told the rpc ended. + /// 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)) diff --git a/protosocket-rpc/src/server/rpc_submitter.rs b/protosocket-rpc/src/server/rpc_submitter.rs index e60dc43..ae21ded 100644 --- a/protosocket-rpc/src/server/rpc_submitter.rs +++ b/protosocket-rpc/src/server/rpc_submitter.rs @@ -5,7 +5,7 @@ use std::{ use futures::stream::{FuturesUnordered, StreamFuture}; use futures::{Stream, StreamExt}; -use protosocket::{MessageReactor, SendManyPermit, SendPermit}; +use protosocket::{MessageReactor, SendBudget}; use crate::{ server::{abortion_tracker::AbortionTracker, ConnectionService, RpcKind}, @@ -33,8 +33,6 @@ where { connection_server: TConnectionService, aborts: AbortionTracker, - /// Message ids of rpcs to reject. Small and self-limiting: bounded by the inbound - /// messages processed between sends. rejections: Vec, rpcs: FuturesUnordered>>, } @@ -96,6 +94,7 @@ impl MessageReactor for RpcSubmitter where TConnectionService: ConnectionService, { + type Codec = TConnectionService::Codec; type Inbound = TConnectionService::Request; type Outbound = TConnectionService::Response; type LogicalOutbound = TConnectionService::Response; @@ -150,23 +149,23 @@ where fn poll_outbound_many( self: Pin<&mut Self>, context: &mut Context<'_>, - outbound: &mut impl protosocket::SendBudget, + outbound: &mut SendBudget<'_, TConnectionService::Codec>, ) -> Poll> { let me = self.get_mut(); let mut produced = 0usize; // First, yield any rejected RPC responses. - if !me.rejections.is_empty() { - let mut permits = outbound.reserve_many(me.rejections.len()); - for message_id in me.rejections.drain(..permits.reserved()) { - permits.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. - loop { + 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 { @@ -174,7 +173,7 @@ where // exhausted, and there's no more work to do for it. continue; }; - let terminal = !matches!(event, RpcStreamEvent::Item(_)); + let terminal = event.is_terminal(); permit.send(me.message_for_event(id, event)); produced += 1; if terminal { @@ -190,7 +189,7 @@ where }; match pin!(&mut rpc).poll_next(context) { Poll::Ready(Some((id, event))) => { - let terminal = !matches!(event, RpcStreamEvent::Item(_)); + let terminal = event.is_terminal(); permit.send(me.message_for_event(id, event)); produced += 1; if terminal { diff --git a/protosocket-rpc/src/server/server_traits.rs b/protosocket-rpc/src/server/server_traits.rs index 1ab1a3e..14441fd 100644 --- a/protosocket-rpc/src/server/server_traits.rs +++ b/protosocket-rpc/src/server/server_traits.rs @@ -11,24 +11,14 @@ use crate::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(&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 @@ -56,6 +46,15 @@ pub trait SocketService: 'static { /// /// 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. @@ -91,10 +90,10 @@ pub trait ConnectionService: Unpin + 'static { /// Type of rpc to be completed pub enum RpcKind { - /// This is a unary rpc. It will complete with a single response. + /// A rpc that produces a single response. Unary(Unary), - /// This is a streaming rpc. It will complete with a stream of responses. + /// A rpc that produces a stream of responses. Streaming(Streaming), - /// Do not process this rpc. The initiating message is answered with a cancellation. + /// 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 cb5ca13..893183f 100644 --- a/protosocket-rpc/src/server/socket_server.rs +++ b/protosocket-rpc/src/server/socket_server.rs @@ -13,6 +13,9 @@ use crate::server::Spawn; use super::rpc_submitter::RpcSubmitter; 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,15 +44,14 @@ 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, @@ -83,7 +85,6 @@ where TSpawnConnection: Spawn< Connection< ::Stream, - TSocketService::Codec, RpcSubmitter, >, >, @@ -129,7 +130,6 @@ where TSpawnConnection: Spawn< Connection< ::Stream, - TSocketService::Codec, RpcSubmitter, >, >, @@ -150,7 +150,6 @@ where #[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..f359407 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,11 +14,9 @@ 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; @@ -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>, ); } From d6c1a4b9120bb321a66b692a970539e323487ea0 Mon Sep 17 00:00:00 2001 From: Dylan Abraham Date: Wed, 15 Jul 2026 16:13:19 -0700 Subject: [PATCH 10/10] refactor: take &mut self in service and connector factory methods --- example-messagepack/src/server.rs | 7 +++++-- example-pool/src/server.rs | 4 ++-- example-proto-spawn-stream/src/server.rs | 7 +++++-- example-proto-spawn/src/server.rs | 7 +++++-- example-proto-stream/src/server.rs | 7 +++++-- example-proto-tls/src/server.rs | 4 ++-- example-proto/src/server.rs | 7 +++++-- example-telnet/src/main.rs | 4 ++-- protosocket-rpc/src/server/server_traits.rs | 4 ++-- protosocket-server/src/connection_server.rs | 4 ++-- 10 files changed, 35 insertions(+), 20 deletions(-) diff --git a/example-messagepack/src/server.rs b/example-messagepack/src/server.rs index 89b5ae8..f01b7aa 100644 --- a/example-messagepack/src/server.rs +++ b/example-messagepack/src/server.rs @@ -58,11 +58,14 @@ impl SocketService for DemoRpcSocketService { type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&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(), diff --git a/example-pool/src/server.rs b/example-pool/src/server.rs index b974550..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 { diff --git a/example-proto-spawn-stream/src/server.rs b/example-proto-spawn-stream/src/server.rs index 8a426b0..679fbb3 100644 --- a/example-proto-spawn-stream/src/server.rs +++ b/example-proto-spawn-stream/src/server.rs @@ -64,11 +64,14 @@ impl SocketService for DemoRpcSocketService { type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&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(), diff --git a/example-proto-spawn/src/server.rs b/example-proto-spawn/src/server.rs index 527594d..a75f4c0 100644 --- a/example-proto-spawn/src/server.rs +++ b/example-proto-spawn/src/server.rs @@ -66,11 +66,14 @@ impl SocketService for DemoRpcSocketService { type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&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(), diff --git a/example-proto-stream/src/server.rs b/example-proto-stream/src/server.rs index 4135e02..87e7a1c 100644 --- a/example-proto-stream/src/server.rs +++ b/example-proto-stream/src/server.rs @@ -64,14 +64,17 @@ impl SocketService for DemoRpcSocketService { type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&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(), diff --git a/example-proto-tls/src/server.rs b/example-proto-tls/src/server.rs index 1feae3c..a1b139c 100644 --- a/example-proto-tls/src/server.rs +++ b/example-proto-tls/src/server.rs @@ -91,12 +91,12 @@ impl SocketService for DemoRpcSocketService { type ConnectionService = DemoRpcConnectionServer; type SocketListener = TlsSocketListener; - fn codec(&self) -> ::Codec { + fn codec(&mut self) -> ::Codec { Default::default() } fn new_stream_service( - &self, + &mut self, stream: &StreamWithAddress>, ) -> Self::ConnectionService { DemoRpcConnectionServer { diff --git a/example-proto/src/server.rs b/example-proto/src/server.rs index 4f9e10f..3adcab4 100644 --- a/example-proto/src/server.rs +++ b/example-proto/src/server.rs @@ -64,14 +64,17 @@ impl SocketService for DemoRpcSocketService { type ConnectionService = DemoRpcConnectionServer; type SocketListener = TcpSocketListener; - fn codec(&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(), diff --git a/example-telnet/src/main.rs b/example-telnet/src/main.rs index 5056a85..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 { diff --git a/protosocket-rpc/src/server/server_traits.rs b/protosocket-rpc/src/server/server_traits.rs index 14441fd..f037977 100644 --- a/protosocket-rpc/src/server/server_traits.rs +++ b/protosocket-rpc/src/server/server_traits.rs @@ -18,13 +18,13 @@ pub trait SocketService: 'static { type SocketListener: SocketListener; /// Create a new message codec for a connection. - fn codec(&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; } diff --git a/protosocket-server/src/connection_server.rs b/protosocket-server/src/connection_server.rs index f359407..f58c966 100644 --- a/protosocket-server/src/connection_server.rs +++ b/protosocket-server/src/connection_server.rs @@ -21,12 +21,12 @@ pub trait ServerConnector: Unpin { 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;