anvil_server/
pubsub.rs

1use crate::{RpcHandler, error::RequestError, handler::handle_request};
2use anvil_rpc::{
3    error::RpcError,
4    request::Request,
5    response::{Response, ResponseResult},
6};
7
8use futures::{FutureExt, Sink, SinkExt, Stream, StreamExt};
9use parking_lot::Mutex;
10use serde::de::DeserializeOwned;
11use std::{
12    collections::VecDeque,
13    fmt,
14    hash::Hash,
15    pin::Pin,
16    sync::Arc,
17    task::{Context, Poll},
18};
19
20/// The general purpose trait for handling RPC requests and subscriptions
21pub trait PubSubRpcHandler: Clone + Send + Sync + Unpin + 'static {
22    /// The request type to expect
23    type Request: DeserializeOwned + Send + Sync + fmt::Debug;
24    /// The identifier to use for subscriptions
25    type SubscriptionId: Hash + PartialEq + Eq + Send + Sync + fmt::Debug;
26    /// The subscription type this handle may create
27    type Subscription: Stream<Item = serde_json::Value> + Send + Sync + Unpin;
28
29    /// Invoked when the request was received
30    fn on_request(
31        &self,
32        request: Self::Request,
33        cx: PubSubContext<Self>,
34    ) -> impl Future<Output = ResponseResult> + Send;
35}
36
37type Subscriptions<SubscriptionId, Subscription> = Arc<Mutex<Vec<(SubscriptionId, Subscription)>>>;
38
39/// Contains additional context and tracks subscriptions
40pub struct PubSubContext<Handler: PubSubRpcHandler> {
41    /// all active subscriptions `id -> Stream`
42    subscriptions: Subscriptions<Handler::SubscriptionId, Handler::Subscription>,
43}
44
45impl<Handler: PubSubRpcHandler> PubSubContext<Handler> {
46    /// Adds new active subscription
47    ///
48    /// Returns the previous subscription, if any
49    pub fn add_subscription(
50        &self,
51        id: Handler::SubscriptionId,
52        subscription: Handler::Subscription,
53    ) -> Option<Handler::Subscription> {
54        let mut subscriptions = self.subscriptions.lock();
55        let mut removed = None;
56        if let Some(idx) = subscriptions.iter().position(|(i, _)| id == *i) {
57            trace!(target: "rpc", ?id,  "removed subscription");
58            removed = Some(subscriptions.swap_remove(idx).1);
59        }
60        trace!(target: "rpc", ?id,  "added subscription");
61        subscriptions.push((id, subscription));
62        removed
63    }
64
65    /// Removes an existing subscription
66    pub fn remove_subscription(
67        &self,
68        id: &Handler::SubscriptionId,
69    ) -> Option<Handler::Subscription> {
70        let mut subscriptions = self.subscriptions.lock();
71        if let Some(idx) = subscriptions.iter().position(|(i, _)| id == i) {
72            trace!(target: "rpc", ?id,  "removed subscription");
73            return Some(subscriptions.swap_remove(idx).1);
74        }
75        None
76    }
77}
78
79impl<Handler: PubSubRpcHandler> Clone for PubSubContext<Handler> {
80    fn clone(&self) -> Self {
81        Self { subscriptions: Arc::clone(&self.subscriptions) }
82    }
83}
84
85impl<Handler: PubSubRpcHandler> Default for PubSubContext<Handler> {
86    fn default() -> Self {
87        Self { subscriptions: Arc::new(Mutex::new(Vec::new())) }
88    }
89}
90
91/// A compatibility helper type to use common `RpcHandler` functions
92struct ContextAwareHandler<Handler: PubSubRpcHandler> {
93    handler: Handler,
94    context: PubSubContext<Handler>,
95}
96
97impl<Handler: PubSubRpcHandler> Clone for ContextAwareHandler<Handler> {
98    fn clone(&self) -> Self {
99        Self { handler: self.handler.clone(), context: self.context.clone() }
100    }
101}
102
103impl<Handler: PubSubRpcHandler> RpcHandler for ContextAwareHandler<Handler> {
104    type Request = Handler::Request;
105
106    fn on_request(&self, request: Self::Request) -> impl Future<Output = ResponseResult> + Send {
107        self.handler.on_request(request, self.context.clone())
108    }
109}
110
111/// Represents a connection to a client via websocket
112///
113/// Contains the state for the entire connection
114pub struct PubSubConnection<Handler: PubSubRpcHandler, Connection> {
115    /// the handler for the websocket connection
116    handler: Handler,
117    /// contains all the subscription related context
118    context: PubSubContext<Handler>,
119    /// The established connection
120    connection: Connection,
121    /// currently in progress requests
122    processing: Vec<Pin<Box<dyn Future<Output = Option<Response>> + Send>>>,
123    /// pending messages to send
124    pending: VecDeque<String>,
125}
126
127impl<Handler: PubSubRpcHandler, Connection> PubSubConnection<Handler, Connection> {
128    pub fn new(connection: Connection, handler: Handler) -> Self {
129        Self {
130            connection,
131            handler,
132            context: Default::default(),
133            pending: Default::default(),
134            processing: Default::default(),
135        }
136    }
137
138    fn process_request(&mut self, req: serde_json::Result<Request>) {
139        let handler =
140            ContextAwareHandler { handler: self.handler.clone(), context: self.context.clone() };
141        self.processing.push(Box::pin(async move {
142            match req {
143                Ok(req) => handle_request(req, handler).await,
144                Err(err) => {
145                    error!(target: "rpc", ?err, "invalid request");
146                    let err = if err.is_syntax() || err.is_eof() {
147                        RpcError::parse_error()
148                    } else {
149                        RpcError::invalid_request()
150                    };
151                    Some(Response::error(err))
152                }
153            }
154        }));
155    }
156}
157
158impl<Handler, Connection> Future for PubSubConnection<Handler, Connection>
159where
160    Handler: PubSubRpcHandler,
161    Connection: Sink<String> + Stream<Item = Result<Option<Request>, RequestError>> + Unpin,
162    <Connection as Sink<String>>::Error: fmt::Debug,
163{
164    type Output = ();
165
166    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
167        let pin = self.get_mut();
168        loop {
169            // drive the websocket
170            while matches!(pin.connection.poll_ready_unpin(cx), Poll::Ready(Ok(()))) {
171                // only start sending if socket is ready
172                if let Some(msg) = pin.pending.pop_front() {
173                    if let Err(err) = pin.connection.start_send_unpin(msg) {
174                        error!(target: "rpc", ?err, "Failed to send message");
175                    }
176                } else {
177                    break;
178                }
179            }
180
181            // Ensure any pending messages are flushed
182            // this needs to be called manually for tungsenite websocket: <https://github.com/foundry-rs/foundry/issues/6345>
183            if let Poll::Ready(Err(err)) = pin.connection.poll_flush_unpin(cx) {
184                trace!(target: "rpc", ?err, "websocket err");
185                // close the connection
186                return Poll::Ready(());
187            }
188
189            loop {
190                match pin.connection.poll_next_unpin(cx) {
191                    Poll::Ready(Some(req)) => match req {
192                        Ok(Some(req)) => {
193                            pin.process_request(Ok(req));
194                        }
195                        Err(err) => match err {
196                            RequestError::Axum(err) => {
197                                trace!(target: "rpc", ?err, "client disconnected");
198                                return Poll::Ready(());
199                            }
200                            RequestError::Io(err) => {
201                                trace!(target: "rpc", ?err, "client disconnected");
202                                return Poll::Ready(());
203                            }
204                            RequestError::Serde(err) => {
205                                pin.process_request(Err(err));
206                            }
207                            RequestError::Disconnect => {
208                                trace!(target: "rpc", "client disconnected");
209                                return Poll::Ready(());
210                            }
211                        },
212                        _ => {}
213                    },
214                    Poll::Ready(None) => {
215                        trace!(target: "rpc", "socket connection finished");
216                        return Poll::Ready(());
217                    }
218                    Poll::Pending => break,
219                }
220            }
221
222            let mut progress = false;
223            for n in (0..pin.processing.len()).rev() {
224                let mut req = pin.processing.swap_remove(n);
225                #[allow(clippy::collapsible_match)]
226                match req.poll_unpin(cx) {
227                    Poll::Ready(Some(resp)) => {
228                        if let Ok(text) = serde_json::to_string(&resp) {
229                            pin.pending.push_back(text);
230                            progress = true;
231                        }
232                    }
233                    Poll::Ready(None) => {}
234                    Poll::Pending => pin.processing.push(req),
235                }
236            }
237
238            {
239                // process subscription events
240                let mut subscriptions = pin.context.subscriptions.lock();
241                'outer: for n in (0..subscriptions.len()).rev() {
242                    let (id, mut sub) = subscriptions.swap_remove(n);
243                    'inner: loop {
244                        #[allow(clippy::collapsible_match)]
245                        match sub.poll_next_unpin(cx) {
246                            Poll::Ready(Some(res)) => {
247                                if let Ok(text) = serde_json::to_string(&res) {
248                                    pin.pending.push_back(text);
249                                    progress = true;
250                                }
251                            }
252                            Poll::Ready(None) => continue 'outer,
253                            Poll::Pending => break 'inner,
254                        }
255                    }
256
257                    subscriptions.push((id, sub));
258                }
259            }
260
261            if !progress {
262                return Poll::Pending;
263            }
264        }
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271    use anvil_rpc::{
272        request::{RequestParams, RpcCall, RpcNotification, Version},
273        response::RpcResponse,
274    };
275    use std::{
276        pin::pin,
277        sync::atomic::{AtomicUsize, Ordering},
278        task::Waker,
279    };
280
281    #[derive(Clone, Default)]
282    struct TestHandler {
283        requests: Arc<AtomicUsize>,
284    }
285
286    impl PubSubRpcHandler for TestHandler {
287        type Request = serde_json::Value;
288        type SubscriptionId = u64;
289        type Subscription = futures::stream::Empty<serde_json::Value>;
290
291        async fn on_request(
292            &self,
293            _request: Self::Request,
294            _cx: PubSubContext<Self>,
295        ) -> ResponseResult {
296            self.requests.fetch_add(1, Ordering::Relaxed);
297            ResponseResult::success(serde_json::Value::Null)
298        }
299    }
300
301    fn notification() -> RpcCall {
302        RpcCall::Notification(RpcNotification {
303            jsonrpc: Some(Version::V2),
304            method: "eth_subscribe".to_owned(),
305            params: RequestParams::None,
306        })
307    }
308
309    fn run_ready<F: Future>(future: F) -> F::Output {
310        let waker = Waker::noop();
311        let mut cx = Context::from_waker(waker);
312        let mut future = pin!(future);
313        match future.as_mut().poll(&mut cx) {
314            Poll::Ready(output) => output,
315            Poll::Pending => panic!("future unexpectedly pending"),
316        }
317    }
318
319    #[test]
320    fn process_request_keeps_empty_batch_invalid() {
321        let mut connection = PubSubConnection::new((), TestHandler::default());
322        connection.process_request(Ok(Request::Batch(vec![])));
323
324        let response = run_ready(connection.processing.pop().unwrap());
325        assert_eq!(
326            response,
327            Some(Response::Single(RpcResponse::from(RpcError::invalid_request())))
328        );
329    }
330
331    #[test]
332    fn process_request_returns_parse_error_for_malformed_json() {
333        let mut connection = PubSubConnection::new((), TestHandler::default());
334        connection.process_request(serde_json::from_str("{"));
335
336        let response = run_ready(connection.processing.pop().unwrap());
337        assert_eq!(response, Some(Response::error(RpcError::parse_error())));
338    }
339
340    #[test]
341    fn process_request_executes_notification_without_response() {
342        let handler = TestHandler::default();
343        let mut connection = PubSubConnection::new((), handler.clone());
344        connection.process_request(Ok(Request::Batch(vec![notification()])));
345
346        let response = run_ready(connection.processing.pop().unwrap());
347        assert_eq!(response, None);
348        assert_eq!(handler.requests.load(Ordering::Relaxed), 1);
349    }
350}
anvil_server/pubsub.rs

anvil_server/
pubsub.rs
