allow transports to see and manipulate client and server contexts.

Author: axos88

example-service/src/client.rs

@@ -7,7 +7,8 @@
 use clap::Parser;
 use service::{WorldClient, init_tracing};
 use std::{net::SocketAddr, time::Duration};
-use tarpc::{client, context, tokio_serde::formats::Json};
+use futures::{future, SinkExt};
+use tarpc::{client, tokio_serde::formats::Json};
 use tokio::time::sleep;
 use tracing::Instrument;
 use tarpc::context::ClientContext;
@@ -30,9 +31,11 @@ async fn main() -> anyhow::Result<()> {
     let mut transport = tarpc::serde_transport::tcp::connect(flags.server_addr, Json::default);
     transport.config_mut().max_frame_length(usize::MAX);
 
+    let transport = transport.await?.with(|msg: tarpc::ClientMessage<ClientContext, _>| future::ok(msg.map_context(|ctx| ctx.shared_context)));
+
     // WorldClient is generated by the service attribute. It has a constructor `new` that takes a
     // config and any Transport as input.
-    let client = WorldClient::new(client::Config::default(), transport.await?).spawn();
+    let client = WorldClient::new(client::Config::default(), transport).spawn();
 
     let hello = async move {
         let mut context = ClientContext::current();

example-service/src/server.rs

@@ -15,12 +15,9 @@ use std::{
     net::{IpAddr, Ipv6Addr, SocketAddr},
     time::Duration,
 };
-use tarpc::{
-    context,
-    server::{self, Channel, incoming::Incoming},
-    tokio_serde::formats::Json,
-};
+use tarpc::{context, server::{self, Channel, incoming::Incoming}, tokio_serde::formats::Json, ClientMessage};
 use tokio::time;
+use tarpc::context::{ServerContext, SharedContext};
 
 #[derive(Parser)]
 struct Flags {
@@ -62,13 +59,14 @@ async fn main() -> anyhow::Result<()> {
     listener
         // Ignore accept errors.
         .filter_map(|r| future::ready(r.ok()))
+        .map(|t| t.map_ok(|msg: ClientMessage<SharedContext, _>| msg.map_context(|ctx| ServerContext::new(ctx))))
         .map(server::BaseChannel::with_defaults)
         // Limit channels to 1 per IP.
-        .max_channels_per_key(1, |t| t.transport().peer_addr().unwrap().ip())
+        .max_channels_per_key(1, |t| t.transport().get_ref().peer_addr().unwrap().ip())
         // serve is generated by the service attribute. It takes as input any type implementing
         // the generated World trait.
         .map(|channel| {
-            let server = HelloServer(channel.transport().peer_addr().unwrap());
+            let server = HelloServer(channel.transport().get_ref().peer_addr().unwrap());
             channel.execute(server.serve()).for_each(spawn)
         })
         // Max 10 channels.

plugins/Cargo.toml

@@ -30,5 +30,6 @@ proc-macro = true
 [dev-dependencies]
 assert-type-eq = "0.1.0"
 futures = "0.3"
+futures-util = "0.3.31"
 serde = { version = "1.0", features = ["derive"] }
 tarpc = { path = "../tarpc", features = ["serde1"] }

plugins/src/lib.rs

@@ -376,6 +376,7 @@ fn collect_cfg_attrs(rpcs: &[RpcMethod]) -> Vec<Vec<&Attribute>> {
 ///
 /// ```no_run
 /// use tarpc::{client, transport, service, server::{self, Channel}, context::ServerContext};
+/// use futures_util::{TryStreamExt, sink::SinkExt};
 ///
 /// #[service]
 /// pub trait Calculator {
@@ -394,6 +395,13 @@ fn collect_cfg_attrs(rpcs: &[RpcMethod]) -> Vec<Vec<&Attribute>> {
 /// // This could be any transport.
 /// let (client_side, server_side) = transport::channel::unbounded();
 ///
+/// let client_side = client_side.with(|msg: tarpc::ClientMessage<tarpc::context::ClientContext, _>| async move {
+///    Ok(msg.map_context(|ctx| ctx.shared_context))
+/// });
+/// let server_side = server_side.map_ok(|msg: tarpc::ClientMessage<tarpc::context::SharedContext, _>|
+///    msg.map_context(tarpc::context::ServerContext::new)
+/// );
+///
 /// // A client can be made like so:
 /// let client = CalculatorClient::new(client::Config::default(), client_side);
 ///
@@ -738,7 +746,7 @@ impl ServiceGenerator<'_> {
                         ::tarpc::client::RequestDispatch<#request_ident, #response_ident, T>
                     >
                 where
-                    T: ::tarpc::Transport<::tarpc::ClientMessage<#request_ident>, ::tarpc::Response<#response_ident>>
+                    T: ::tarpc::Transport<::tarpc::ClientMessage<::tarpc::context::ClientContext, #request_ident>, ::tarpc::Response<#response_ident>>
                 {
                     let new_client = ::tarpc::client::new(config, transport);
                     ::tarpc::client::NewClient {

tarpc/examples/compression.rs

@@ -9,12 +9,8 @@ use futures::{Sink, SinkExt, Stream, StreamExt, TryStreamExt, prelude::*};
 use serde::{Deserialize, Serialize};
 use serde_bytes::ByteBuf;
 use std::{io, io::Read, io::Write};
-use tarpc::{
-    client, context,
-    serde_transport::tcp,
-    server::{BaseChannel, Channel},
-    tokio_serde::formats::Bincode,
-};
+use tarpc::{client, context, serde_transport::tcp, server::{BaseChannel, Channel}, tokio_serde::formats::Bincode, ClientMessage};
+use tarpc::context::{ClientContext, ServerContext, SharedContext};
 
 /// Type of compression that should be enabled on the request. The transport is free to ignore this.
 #[derive(Debug, PartialEq, Eq, Clone, Copy, Deserialize, Serialize)]
@@ -120,17 +116,22 @@ async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     let mut incoming = tcp::listen("localhost:0", Bincode::default).await?;
+
     let addr = incoming.local_addr();
     tokio::spawn(async move {
         let transport = incoming.next().await.unwrap().unwrap();
-        BaseChannel::with_defaults(add_compression(transport))
+        let transport = add_compression(transport);
+        let transport = transport.map_ok(|msg: ClientMessage<SharedContext, _>| msg.map_context(|ctx| ServerContext::new(ctx)));
+        BaseChannel::with_defaults(transport)
             .execute(HelloServer.serve())
             .for_each(spawn)
             .await;
     });
 
     let transport = tcp::connect(addr, Bincode::default).await?;
-    let client = WorldClient::new(client::Config::default(), add_compression(transport)).spawn();
+    let transport = add_compression(transport);
+    let transport = transport.with(|msg: ClientMessage<ClientContext, _>| future::ok(msg.map_context(|ctx| ctx.shared_context)));
+    let client = WorldClient::new(client::Config::default(), transport).spawn();
 
     println!(
         "{}",

tarpc/examples/custom_transport.rs

@@ -5,8 +5,8 @@
 // https://opensource.org/licenses/MIT.
 
 use futures::prelude::*;
-use tarpc::context::{ClientContext, ServerContext};
-use tarpc::serde_transport as transport;
+use tarpc::context::{ClientContext, ServerContext, SharedContext};
+use tarpc::{serde_transport as transport, ClientMessage};
 use tarpc::server::{BaseChannel, Channel};
 use tarpc::tokio_serde::formats::Bincode;
 use tarpc::tokio_util::codec::length_delimited::LengthDelimitedCodec;
@@ -23,7 +23,6 @@ struct Service;
 impl PingService for Service {
     async fn ping(self, _: &mut ServerContext) {}
 }
-
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     let bind_addr = "/tmp/tarpc_on_unix_example.sock";
@@ -40,6 +39,7 @@ async fn main() -> anyhow::Result<()> {
             let (conn, _addr) = listener.accept().await.unwrap();
             let framed = codec_builder.new_framed(conn);
             let transport = transport::new(framed, Bincode::default());
+            let transport = transport.map_ok(|c: ClientMessage<SharedContext, _>| c.map_context(ServerContext::new));
 
             let fut = BaseChannel::with_defaults(transport)
                 .execute(Service.serve())
@@ -50,6 +50,7 @@ async fn main() -> anyhow::Result<()> {
 
     let conn = UnixStream::connect(bind_addr).await?;
     let transport = transport::new(codec_builder.new_framed(conn), Bincode::default());
+    let transport = transport.with(|msg: ClientMessage<ClientContext, _>| future::ok(msg.map_context(|ctx| ctx.shared_context)));
     PingServiceClient::new(Default::default(), transport)
         .spawn()
         .ping(&mut ClientContext::current())

tarpc/examples/pubsub.rs

@@ -48,15 +48,11 @@ use std::{
     sync::{Arc, Mutex, RwLock},
 };
 use subscriber::Subscriber as _;
-use tarpc::{
-    client, context,
-    serde_transport::tcp,
-    server::{self, Channel},
-    tokio_serde::formats::Json,
-};
+use tarpc::{client, context, serde_transport::tcp, server::{self, Channel}, tokio_serde::formats::Json, ClientMessage};
 use tokio::net::ToSocketAddrs;
 use tracing::info;
 use tracing_subscriber::prelude::*;
+use tarpc::context::{ServerContext, SharedContext};
 
 pub mod subscriber {
     #[tarpc::service]
@@ -104,6 +100,7 @@ impl Subscriber {
     ) -> anyhow::Result<SubscriberHandle> {
         let publisher = tcp::connect(publisher_addr, Json::default).await?;
         let local_addr = publisher.local_addr()?;
+        let publisher = publisher.map_ok(|msg: ClientMessage<SharedContext, _>| msg.map_context(|ctx| ServerContext::new(ctx)));
         let mut handler = server::BaseChannel::with_defaults(publisher).requests();
         let subscriber = Subscriber { local_addr, topics };
         // The first request is for the topics being subscribed to.
@@ -164,6 +161,8 @@ impl Publisher {
             let publisher = connecting_publishers.next().await.unwrap().unwrap();
             info!(publisher.peer_addr = ?publisher.peer_addr(), "publisher connected.");
 
+            let publisher = publisher.map_ok(|msg: ClientMessage<SharedContext, _>| msg.map_context(|ctx| ServerContext::new(ctx)));
+
             server::BaseChannel::with_defaults(publisher)
                 .execute(self.serve())
                 .for_each(spawn)
@@ -183,6 +182,7 @@ impl Publisher {
         tokio::spawn(async move {
             while let Some(conn) = connecting_subscribers.next().await {
                 let subscriber_addr = conn.peer_addr().unwrap();
+                let conn = conn.with(|msg: tarpc::ClientMessage<tarpc::context::ClientContext, _>| future::ok(msg.map_context(|ctx| ctx.shared_context)));
 
                 let tarpc::client::NewClient {
                     client: subscriber,
@@ -341,7 +341,7 @@ async fn main() -> anyhow::Result<()> {
 
     let publisher = publisher::PublisherClient::new(
         client::Config::default(),
-        tcp::connect(addrs.publisher, Json::default).await?,
+        tcp::connect(addrs.publisher, Json::default).await?.with(|msg: tarpc::ClientMessage<tarpc::context::ClientContext, _>| future::ok(msg.map_context(|ctx| ctx.shared_context)))
     )
     .spawn();
 

tarpc/examples/readme.rs

@@ -5,10 +5,8 @@
 // https://opensource.org/licenses/MIT.
 
 use futures::prelude::*;
-use tarpc::{
-    client, context,
-    server::{self, Channel},
-};
+use tarpc::{client, context, server::{self, Channel}, transport, ClientMessage};
+use tarpc::context::{ClientContext, ServerContext, SharedContext};
 
 /// This is the service definition. It looks a lot like a trait definition.
 /// It defines one RPC, hello, which takes one arg, name, and returns a String.
@@ -34,7 +32,10 @@ async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let (client_transport, server_transport) = tarpc::transport::channel::unbounded();
+    let (client_transport, server_transport) = transport::channel::unbounded_mapped(
+        |msg: ClientMessage<ClientContext, _>| msg.map_context(|ctx| ctx.shared_context),
+        |msg: ClientMessage<SharedContext, _>| msg.map_context(ServerContext::new),
+    );
 
     let server = server::BaseChannel::with_defaults(server_transport);
     tokio::spawn(server.execute(HelloServer.serve()).for_each(spawn));

tarpc/examples/tls_over_tcp.rs

@@ -17,8 +17,7 @@ use tokio_rustls::rustls::{
     server::{WebPkiClientVerifier, danger::ClientCertVerifier},
 };
 use tokio_rustls::{TlsAcceptor, TlsConnector};
-
-use tarpc::context::{ClientContext, ServerContext};
+use tarpc::context::{ClientContext, ServerContext, SharedContext};
 use tarpc::serde_transport as transport;
 use tarpc::server::{BaseChannel, Channel};
 use tarpc::tokio_serde::formats::Bincode;
@@ -115,6 +114,7 @@ async fn main() -> anyhow::Result<()> {
             let framed = codec_builder.new_framed(tls_stream);
 
             let transport = transport::new(framed, Bincode::default());
+            let transport = transport.map_ok(|c: tarpc::ClientMessage<SharedContext, _>| c.map_context(|ctx| ServerContext::new(ctx)));
 
             let fut = BaseChannel::with_defaults(transport)
                 .execute(Service.serve())
@@ -144,6 +144,7 @@ async fn main() -> anyhow::Result<()> {
     let stream = connector.connect(domain, stream).await?;
 
     let transport = transport::new(codec_builder.new_framed(stream), Bincode::default());
+    let transport = transport.with(|msg: tarpc::ClientMessage<ClientContext, _>| future::ok(msg.map_context(|ctx| ctx.shared_context)));
     let answer = PingServiceClient::new(Default::default(), transport)
         .spawn()
         .ping(&mut ClientContext::current())

tarpc/examples/tracing.rs

@@ -35,6 +35,7 @@ use tarpc::{
 };
 use tokio::net::TcpStream;
 use tracing_subscriber::prelude::*;
+use tarpc::context::{ClientContext, ServerContext, SharedContext};
 
 pub mod add {
     #[tarpc::service]
@@ -124,7 +125,7 @@ where
 }
 
 fn make_stub<Req, Resp, const N: usize>(
-    backends: [impl Transport<ClientMessage<Arc<Req>>, Response<Resp>> + Send + Sync + 'static; N],
+    backends: [impl Transport<ClientMessage<ClientContext, Arc<Req>>, Response<Resp>> + Send + Sync + 'static; N],
 ) -> retry::Retry<
     impl Fn(&Result<Resp, RpcError>, u32) -> bool + Clone,
     load_balance::RoundRobin<client::Channel<Arc<Req>, Resp>>,
@@ -173,23 +174,28 @@ async fn main() -> anyhow::Result<()> {
         .serving(AddServer.serve());
     let add_server = add_listener1
         .chain(add_listener2)
+        .map(|t| t.map_ok(|msg: ClientMessage<SharedContext, _>| msg.map_context(|ctx| ServerContext::new(ctx))))
         .map(BaseChannel::with_defaults);
     tokio::spawn(spawn_incoming(add_server.execute(server)));
 
+    let map_context = |msg: ClientMessage<ClientContext, _>| future::ok(msg.map_context(|ctx| ctx.shared_context));
+
     let add_client = add::AddClient::from(make_stub([
-        tarpc::serde_transport::tcp::connect(addr1, Json::default).await?,
-        tarpc::serde_transport::tcp::connect(addr2, Json::default).await?,
+        tarpc::serde_transport::tcp::connect(addr1, Json::default).await?.with(map_context),
+        tarpc::serde_transport::tcp::connect(addr2, Json::default).await?.with(map_context),
     ]));
 
     let double_listener = tarpc::serde_transport::tcp::listen("localhost:0", Json::default)
         .await?
-        .filter_map(|r| future::ready(r.ok()));
-    let addr = double_listener.get_ref().local_addr();
+        .filter_map(|r| future::ready(r.ok()))
+        .map(|t| t.map_ok(|msg: ClientMessage<SharedContext, _>| msg.map_context(|ctx| ServerContext::new(ctx))));
+    let addr = double_listener.get_ref().get_ref().local_addr();
     let double_server = double_listener.map(BaseChannel::with_defaults).take(1);
     let server = DoubleServer { add_client }.serve();
     tokio::spawn(spawn_incoming(double_server.execute(server)));
 
     let to_double_server = tarpc::serde_transport::tcp::connect(addr, Json::default).await?;
+    let to_double_server = to_double_server.with(|msg: ClientMessage<ClientContext, _>| future::ok(msg.map_context(|ctx| ctx.shared_context)));
     let double_client =
         double::DoubleClient::new(client::Config::default(), to_double_server).spawn();