Skip to main content

workflow_rpc/client/protocol/
serde_json.rs

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