alloy_pubsub/
frontend.rs

1use crate::{ix::PubSubInstruction, managers::InFlight, RawSubscription};
2use alloy_json_rpc::{RequestPacket, Response, ResponsePacket, SerializedRequest};
3use alloy_primitives::B256;
4use alloy_transport::{TransportError, TransportErrorKind, TransportFut, TransportResult};
5use futures::{future::try_join_all, FutureExt, TryFutureExt};
6use std::{
7    future::Future,
8    sync::{
9        atomic::{AtomicUsize, Ordering},
10        Arc,
11    },
12    task::{Context, Poll},
13};
14use tokio::sync::{mpsc, oneshot};
15use tracing::{debug, debug_span, Instrument};
16
17/// A `PubSubFrontend` is [`Transport`] composed of a channel to a running
18/// PubSub service.
19///
20/// [`Transport`]: alloy_transport::Transport
21#[derive(Debug, Clone)]
22pub struct PubSubFrontend {
23    tx: mpsc::UnboundedSender<PubSubInstruction>,
24    /// The number of items to buffer in new subscription channels. Defaults to
25    /// 16. See [`tokio::sync::broadcast::channel`] for a description.
26    channel_size: Arc<AtomicUsize>,
27}
28
29impl PubSubFrontend {
30    /// Create a new frontend.
31    pub fn new(tx: mpsc::UnboundedSender<PubSubInstruction>) -> Self {
32        Self { tx, channel_size: Arc::new(AtomicUsize::new(16)) }
33    }
34
35    /// Get the subscription ID for a local ID.
36    pub fn get_subscription(
37        &self,
38        id: B256,
39    ) -> impl Future<Output = TransportResult<RawSubscription>> + Send + 'static {
40        let backend_tx = self.tx.clone();
41        async move {
42            let (tx, rx) = oneshot::channel();
43            backend_tx
44                .send(PubSubInstruction::GetSub(id, tx))
45                .map_err(|_| TransportErrorKind::backend_gone())?;
46            rx.await
47                .map_err(|_| TransportErrorKind::backend_gone())?
48                .map_or_else(|| Err(TransportErrorKind::custom_str("subscription not found")), Ok)
49        }
50    }
51
52    /// Unsubscribe from a subscription.
53    pub fn unsubscribe(&self, id: B256) -> TransportResult<()> {
54        self.tx
55            .send(PubSubInstruction::Unsubscribe(id))
56            .map_err(|_| TransportErrorKind::backend_gone())
57    }
58
59    /// Send a request.
60    pub fn send(
61        &self,
62        req: SerializedRequest,
63    ) -> impl Future<Output = TransportResult<Response>> + Send + 'static {
64        let tx = self.tx.clone();
65        let channel_size = self.channel_size.load(Ordering::Relaxed);
66        let method_name = req.method_clone();
67
68        async move {
69            debug!("sending request to backend");
70            let (in_flight, rx) = InFlight::new(req, channel_size);
71            tx.send(PubSubInstruction::Request(in_flight))
72                .map_err(|_| TransportErrorKind::backend_gone())?;
73            let resp = rx.await.map_err(|_| TransportErrorKind::backend_gone())?;
74            if tracing::enabled!(tracing::Level::TRACE) {
75                trace!(?resp, "retrieved response");
76            } else {
77                debug!(resp=?resp.as_ref().map(|_| ()), "retrieved response");
78            };
79            resp
80        }
81        .instrument(debug_span!("request", %method_name))
82    }
83
84    /// Send a packet of requests, by breaking it up into individual requests.
85    ///
86    /// Once all responses are received, we return a single response packet.
87    pub fn send_packet(&self, req: RequestPacket) -> TransportFut<'static> {
88        match req {
89            RequestPacket::Single(req) => self.send(req).map_ok(ResponsePacket::Single).boxed(),
90            RequestPacket::Batch(reqs) => try_join_all(reqs.into_iter().map(|req| self.send(req)))
91                .map_ok(ResponsePacket::Batch)
92                .boxed(),
93        }
94    }
95
96    /// Get the currently configured channel size. This is the number of items
97    /// to buffer in new subscription channels. Defaults to 16. See
98    /// [`tokio::sync::broadcast`] for a description of relevant
99    /// behavior.
100    pub fn channel_size(&self) -> usize {
101        self.channel_size.load(Ordering::Relaxed)
102    }
103
104    /// Set the channel size. This is the number of items to buffer in new
105    /// subscription channels. Defaults to 16. See
106    /// [`tokio::sync::broadcast`] for a description of relevant
107    /// behavior.
108    pub fn set_channel_size(&self, channel_size: usize) {
109        debug_assert_ne!(channel_size, 0, "channel size must be non-zero");
110        self.channel_size.store(channel_size, Ordering::Relaxed);
111    }
112}
113
114impl tower::Service<RequestPacket> for PubSubFrontend {
115    type Response = ResponsePacket;
116    type Error = TransportError;
117    type Future = TransportFut<'static>;
118
119    #[inline]
120    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
121        let result =
122            if self.tx.is_closed() { Err(TransportErrorKind::backend_gone()) } else { Ok(()) };
123        Poll::Ready(result)
124    }
125
126    #[inline]
127    fn call(&mut self, req: RequestPacket) -> Self::Future {
128        self.send_packet(req)
129    }
130}
alloy_pubsub/frontend.rs

alloy_pubsub/
frontend.rs
