Skip to main content

wrpc_transport/frame/conn/
client.rs

1use core::marker::PhantomData;
2
3use anyhow::Context as _;
4use bytes::{Buf as _, BufMut as _, BytesMut};
5use tokio::io::{AsyncBufRead, AsyncWrite, AsyncWriteExt as _};
6use tokio_util::codec::Encoder;
7use tracing::{instrument, trace};
8use wasm_tokio::{CoreNameEncoder, Leb128Encoder};
9
10use crate::frame::conn::{Incoming, Outgoing};
11use crate::frame::{ConnHandler, PROTOCOL};
12
13/// Encodes the header of an invocation, which is followed by `params_len` bytes of
14/// parameters, to a `BytesMut`
15///
16/// This is low-level API, most users should use [`invoke`].
17#[instrument(level = "trace", skip_all)]
18pub fn encode_invocation(
19    buf: &mut BytesMut,
20    instance: &str,
21    func: &str,
22    params_len: usize,
23) -> std::io::Result<()> {
24    let params_len = u32::try_from(params_len)
25        .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
26    buf.reserve(
27        17_usize // len(PROTOCOL) + len(instance) + len(func) + len([]) + len(params)
28            .saturating_add(instance.len())
29            .saturating_add(func.len()),
30    );
31    buf.put_u8(PROTOCOL);
32    CoreNameEncoder.encode(instance, buf)?;
33    CoreNameEncoder.encode(func, buf)?;
34    buf.put_u8(0);
35    Leb128Encoder.encode(params_len, buf)?;
36    Ok(())
37}
38
39/// Defines invocation behavior
40#[derive(Clone)]
41pub struct InvokeBuilder<H = ()>(PhantomData<H>)
42where
43    H: ?Sized;
44
45impl<H> InvokeBuilder<H> {
46    /// Invoke function `func` on instance `instance`
47    #[instrument(level = "trace", skip_all)]
48    pub async fn invoke<P, I, O>(
49        self,
50        mut tx: O,
51        rx: I,
52        instance: &str,
53        func: &str,
54        params: impl AsRef<[u8]>,
55        paths: impl AsRef<[P]> + Send,
56    ) -> anyhow::Result<(Outgoing, Incoming)>
57    where
58        P: AsRef<[Option<usize>]> + Send + Sync,
59        I: AsyncBufRead + Unpin + Send + 'static,
60        O: AsyncWrite + Unpin + Send + 'static,
61        H: ConnHandler<I, O>,
62    {
63        let params = params.as_ref();
64        let mut buf = BytesMut::default();
65        encode_invocation(&mut buf, instance, func, params.len())
66            .context("failed to encode invocation")?;
67        trace!(?buf, "writing invocation");
68        tx.write_all_buf(&mut buf.chain(params))
69            .await
70            .context("failed to initialize connection")?;
71        tx.flush().await.context("failed to flush invocation")?;
72
73        let tx = Outgoing::new(tx, |tx, res| H::on_egress(tx, res));
74        let rx = Incoming::new(rx, paths.as_ref(), H::DRAIN, |rx, res| {
75            H::on_ingress(rx, res)
76        });
77        Ok((tx, rx))
78    }
79}
80
81impl<H> Default for InvokeBuilder<H> {
82    fn default() -> Self {
83        Self(PhantomData)
84    }
85}
86
87/// Invoke function `func` on instance `instance`
88#[instrument(level = "trace", skip_all)]
89pub async fn invoke<P, I, O>(
90    tx: O,
91    rx: I,
92    instance: &str,
93    func: &str,
94    params: impl AsRef<[u8]>,
95    paths: impl AsRef<[P]> + Send,
96) -> anyhow::Result<(Outgoing, Incoming)>
97where
98    P: AsRef<[Option<usize>]> + Send + Sync,
99    I: AsyncBufRead + Unpin + Send + 'static,
100    O: AsyncWrite + Unpin + Send + 'static,
101{
102    InvokeBuilder::<()>::default()
103        .invoke(tx, rx, instance, func, params, paths)
104        .await
105}