anvil_server/
pubsub.rs

1use crate::{RpcHandler, error::RequestError, handler::handle_request};
2use anvil_rpc::{
3    error::RpcError,
4    request::{Id, 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    /// Invoked when a deserialized request returns an error from [`Self::on_request`].
37    /// Parameters retain their original shape with fork URL credentials redacted.
38    /// Deserialization errors are excluded because their messages can contain request parameters.
39    fn on_request_error(
40        &self,
41        _method: &str,
42        _id: &Id,
43        _params: &serde_json::Value,
44        _error: &RpcError,
45    ) {
46    }
47
48    /// Whether an ended subscription requires closing the client connection.
49    fn close_on_subscription_end(&self, _subscription: &Self::Subscription) -> bool {
50        false
51    }
52}
53
54type Subscriptions<SubscriptionId, Subscription> =
55    Arc<Mutex<VecDeque<(SubscriptionId, Subscription)>>>;
56
57/// Contains additional context and tracks subscriptions
58pub struct PubSubContext<Handler: PubSubRpcHandler> {
59    /// all active subscriptions `id -> Stream`
60    subscriptions: Subscriptions<Handler::SubscriptionId, Handler::Subscription>,
61}
62
63impl<Handler: PubSubRpcHandler> PubSubContext<Handler> {
64    /// Adds new active subscription
65    ///
66    /// Returns the previous subscription, if any
67    pub fn add_subscription(
68        &self,
69        id: Handler::SubscriptionId,
70        subscription: Handler::Subscription,
71    ) -> Option<Handler::Subscription> {
72        let mut subscriptions = self.subscriptions.lock();
73        let mut removed = None;
74        if let Some(idx) = subscriptions.iter().position(|(i, _)| id == *i) {
75            trace!(target: "rpc", ?id,  "removed subscription");
76            removed = Some(subscriptions.swap_remove_back(idx).unwrap().1);
77        }
78        trace!(target: "rpc", ?id,  "added subscription");
79        subscriptions.push_back((id, subscription));
80        removed
81    }
82
83    /// Removes an existing subscription
84    pub fn remove_subscription(
85        &self,
86        id: &Handler::SubscriptionId,
87    ) -> Option<Handler::Subscription> {
88        let mut subscriptions = self.subscriptions.lock();
89        if let Some(idx) = subscriptions.iter().position(|(i, _)| id == i) {
90            trace!(target: "rpc", ?id,  "removed subscription");
91            return Some(subscriptions.swap_remove_back(idx).unwrap().1);
92        }
93        None
94    }
95}
96
97impl<Handler: PubSubRpcHandler> Clone for PubSubContext<Handler> {
98    fn clone(&self) -> Self {
99        Self { subscriptions: Arc::clone(&self.subscriptions) }
100    }
101}
102
103impl<Handler: PubSubRpcHandler> Default for PubSubContext<Handler> {
104    fn default() -> Self {
105        Self { subscriptions: Arc::new(Mutex::new(VecDeque::new())) }
106    }
107}
108
109/// A compatibility helper type to use common `RpcHandler` functions
110struct ContextAwareHandler<Handler: PubSubRpcHandler> {
111    handler: Handler,
112    context: PubSubContext<Handler>,
113}
114
115impl<Handler: PubSubRpcHandler> Clone for ContextAwareHandler<Handler> {
116    fn clone(&self) -> Self {
117        Self { handler: self.handler.clone(), context: self.context.clone() }
118    }
119}
120
121impl<Handler: PubSubRpcHandler> RpcHandler for ContextAwareHandler<Handler> {
122    type Request = Handler::Request;
123
124    fn on_request(&self, request: Self::Request) -> impl Future<Output = ResponseResult> + Send {
125        self.handler.on_request(request, self.context.clone())
126    }
127
128    fn on_request_error(
129        &self,
130        method: &str,
131        id: &Id,
132        params: &serde_json::Value,
133        error: &RpcError,
134    ) {
135        self.handler.on_request_error(method, id, params, error);
136    }
137}
138
139/// Represents a connection to a client via websocket
140///
141/// Contains the state for the entire connection
142pub struct PubSubConnection<Handler: PubSubRpcHandler, Connection> {
143    /// the handler for the websocket connection
144    handler: Handler,
145    /// contains all the subscription related context
146    context: PubSubContext<Handler>,
147    /// The established connection
148    connection: Connection,
149    /// currently in progress requests
150    processing: Vec<Pin<Box<dyn Future<Output = Option<Response>> + Send>>>,
151    /// pending messages to send
152    pending: PendingMessages,
153}
154
155impl<Handler: PubSubRpcHandler, Connection> PubSubConnection<Handler, Connection> {
156    pub fn new(connection: Connection, handler: Handler) -> Self {
157        Self {
158            connection,
159            handler,
160            context: Default::default(),
161            pending: Default::default(),
162            processing: Default::default(),
163        }
164    }
165
166    fn process_request(&mut self, req: serde_json::Result<Request>) {
167        let handler =
168            ContextAwareHandler { handler: self.handler.clone(), context: self.context.clone() };
169        self.processing.push(Box::pin(async move {
170            match req {
171                Ok(req) => handle_request(req, handler).await,
172                Err(err) => {
173                    error!(target: "rpc", ?err, "invalid request");
174                    let err = if err.is_syntax() || err.is_eof() {
175                        RpcError::parse_error()
176                    } else {
177                        RpcError::invalid_request()
178                    };
179                    Some(Response::error(err))
180                }
181            }
182        }));
183    }
184}
185
186impl<Handler, Connection> Future for PubSubConnection<Handler, Connection>
187where
188    Handler: PubSubRpcHandler,
189    Connection: Sink<String> + Stream<Item = Result<Option<Request>, RequestError>> + Unpin,
190    <Connection as Sink<String>>::Error: fmt::Debug,
191{
192    type Output = ();
193
194    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
195        let pin = self.get_mut();
196        // Cooperate with other connections even when subscription streams stay ready.
197        let mut remaining = 256;
198        loop {
199            // drive the websocket
200            let mut writable = true;
201            loop {
202                match pin.connection.poll_ready_unpin(cx) {
203                    Poll::Ready(Ok(())) => {}
204                    Poll::Ready(Err(err)) => {
205                        trace!(target: "rpc", ?err, "connection err");
206                        return Poll::Ready(());
207                    }
208                    Poll::Pending => {
209                        writable = false;
210                        break;
211                    }
212                }
213                // only start sending if socket is ready
214                if let Some(msg) = pin.pending.pop_front() {
215                    if let Err(err) = pin.connection.start_send_unpin(msg) {
216                        error!(target: "rpc", ?err, "Failed to send message");
217                        return Poll::Ready(());
218                    }
219                } else {
220                    break;
221                }
222            }
223
224            // Ensure any pending messages are flushed
225            // this needs to be called manually for tungsenite websocket: <https://github.com/foundry-rs/foundry/issues/6345>
226            if let Poll::Ready(Err(err)) = pin.connection.poll_flush_unpin(cx) {
227                trace!(target: "rpc", ?err, "websocket err");
228                // close the connection
229                return Poll::Ready(());
230            }
231
232            loop {
233                match pin.connection.poll_next_unpin(cx) {
234                    Poll::Ready(Some(req)) => match req {
235                        Ok(Some(req)) => {
236                            pin.process_request(Ok(req));
237                        }
238                        Err(err) => match err {
239                            RequestError::Axum(err) => {
240                                trace!(target: "rpc", ?err, "client disconnected");
241                                return Poll::Ready(());
242                            }
243                            RequestError::Io(err) => {
244                                trace!(target: "rpc", ?err, "client disconnected");
245                                return Poll::Ready(());
246                            }
247                            RequestError::Serde(err) => {
248                                pin.process_request(Err(err));
249                            }
250                            RequestError::Disconnect => {
251                                trace!(target: "rpc", "client disconnected");
252                                return Poll::Ready(());
253                            }
254                        },
255                        _ => {}
256                    },
257                    Poll::Ready(None) => {
258                        trace!(target: "rpc", "socket connection finished");
259                        return Poll::Ready(());
260                    }
261                    Poll::Pending => break,
262                }
263            }
264
265            let mut progress = false;
266            for n in (0..pin.processing.len()).rev() {
267                let mut req = pin.processing.swap_remove(n);
268                #[allow(clippy::collapsible_match)]
269                match req.poll_unpin(cx) {
270                    Poll::Ready(Some(resp)) => {
271                        if let Ok(text) = serde_json::to_string(&resp) {
272                            if !pin.pending.send(&mut pin.connection, cx, text, &mut writable) {
273                                return Poll::Ready(());
274                            }
275                            progress = true;
276                        }
277                    }
278                    Poll::Ready(None) => {}
279                    Poll::Pending => pin.processing.push(req),
280                }
281            }
282
283            {
284                // Materialize events before upstream queues overflow or their backing data expires.
285                let mut subscriptions = pin.context.subscriptions.lock();
286                for _ in 0..subscriptions.len() {
287                    let (id, mut sub) = subscriptions.pop_back().unwrap();
288                    #[allow(clippy::collapsible_match)]
289                    match sub.poll_next_unpin(cx) {
290                        Poll::Ready(Some(res)) => {
291                            if let Ok(text) = serde_json::to_string(&res) {
292                                if !pin.pending.send(&mut pin.connection, cx, text, &mut writable) {
293                                    return Poll::Ready(());
294                                }
295                                progress = true;
296                            }
297                            remaining -= 1;
298                        }
299                        Poll::Ready(None) => {
300                            if pin.handler.close_on_subscription_end(&sub) {
301                                warn!(target: "rpc", ?id, "Closing connection after subscription ended");
302                                return Poll::Ready(());
303                            }
304                            continue;
305                        }
306                        Poll::Pending => {}
307                    }
308
309                    subscriptions.push_front((id, sub));
310                    if remaining == 0 {
311                        cx.waker().wake_by_ref();
312                        return Poll::Pending;
313                    }
314                }
315            }
316
317            if !progress {
318                return Poll::Pending;
319            }
320        }
321    }
322}
323
324/// Bounded serialized backlog shared by responses and subscription notifications.
325#[derive(Default)]
326struct PendingMessages {
327    messages: VecDeque<String>,
328    bytes: usize,
329}
330
331impl PendingMessages {
332    fn pop_front(&mut self) -> Option<String> {
333        let message = self.messages.pop_front()?;
334        self.bytes -= message.len();
335        Some(message)
336    }
337
338    /// Sends directly while writable, otherwise queues within the connection's backlog limits.
339    /// Returns false when the connection must terminate.
340    fn send<Connection>(
341        &mut self,
342        connection: &mut Connection,
343        cx: &mut Context<'_>,
344        message: String,
345        writable: &mut bool,
346    ) -> bool
347    where
348        Connection: Sink<String> + Unpin,
349        Connection::Error: fmt::Debug,
350    {
351        if *writable {
352            debug_assert!(self.messages.is_empty());
353            match connection.poll_ready_unpin(cx) {
354                Poll::Ready(Ok(())) => {
355                    if let Err(err) = connection.start_send_unpin(message) {
356                        error!(target: "rpc", ?err, "Failed to send message");
357                        return false;
358                    }
359                    return true;
360                }
361                Poll::Ready(Err(err)) => {
362                    trace!(target: "rpc", ?err, "connection err");
363                    return false;
364                }
365                Poll::Pending => *writable = false,
366            }
367        }
368
369        if self.messages.len() >= MAX_PENDING_MESSAGES
370            || message.len() > MAX_PENDING_BYTES - self.bytes
371        {
372            warn!(target: "rpc", "Closing connection because its outgoing backlog is full");
373            return false;
374        }
375        self.bytes += message.len();
376        self.messages.push_back(message);
377        true
378    }
379}
380
381// Bound both serialized payloads and per-message allocation overhead for slow clients.
382const MAX_PENDING_BYTES: usize = 16 * 1024 * 1024;
383const MAX_PENDING_MESSAGES: usize = 4096;
384
385#[cfg(test)]
386mod tests {
387    use super::*;
388    use anvil_rpc::{
389        request::{RequestParams, RpcCall, RpcNotification, Version},
390        response::RpcResponse,
391    };
392    use futures::channel::mpsc;
393    use std::{
394        pin::pin,
395        sync::atomic::{AtomicUsize, Ordering},
396        task::{Wake, Waker},
397    };
398
399    #[derive(Clone, Default)]
400    struct TestHandler {
401        requests: Arc<AtomicUsize>,
402        remove_on_request: Option<u64>,
403        close_on_end: bool,
404    }
405
406    impl PubSubRpcHandler for TestHandler {
407        type Request = serde_json::Value;
408        type SubscriptionId = u64;
409        type Subscription = Pin<Box<dyn Stream<Item = serde_json::Value> + Send + Sync>>;
410
411        async fn on_request(
412            &self,
413            _request: Self::Request,
414            cx: PubSubContext<Self>,
415        ) -> ResponseResult {
416            self.requests.fetch_add(1, Ordering::Relaxed);
417            if let Some(id) = self.remove_on_request {
418                cx.remove_subscription(&id);
419            }
420            ResponseResult::success(serde_json::Value::Null)
421        }
422
423        fn close_on_subscription_end(&self, _subscription: &Self::Subscription) -> bool {
424            self.close_on_end
425        }
426    }
427
428    fn notification() -> RpcCall {
429        RpcCall::Notification(RpcNotification {
430            jsonrpc: Some(Version::V2),
431            method: "eth_subscribe".to_owned(),
432            params: RequestParams::None,
433        })
434    }
435
436    fn run_ready<F: Future>(future: F) -> F::Output {
437        let waker = Waker::noop();
438        let mut cx = Context::from_waker(waker);
439        let mut future = pin!(future);
440        match future.as_mut().poll(&mut cx) {
441            Poll::Ready(output) => output,
442            Poll::Pending => panic!("future unexpectedly pending"),
443        }
444    }
445
446    #[test]
447    fn process_request_keeps_empty_batch_invalid() {
448        let mut connection = PubSubConnection::new((), TestHandler::default());
449        connection.process_request(Ok(Request::Batch(vec![])));
450
451        let response = run_ready(connection.processing.pop().unwrap());
452        assert_eq!(
453            response,
454            Some(Response::Single(RpcResponse::from(RpcError::invalid_request())))
455        );
456    }
457
458    #[test]
459    fn process_request_returns_parse_error_for_malformed_json() {
460        let mut connection = PubSubConnection::new((), TestHandler::default());
461        connection.process_request(serde_json::from_str("{"));
462
463        let response = run_ready(connection.processing.pop().unwrap());
464        assert_eq!(response, Some(Response::error(RpcError::parse_error())));
465    }
466
467    #[test]
468    fn process_request_executes_notification_without_response() {
469        let handler = TestHandler::default();
470        let mut connection = PubSubConnection::new((), handler.clone());
471        connection.process_request(Ok(Request::Batch(vec![notification()])));
472
473        let response = run_ready(connection.processing.pop().unwrap());
474        assert_eq!(response, None);
475        assert_eq!(handler.requests.load(Ordering::Relaxed), 1);
476    }
477
478    #[derive(Default)]
479    struct TestConnection {
480        ready: bool,
481        fail_next_send: bool,
482        fail_after: Option<usize>,
483        ready_polls: usize,
484        sent: Vec<String>,
485        requests: VecDeque<Request>,
486        send_budget: Option<usize>,
487        ready_waker: Option<Waker>,
488        disconnected: bool,
489    }
490
491    impl Stream for TestConnection {
492        type Item = Result<Option<Request>, RequestError>;
493
494        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
495            if self.disconnected {
496                return Poll::Ready(None);
497            }
498            match self.requests.pop_front() {
499                Some(request) => Poll::Ready(Some(Ok(Some(request)))),
500                None => Poll::Pending,
501            }
502        }
503    }
504
505    impl Sink<String> for TestConnection {
506        type Error = &'static str;
507
508        fn poll_ready(
509            mut self: Pin<&mut Self>,
510            cx: &mut Context<'_>,
511        ) -> Poll<Result<(), Self::Error>> {
512            let polls = self.ready_polls;
513            self.ready_polls += 1;
514            if self.fail_after.is_some_and(|after| polls >= after) {
515                return Poll::Ready(Err("sink failed"));
516            }
517            if self.ready {
518                Poll::Ready(Ok(()))
519            } else {
520                self.ready_waker = Some(cx.waker().clone());
521                Poll::Pending
522            }
523        }
524
525        fn start_send(mut self: Pin<&mut Self>, item: String) -> Result<(), Self::Error> {
526            assert!(self.ready);
527            if self.fail_next_send {
528                self.fail_next_send = false;
529                return Err("sink failed");
530            }
531            self.sent.push(item);
532            if let Some(budget) = &mut self.send_budget {
533                *budget -= 1;
534                self.ready = *budget != 0;
535            }
536            Ok(())
537        }
538
539        fn poll_flush(
540            self: Pin<&mut Self>,
541            _cx: &mut Context<'_>,
542        ) -> Poll<Result<(), Self::Error>> {
543            Poll::Ready(Ok(()))
544        }
545
546        fn poll_close(
547            self: Pin<&mut Self>,
548            _cx: &mut Context<'_>,
549        ) -> Poll<Result<(), Self::Error>> {
550            Poll::Ready(Ok(()))
551        }
552    }
553
554    #[test]
555    fn sink_readiness_error_terminates_connection() {
556        for fail_after in [0, 1] {
557            let transport =
558                TestConnection { ready: true, fail_after: Some(fail_after), ..Default::default() };
559            let mut connection = PubSubConnection::new(transport, TestHandler::default());
560            connection.context.add_subscription(
561                1,
562                Box::pin(futures::stream::iter(vec![serde_json::Value::Null])),
563            );
564
565            assert!(poll_once(&mut connection).is_ready());
566            assert!(connection.connection.sent.is_empty());
567        }
568    }
569
570    #[test]
571    fn response_start_send_error_terminates_connection() {
572        let transport = TestConnection { ready: true, fail_next_send: true, ..Default::default() };
573        let mut connection = PubSubConnection::new(transport, TestHandler::default());
574        connection.connection.requests.push_back(Request::Batch(vec![]));
575
576        assert!(poll_once(&mut connection).is_ready());
577        assert!(!connection.connection.fail_next_send);
578    }
579
580    #[test]
581    fn subscription_backpressure_limits_readiness_scan_and_allows_disconnect() {
582        for count in [1, 64] {
583            let mut connection = connection();
584            for id in 0..count {
585                connection.context.add_subscription(
586                    id,
587                    Box::pin(futures::stream::iter(vec![serde_json::Value::from(id)])),
588                );
589            }
590
591            assert!(poll_once(&mut connection).is_pending());
592            // One readiness check per round, regardless of the number of subscriptions.
593            assert_eq!(connection.connection.ready_polls, 2);
594            assert!(connection.connection.sent.is_empty());
595            assert_eq!(connection.pending.messages.len(), count as usize);
596            assert!(connection.context.subscriptions.lock().is_empty());
597
598            connection.connection.disconnected = true;
599            assert!(poll_once(&mut connection).is_ready());
600            assert!(connection.connection.sent.is_empty());
601        }
602    }
603
604    #[derive(Default)]
605    struct TestWake(AtomicUsize);
606
607    impl Wake for TestWake {
608        fn wake(self: Arc<Self>) {
609            self.0.fetch_add(1, Ordering::Relaxed);
610        }
611    }
612
613    #[test]
614    fn subscription_backpressure_wakes_and_rotates_fairly() {
615        let count = 3;
616        let mut connection = connection();
617        for id in 1..=count {
618            let events = (0..if id == count { 12 } else { 3 })
619                .map(|event| serde_json::json!([id, event]))
620                .collect::<Vec<_>>();
621            connection.context.add_subscription(id, Box::pin(futures::stream::iter(events)));
622        }
623        connection.context.add_subscription(count + 1, Box::pin(futures::stream::pending()));
624        connection.context.add_subscription(count + 2, Box::pin(futures::stream::empty()));
625        let wake = Arc::new(TestWake::default());
626        let waker = Waker::from(wake.clone());
627        let mut cx = Context::from_waker(&waker);
628        assert!(Pin::new(&mut connection).poll(&mut cx).is_pending());
629
630        for window in 0..count * 3 {
631            connection.connection.ready = true;
632            connection.connection.send_budget = Some(1);
633            connection.connection.ready_waker.take().unwrap().wake();
634            assert_eq!(wake.0.swap(0, Ordering::Relaxed), 1);
635            assert!(Pin::new(&mut connection).poll(&mut cx).is_pending());
636            assert_eq!(connection.connection.sent.len(), window as usize + 1);
637        }
638        let mut sent = sent_events(&connection.connection);
639        for (event, window) in sent.as_chunks_mut::<3>().0.iter_mut().enumerate() {
640            window.sort_unstable();
641            assert_eq!(*window, [[1, event as u64], [2, event as u64], [3, event as u64]]);
642        }
643
644        connection.connection.ready = true;
645        connection.connection.send_budget = None;
646        connection.connection.ready_waker.take().unwrap().wake();
647        assert_eq!(wake.0.swap(0, Ordering::Relaxed), 1);
648        assert!(Pin::new(&mut connection).poll(&mut cx).is_pending());
649        let sent = sent_events(&connection.connection);
650        assert_eq!(sent.len(), 18);
651        for id in 1..=count {
652            let events = sent.iter().filter(|event| event[0] == id).map(|event| event[1]);
653            assert_eq!(
654                events.collect::<Vec<_>>(),
655                (0..if id == count { 12 } else { 3 }).collect::<Vec<_>>()
656            );
657        }
658        assert_eq!(sent[9..], (3..12).map(|event| [count, event]).collect::<Vec<_>>());
659        assert_eq!(connection.context.subscriptions.lock().len(), 1);
660        assert!(connection.context.remove_subscription(&(count + 1)).is_some());
661    }
662
663    #[test]
664    fn subscription_end_closes_only_when_requested() {
665        for close_on_end in [false, true] {
666            let handler = TestHandler { close_on_end, ..Default::default() };
667            let mut connection = PubSubConnection::new(TestConnection::default(), handler);
668            connection.context.add_subscription(1, Box::pin(futures::stream::empty()));
669
670            assert_eq!(poll_once(&mut connection).is_ready(), close_on_end);
671            assert!(connection.connection.sent.is_empty());
672        }
673    }
674
675    #[test]
676    fn subscription_backpressure_allows_unsubscribe_and_replacement() {
677        let handler = TestHandler { remove_on_request: Some(1), ..Default::default() };
678        let mut connection = PubSubConnection::new(TestConnection::default(), handler.clone());
679        let mut senders = Vec::new();
680        for id in 1..=3 {
681            let (tx, rx) = mpsc::unbounded();
682            tx.unbounded_send(serde_json::json!(id)).unwrap();
683            senders.push(tx);
684            connection.context.add_subscription(id, Box::pin(rx));
685        }
686        let mut cx = Context::from_waker(Waker::noop());
687        assert!(poll_once(&mut connection).is_pending());
688        assert_eq!(connection.pending.messages.len(), 3);
689
690        let mut replaced = connection
691            .context
692            .add_subscription(2, Box::pin(futures::stream::iter(vec![serde_json::json!(20)])))
693            .unwrap();
694        assert!(replaced.poll_next_unpin(&mut cx).is_pending());
695        drop(replaced);
696        connection.connection.requests.push_back(Request::Single(notification()));
697        assert!(poll_once(&mut connection).is_pending());
698        assert_eq!(handler.requests.load(Ordering::Relaxed), 1);
699        assert_eq!(connection.pending.messages.len(), 4);
700        assert!(connection.context.remove_subscription(&1).is_none());
701        assert!(connection.connection.sent.is_empty());
702        assert!(senders[0].unbounded_send(serde_json::json!(10)).is_err());
703        assert!(senders[1].unbounded_send(serde_json::json!(20)).is_err());
704        senders[2].unbounded_send(serde_json::json!(30)).unwrap();
705
706        connection.connection.ready = true;
707        assert!(poll_once(&mut connection).is_pending());
708        // Already materialized notifications precede later cancellation and replacement.
709        assert_eq!(connection.connection.sent.len(), 5);
710        let (initial, later) = connection.connection.sent.split_at_mut(3);
711        initial.sort_unstable();
712        assert_eq!(initial, ["1", "2", "3"]);
713        assert_eq!(later, ["20", "30"]);
714        assert!(connection.pending.messages.is_empty());
715        assert_eq!(connection.context.subscriptions.lock().len(), 1);
716    }
717
718    #[test]
719    fn subscription_backpressure_drains_bounded_upstream() {
720        let (mut tx, rx) = mpsc::channel(2048);
721        let mut connection = connection();
722        connection.context.add_subscription(1, Box::pin(rx.map(serde_json::Value::from)));
723
724        for event in 0..3000 {
725            tx.try_send(event).expect("the blocked connection must keep draining its upstream");
726            assert!(poll_once(&mut connection).is_pending());
727        }
728
729        connection.connection.ready = true;
730        assert!(poll_once(&mut connection).is_pending());
731        assert_eq!(
732            connection.connection.sent,
733            (0..3000).map(|event| event.to_string()).collect::<Vec<_>>()
734        );
735    }
736
737    #[test]
738    fn subscription_backlog_overflow_terminates_connection() {
739        let mut connection = connection();
740        connection
741            .context
742            .add_subscription(1, Box::pin(futures::stream::repeat(serde_json::Value::Null)));
743
744        assert!((0..=MAX_PENDING_MESSAGES).any(|_| poll_once(&mut connection).is_ready()));
745        assert_eq!(connection.pending.messages.len(), MAX_PENDING_MESSAGES);
746        assert!(connection.connection.sent.is_empty());
747    }
748
749    #[test]
750    fn responses_and_subscriptions_share_backlog_limit() {
751        let (tx, rx) = mpsc::unbounded();
752        let mut connection = connection();
753        connection.context.add_subscription(1, Box::pin(rx));
754
755        for _ in 0..MAX_PENDING_MESSAGES - 1 {
756            tx.unbounded_send(serde_json::Value::Null).unwrap();
757            assert!(poll_once(&mut connection).is_pending());
758        }
759        connection.connection.requests.push_back(Request::Batch(vec![]));
760        assert!(poll_once(&mut connection).is_pending());
761        assert_eq!(connection.pending.messages.len(), MAX_PENDING_MESSAGES);
762
763        tx.unbounded_send(serde_json::Value::Null).unwrap();
764        assert!(poll_once(&mut connection).is_ready());
765        assert!(connection.connection.sent.is_empty());
766    }
767
768    #[test]
769    fn backlog_byte_limit_reclaims_drained_messages() {
770        let (tx, rx) = mpsc::unbounded();
771        let mut connection = connection();
772        connection.context.add_subscription(1, Box::pin(rx));
773        // Include the JSON string's two quotes in the byte budget.
774        let message = serde_json::Value::String("x".repeat(MAX_PENDING_BYTES / 2 - 2));
775        for _ in 0..2 {
776            tx.unbounded_send(message.clone()).unwrap();
777            assert!(poll_once(&mut connection).is_pending());
778        }
779        assert_eq!(connection.pending.bytes, MAX_PENDING_BYTES);
780
781        connection.connection.ready = true;
782        connection.connection.send_budget = Some(1);
783        assert!(poll_once(&mut connection).is_pending());
784        assert_eq!(connection.connection.sent.len(), 1);
785        assert_eq!(connection.pending.bytes, MAX_PENDING_BYTES / 2);
786
787        tx.unbounded_send(message).unwrap();
788        assert!(poll_once(&mut connection).is_pending());
789        assert_eq!(connection.pending.bytes, MAX_PENDING_BYTES);
790
791        tx.unbounded_send(serde_json::Value::Null).unwrap();
792        assert!(poll_once(&mut connection).is_ready());
793        assert_eq!(connection.pending.messages.len(), 2);
794        assert_eq!(connection.pending.bytes, MAX_PENDING_BYTES);
795    }
796
797    #[test]
798    fn oversized_message_requires_writable_connection() {
799        for ready in [true, false] {
800            let transport = TestConnection { ready, ..Default::default() };
801            let mut connection = PubSubConnection::new(transport, TestHandler::default());
802            connection.context.add_subscription(
803                1,
804                Box::pin(futures::stream::iter([serde_json::Value::String(
805                    "x".repeat(MAX_PENDING_BYTES),
806                )])),
807            );
808
809            assert_eq!(poll_once(&mut connection).is_pending(), ready);
810            if ready {
811                assert_eq!(connection.connection.sent[0].len(), MAX_PENDING_BYTES + 2);
812            } else {
813                assert!(connection.connection.sent.is_empty());
814            }
815            assert_eq!(connection.pending.bytes, 0);
816            assert!(connection.pending.messages.is_empty());
817        }
818    }
819
820    #[test]
821    fn writable_subscription_burst_yields_and_rotates_without_buffering() {
822        let transport = TestConnection { ready: true, ..Default::default() };
823        let mut connection = PubSubConnection::new(transport, TestHandler::default());
824        connection
825            .context
826            .add_subscription(1, Box::pin(futures::stream::iter([serde_json::json!(2)])));
827        connection.context.add_subscription(
828            2,
829            Box::pin(futures::stream::iter(vec![serde_json::json!(1); MAX_PENDING_MESSAGES + 1])),
830        );
831        let wake = Arc::new(TestWake::default());
832        let waker = Waker::from(wake.clone());
833        let mut cx = Context::from_waker(&waker);
834
835        assert!(Pin::new(&mut connection).poll(&mut cx).is_pending());
836        assert_eq!(wake.0.swap(0, Ordering::Relaxed), 1);
837        let mut first = connection.connection.sent[..2].to_vec();
838        first.sort_unstable();
839        assert_eq!(first, ["1", "2"]);
840        assert!(connection.connection.sent.len() < MAX_PENDING_MESSAGES);
841
842        for _ in 0..MAX_PENDING_MESSAGES {
843            assert!(Pin::new(&mut connection).poll(&mut cx).is_pending());
844            if connection.context.subscriptions.lock().is_empty() {
845                break;
846            }
847        }
848        assert!(connection.context.subscriptions.lock().is_empty());
849        assert_eq!(connection.connection.sent.len(), MAX_PENDING_MESSAGES + 2);
850        assert_eq!(connection.connection.sent[2..], vec!["1"; MAX_PENDING_MESSAGES]);
851        assert!(connection.pending.messages.is_empty());
852        assert_eq!(connection.pending.bytes, 0);
853    }
854
855    fn connection() -> PubSubConnection<TestHandler, TestConnection> {
856        PubSubConnection::new(TestConnection::default(), TestHandler::default())
857    }
858
859    fn poll_once(connection: &mut PubSubConnection<TestHandler, TestConnection>) -> Poll<()> {
860        Pin::new(connection).poll(&mut Context::from_waker(Waker::noop()))
861    }
862
863    fn sent_events(connection: &TestConnection) -> Vec<[u64; 2]> {
864        connection.sent.iter().map(|event| serde_json::from_str(event).unwrap()).collect()
865    }
866}
anvil_server/pubsub.rs

anvil_server/
pubsub.rs
