Connection driven rpc backpressure

Cargo.toml

@@ -6,6 +6,8 @@ members = [
     "example-messagepack",
     "example-pool",
     "example-proto",
+    "example-proto-spawn",
+    "example-proto-spawn-stream",
     "example-proto-stream",
     "example-proto-tls",
     "example-telnet",

example-messagepack/src/server.rs

@@ -1,11 +1,11 @@
 use std::sync::atomic::AtomicUsize;
 
-use futures::Stream;
+use futures::{stream::BoxStream, Stream, StreamExt};
 use messages::{EchoRequest, EchoResponse, EchoStream, Request, Response, ResponseBehavior};
 use protosocket::{PooledEncoder, StreamWithAddress, TcpSocketListener};
 use protosocket_rpc::{
-    server::{ConnectionService, RpcResponder, SocketService},
-    Message, ProtosocketControlCode,
+    server::{ConnectionService, RpcKind, SocketService},
+    ProtosocketControlCode,
 };
 use tokio::net::TcpStream;
 
@@ -55,20 +55,17 @@ async fn run_main() -> Result<(), Box<dyn std::error::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<protosocket_messagepack::MessagePackSerializer<Response>>,
-        protosocket_messagepack::MessagePackDecoder<Request>,
-    );
     type ConnectionService = DemoRpcConnectionServer;
     type SocketListener = TcpSocketListener;
 
-    fn codec(&self) -> Self::Codec {
+    fn codec(&mut self) -> <Self::ConnectionService as ConnectionService>::Codec {
         Default::default()
     }
 
-    fn new_stream_service(&self, stream: &StreamWithAddress<TcpStream>) -> Self::ConnectionService {
+    fn new_stream_service(
+        &mut self,
+        stream: &StreamWithAddress<TcpStream>,
+    ) -> Self::ConnectionService {
         log::info!("new connection server {}", stream.address());
         DemoRpcConnectionServer {
             address: stream.address(),
@@ -82,32 +79,36 @@ struct DemoRpcConnectionServer {
     address: std::net::SocketAddr,
 }
 impl ConnectionService for DemoRpcConnectionServer {
+    type Codec = (
+        // Use a pooled encoder to amortize memory allocation cost.
+        // Each connection gets its own little memory pool.
+        PooledEncoder<protosocket_messagepack::MessagePackSerializer<Response>>,
+        protosocket_messagepack::MessagePackDecoder<Request>,
+    );
     type Request = Request;
     type Response = Response;
+    type UnaryFutureType = futures::future::Ready<Response>;
+    type StreamType = BoxStream<'static, Response>;
 
     fn new_rpc(
         &mut self,
         initiating_message: Self::Request,
-        responder: RpcResponder<'_, Self::Response>,
-    ) {
+    ) -> RpcKind<Self::UnaryFutureType, Self::StreamType> {
         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
             }
         }
     }

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<<Self::Codec as Encoder>::Message>,
         _address: &StreamWithAddress<TcpStream>,
     ) -> Self::Reactor {
@@ -46,7 +46,6 @@ impl ServerConnector for ServerContext {
         &self,
         connection: protosocket::Connection<
             <Self::SocketListener as SocketListener>::Stream,
-            Self::Codec,
             Self::Reactor,
         >,
     ) {
@@ -58,6 +57,7 @@ struct PooledReactor {
     outbound: spillway::Sender<String>,
 }
 impl MessageReactor for PooledReactor {
+    type Codec = ByteBufferRingCodec;
     type Inbound = String;
     type Outbound = String;
     type LogicalOutbound = String;

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" }

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<dyn std::error::Error>> {
+    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<dyn std::error::Error>> {
+    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<ProstSerializer<Request>>,
+                ProstDecoder<Response>,
+            ),
+            _,
+        >(
+            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<AtomicUsize>,
+    latency: Arc<histogram::AtomicHistogram>,
+    concurrent_count: Arc<AtomicUsize>,
+) {
+    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<AtomicUsize>,
+    request_ids: Arc<AtomicU64>,
+    client: RpcClient<Request, Response>,
+    metrics_count: Arc<AtomicUsize>,
+    metrics_latency: Arc<histogram::AtomicHistogram>,
+) {
+    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");
+        }
+    }
+}