workflow_rpc/client/protocol/
borsh.rs

1use super::{Pending, PendingMap, ProtocolHandler};
2use crate::client::Interface;
3pub use crate::client::error::Error;
4pub use crate::client::result::Result;
5use crate::imports::*;
6use crate::messages::borsh::*;
7use core::marker::PhantomData;
8
9pub type BorshResponseFn =
10    Arc<Box<dyn Fn(Result<&[u8]>, Option<&Duration>) -> Result<()> + Sync + Send>>;
11
12/// Borsh RPC message handler and dispatcher
13pub struct BorshProtocol<Ops, Id>
14where
15    Ops: OpsT,
16    Id: IdT,
17{
18    ws: Arc<WebSocket>,
19    pending: PendingMap<Id, BorshResponseFn>,
20    interface: Option<Arc<Interface<Ops>>>,
21    ops: PhantomData<Ops>,
22    id: PhantomData<Id>,
23}
24
25impl<Ops, Id> BorshProtocol<Ops, Id>
26where
27    Ops: OpsT,
28    Id: IdT,
29{
30    fn new(ws: Arc<WebSocket>, interface: Option<Arc<Interface<Ops>>>) -> Self {
31        BorshProtocol {
32            ws,
33            pending: Arc::new(Mutex::new(AHashMap::new())),
34            interface,
35            ops: PhantomData,
36            id: PhantomData,
37        }
38    }
39}
40
41type MessageInfo<'l, Ops, Id> = (Option<Id>, Option<Ops>, Result<&'l [u8]>);
42
43impl<Ops, Id> BorshProtocol<Ops, Id>
44where
45    Id: IdT,
46    Ops: OpsT,
47{
48    fn decode<'l>(&self, server_message: &'l [u8]) -> ServerResult<MessageInfo<'l, Ops, Id>> {
49        match BorshServerMessage::try_from(server_message) {
50            Ok(msg) => {
51                let header = msg.header;
52                match header.kind {
53                    ServerMessageKind::Success => {
54                        Ok((header.id, header.op, Ok(msg.payload)))
55                        // Ok((Some(header.id), header.op.clone(), Ok(msg.data)))
56                    }
57                    ServerMessageKind::Error => {
58                        if let Ok(err) = ServerError::try_from_slice(msg.payload) {
59                            Ok((header.id, None, Err(Error::RpcCall(err))))
60                        } else {
61                            Ok((header.id, None, Err(Error::ErrorDeserializingResponseData)))
62                        }
63                    }
64                    ServerMessageKind::Notification => Ok((None, header.op, Ok(msg.payload))),
65                }
66            }
67            Err(err) => Err(ServerError::RespDeserialize(err.to_string())),
68        }
69    }
70
71    /// Send a Borsh-encoded request for `op` and await the deserialized
72    /// response, returning an error on transport or (de)serialization failure
73    /// or if the server responds with an error.
74    pub async fn request<Req, Resp>(&self, op: Ops, req: Req) -> Result<Resp>
75    where
76        Req: MsgT,
77        Resp: MsgT,
78    {
79        let payload = borsh::to_vec(&req).map_err(|_| Error::BorshSerialize)?;
80
81        let id = Id::generate();
82        let (sender, receiver) = oneshot();
83
84        {
85            let mut pending = self.pending.lock().unwrap();
86            pending.insert(
87                id.clone(),
88                Pending::new(Arc::new(Box::new(move |result, _duration| {
89                    sender.try_send(result.map(|data| data.to_vec()))?;
90                    Ok(())
91                }))),
92            );
93        }
94
95        // TODO - post error into sender if ws.send() fails
96        self.ws
97            .post(to_ws_msg(BorshReqHeader::new(Some(id), op), &payload))
98            .await?;
99
100        let data = receiver.recv().await??;
101        let resp = ServerResult::<Resp>::try_from_slice(data.as_ref())
102            .map_err(|e| Error::BorshDeserialize(e.to_string()))?;
103
104        Ok(resp?)
105    }
106
107    /// Send a Borsh-encoded notification for `op` without awaiting a response.
108    pub async fn notify<Msg>(&self, op: Ops, payload: Msg) -> Result<()>
109    where
110        Msg: BorshSerialize + Send + Sync + 'static,
111    {
112        let payload = borsh::to_vec(&payload).map_err(|_| Error::BorshSerialize)?;
113        self.ws
114            .post(to_ws_msg(
115                BorshReqHeader::<Ops, Id>::new(None, op),
116                &payload,
117            ))
118            .await?;
119        Ok(())
120    }
121
122    async fn handle_notification(&self, op: &Ops, payload: &[u8]) -> Result<()> {
123        if let Some(interface) = &self.interface {
124            interface
125                .call_notification_with_borsh(op, payload)
126                .await
127                .unwrap_or_else(|err| log_trace!("error handling server notification {}", err));
128        } else {
129            log_trace!("unable to handle server notification - interface is not initialized");
130        }
131
132        Ok(())
133    }
134}
135
136#[async_trait]
137impl<Ops, Id> ProtocolHandler<Ops> for BorshProtocol<Ops, Id>
138where
139    Id: IdT,
140    Ops: OpsT,
141{
142    fn new(ws: Arc<WebSocket>, interface: Option<Arc<Interface<Ops>>>) -> Self
143    where
144        Self: Sized,
145    {
146        BorshProtocol::new(ws, interface)
147    }
148
149    async fn handle_timeout(&self, timeout: Duration) {
150        self.pending.lock().unwrap().retain(|_, pending| {
151            if pending.timestamp.elapsed() > timeout {
152                (pending.callback)(Err(Error::Timeout), None).unwrap_or_else(|err| {
153                    log_trace!("Error in RPC callback during timeout: `{err}`")
154                });
155                false
156            } else {
157                true
158            }
159        });
160    }
161
162    async fn handle_disconnect(&self) -> Result<()> {
163        self.pending.lock().unwrap().retain(|_, pending| {
164            (pending.callback)(Err(Error::Disconnect), None)
165                .unwrap_or_else(|err| log_trace!("Error in RPC callback during timeout: `{err}`"));
166            false
167        });
168
169        Ok(())
170    }
171
172    async fn handle_message(&self, message: WebSocketMessage) -> Result<()> {
173        if let WebSocketMessage::Binary(server_message) = message {
174            let (id, op, result) = self.decode(server_message.as_slice())?;
175            if let Some(id) = id {
176                match self.pending.lock().unwrap().remove(&id) {
177                    Some(pending) => (pending.callback)(result, Some(&pending.timestamp.elapsed())),
178                    _ => Err(Error::ResponseHandler(format!("{id:?}"))),
179                }
180            } else if let Some(op) = op {
181                match result {
182                    Ok(data) => self.handle_notification(&op, data).await,
183                    _ => Ok(()),
184                }
185            } else {
186                Err(Error::NotificationMethod)
187            }
188        } else {
189            return Err(Error::WebSocketMessageType);
190        }
191    }
192}
workflow_rpc/client/protocol/borsh.rs

workflow_rpc/client/protocol/
borsh.rs
