Skip to main content

wrpc_transport/frame/tcp/
tokio.rs

1//! wRPC TCP transport using [tokio]
2
3use bytes::Bytes;
4use tokio::io::BufReader;
5use tokio::net::{TcpStream, ToSocketAddrs};
6use tracing::instrument;
7
8use crate::Invoke;
9use crate::frame::{Incoming, Outgoing, invoke};
10
11/// [Invoke] implementation of a TCP transport using [tokio]
12#[derive(Clone, Debug)]
13pub struct Client<T> {
14    addr: T,
15    nodelay: Option<bool>,
16    zero_linger: bool,
17}
18
19impl<T> Client<T> {
20    /// Sets whether `TCP_NODELAY` is enabled on connections established by this client,
21    /// `None` leaves the system default unchanged.
22    /// It is enabled by default, since invocations consist of small writes that are each
23    /// waited on by the peer.
24    #[must_use]
25    pub fn nodelay(self, nodelay: Option<bool>) -> Self {
26        Self { nodelay, ..self }
27    }
28
29    /// Sets `SO_LINGER` to zero on connections established by this client, so that closing
30    /// them aborts the connection instead of leaving it in `TIME_WAIT`.
31    /// Aborting the connection discards any data not yet sent to the peer, so this should only be
32    /// used if the peer does not depend on receiving all data, for example in benchmarks.
33    /// It is not set by default.
34    #[must_use]
35    pub fn set_zero_linger(self) -> Self {
36        Self {
37            zero_linger: true,
38            ..self
39        }
40    }
41}
42
43impl<T> From<T> for Client<T>
44where
45    T: ToSocketAddrs + Clone,
46{
47    fn from(addr: T) -> Self {
48        Self {
49            addr,
50            nodelay: Some(true),
51            zero_linger: false,
52        }
53    }
54}
55
56impl<T> Invoke for Client<T>
57where
58    T: ToSocketAddrs + Clone + Send + Sync,
59{
60    type Context = ();
61
62    #[instrument(level = "trace", skip(self, paths, params), fields(params = format!("{params:02x?}")))]
63    async fn invoke<P>(
64        &self,
65        (): Self::Context,
66        instance: &str,
67        func: &str,
68        params: Bytes,
69        paths: impl AsRef<[P]> + Send,
70    ) -> anyhow::Result<(Outgoing, Incoming)>
71    where
72        P: AsRef<[Option<usize>]> + Send + Sync,
73    {
74        let stream = TcpStream::connect(self.addr.clone()).await?;
75        if let Some(nodelay) = self.nodelay {
76            stream.set_nodelay(nodelay)?;
77        }
78        if self.zero_linger {
79            stream.set_zero_linger()?;
80        }
81        let (rx, tx) = stream.into_split();
82        invoke(tx, BufReader::new(rx), instance, func, params, paths).await
83    }
84}