alloy_rpc_client/
batch.rs

1use crate::{client::RpcClientInner, ClientRef};
2use alloy_json_rpc::{
3    transform_response, try_deserialize_ok, Id, Request, RequestPacket, ResponsePacket, RpcRecv,
4    RpcSend, SerializedRequest,
5};
6use alloy_primitives::map::HashMap;
7use alloy_transport::{
8    BoxTransport, TransportError, TransportErrorKind, TransportFut, TransportResult,
9};
10use futures::FutureExt;
11use pin_project::pin_project;
12use serde_json::value::RawValue;
13use std::{
14    borrow::Cow,
15    future::{Future, IntoFuture},
16    marker::PhantomData,
17    pin::Pin,
18    task::{
19        self, ready,
20        Poll::{self, Ready},
21    },
22};
23use tokio::sync::oneshot;
24use tower::Service;
25
26pub(crate) type Channel = oneshot::Sender<TransportResult<Box<RawValue>>>;
27pub(crate) type ChannelMap = HashMap<Id, Channel>;
28
29/// A batch JSON-RPC request, used to bundle requests into a single transport
30/// call.
31///
32/// Calls are serialized when they are added. Sending the batch is still lazy:
33/// [`Self::send`] returns a future, and the transport is not called until that
34/// future is polled. Await the batch before awaiting its [`Waiter`]s.
35///
36/// Responses are matched to waiters by JSON-RPC ID, not response order.
37/// Missing IDs produce a per-waiter error; responses with unknown IDs are
38/// ignored.
39#[derive(Debug)]
40#[must_use = "a BatchRequest does nothing unless it is awaited or `send()` is awaited"]
41pub struct BatchRequest<'a> {
42    /// The transport via which the batch will be sent.
43    transport: ClientRef<'a>,
44
45    /// The requests to be sent.
46    requests: RequestPacket,
47
48    /// The channels to send the responses through.
49    channels: ChannelMap,
50}
51
52/// Awaits a single response for a request that has been included in a batch.
53///
54/// Await the corresponding [`BatchRequest`] first so transport-level failures
55/// are observed and the response channels are populated.
56#[must_use = "a Waiter requires its BatchRequest to be sent and the Waiter to be awaited"]
57#[pin_project]
58#[derive(Debug)]
59pub struct Waiter<Resp, Output = Resp, Map = fn(Resp) -> Output> {
60    #[pin]
61    rx: oneshot::Receiver<TransportResult<Box<RawValue>>>,
62    map: Option<Map>,
63    _resp: PhantomData<fn() -> (Output, Resp)>,
64}
65
66impl<Resp, Output, Map> Waiter<Resp, Output, Map> {
67    /// Map the response to a different type. This is usable for converting
68    /// the response to a more usable type, e.g. changing `U64` to `u64`.
69    ///
70    /// ## Note
71    ///
72    /// Carefully review the rust documentation on [fn pointers] before passing
73    /// them to this function. Unless the pointer is specifically coerced to a
74    /// `fn(_) -> _`, the `NewMap` will be inferred as that function's unique
75    /// type. This can lead to confusing error messages.
76    ///
77    /// [fn pointers]: https://doc.rust-lang.org/std/primitive.fn.html#creating-function-pointers
78    pub fn map_resp<NewOutput, NewMap>(self, map: NewMap) -> Waiter<Resp, NewOutput, NewMap>
79    where
80        NewMap: FnOnce(Resp) -> NewOutput,
81    {
82        Waiter { rx: self.rx, map: Some(map), _resp: PhantomData }
83    }
84}
85
86impl<Resp> From<oneshot::Receiver<TransportResult<Box<RawValue>>>> for Waiter<Resp> {
87    fn from(rx: oneshot::Receiver<TransportResult<Box<RawValue>>>) -> Self {
88        Self { rx, map: Some(std::convert::identity), _resp: PhantomData }
89    }
90}
91
92impl<Resp, Output, Map> std::future::Future for Waiter<Resp, Output, Map>
93where
94    Resp: RpcRecv,
95    Map: FnOnce(Resp) -> Output,
96{
97    type Output = TransportResult<Output>;
98
99    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
100        let this = self.get_mut();
101
102        match ready!(this.rx.poll_unpin(cx)) {
103            Ok(resp) => {
104                let resp: Result<Resp, _> = try_deserialize_ok(resp);
105                Ready(resp.map(this.map.take().expect("polled after completion")))
106            }
107            Err(e) => Poll::Ready(Err(TransportErrorKind::custom(e))),
108        }
109    }
110}
111
112#[pin_project::pin_project(project = CallStateProj)]
113#[expect(unnameable_types, missing_debug_implementations)]
114pub enum BatchFuture {
115    Prepared {
116        transport: BoxTransport,
117        requests: RequestPacket,
118        channels: ChannelMap,
119    },
120    AwaitingResponse {
121        channels: ChannelMap,
122        #[pin]
123        fut: TransportFut<'static>,
124    },
125    Complete,
126}
127
128impl<'a> BatchRequest<'a> {
129    /// Create a new batch request.
130    pub fn new(transport: &'a RpcClientInner) -> Self {
131        Self {
132            transport,
133            requests: RequestPacket::Batch(Vec::with_capacity(10)),
134            channels: HashMap::with_capacity_and_hasher(10, Default::default()),
135        }
136    }
137
138    fn push_raw(
139        &mut self,
140        request: SerializedRequest,
141    ) -> oneshot::Receiver<TransportResult<Box<RawValue>>> {
142        let (tx, rx) = oneshot::channel();
143        self.channels.insert(request.id().clone(), tx);
144        self.requests.push(request);
145        rx
146    }
147
148    fn push<Params: RpcSend, Resp: RpcRecv>(
149        &mut self,
150        request: Request<Params>,
151    ) -> TransportResult<Waiter<Resp>> {
152        let ser = request.serialize().map_err(TransportError::ser_err)?;
153        Ok(self.push_raw(ser).into())
154    }
155
156    /// Add a call to the batch.
157    ///
158    /// Unlike [`RpcCall`](crate::RpcCall), this serializes the parameters
159    /// immediately so calls with different parameter types can share a batch.
160    ///
161    /// ### Errors
162    ///
163    /// If the request cannot be serialized, this will return an error.
164    pub fn add_call<Params: RpcSend, Resp: RpcRecv>(
165        &mut self,
166        method: impl Into<Cow<'static, str>>,
167        params: &Params,
168    ) -> TransportResult<Waiter<Resp>> {
169        let request = self.transport.make_request(method, Cow::Borrowed(params));
170        self.push(request)
171    }
172
173    /// Return a future that sends the batch when polled.
174    ///
175    /// This method alone does not dispatch anything; await the returned future.
176    pub fn send(self) -> BatchFuture {
177        BatchFuture::Prepared {
178            transport: self.transport.transport.clone(),
179            requests: self.requests,
180            channels: self.channels,
181        }
182    }
183}
184
185impl IntoFuture for BatchRequest<'_> {
186    type Output = <BatchFuture as Future>::Output;
187    type IntoFuture = BatchFuture;
188
189    fn into_future(self) -> Self::IntoFuture {
190        self.send()
191    }
192}
193
194impl BatchFuture {
195    fn poll_prepared(
196        mut self: Pin<&mut Self>,
197        cx: &mut task::Context<'_>,
198    ) -> Poll<<Self as Future>::Output> {
199        let CallStateProj::Prepared { transport, requests, channels } = self.as_mut().project()
200        else {
201            unreachable!("Called poll_prepared in incorrect state")
202        };
203
204        if let Err(e) = task::ready!(transport.poll_ready(cx)) {
205            self.set(Self::Complete);
206            return Poll::Ready(Err(e));
207        }
208
209        // We only have mut refs, and we want ownership, so we just replace with 0-capacity
210        // collections.
211        let channels = std::mem::take(channels);
212        let req = std::mem::replace(requests, RequestPacket::Batch(Vec::new()));
213
214        let fut = transport.call(req);
215        self.set(Self::AwaitingResponse { channels, fut });
216        cx.waker().wake_by_ref();
217        Poll::Pending
218    }
219
220    fn poll_awaiting_response(
221        mut self: Pin<&mut Self>,
222        cx: &mut task::Context<'_>,
223    ) -> Poll<<Self as Future>::Output> {
224        let CallStateProj::AwaitingResponse { channels, fut } = self.as_mut().project() else {
225            unreachable!("Called poll_awaiting_response in incorrect state")
226        };
227
228        // Has the service responded yet?
229        let responses = match ready!(fut.poll(cx)) {
230            Ok(responses) => responses,
231            Err(e) => {
232                self.set(Self::Complete);
233                return Poll::Ready(Err(e));
234            }
235        };
236
237        // Send all responses via channels
238        match responses {
239            ResponsePacket::Single(single) => {
240                if let Some(tx) = channels.remove(&single.id) {
241                    let _ = tx.send(transform_response(single));
242                }
243            }
244            ResponsePacket::Batch(responses) => {
245                for response in responses {
246                    if let Some(tx) = channels.remove(&response.id) {
247                        let _ = tx.send(transform_response(response));
248                    }
249                }
250            }
251        }
252
253        // Any channels remaining in the map are missing responses.
254        // To avoid hanging futures, we send an error.
255        for (id, tx) in channels.drain() {
256            let _ = tx.send(Err(TransportErrorKind::missing_batch_response(id)));
257        }
258
259        self.set(Self::Complete);
260        Poll::Ready(Ok(()))
261    }
262}
263
264impl Future for BatchFuture {
265    type Output = TransportResult<()>;
266
267    fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
268        if matches!(*self.as_mut(), Self::Prepared { .. }) {
269            return self.poll_prepared(cx);
270        }
271
272        if matches!(*self.as_mut(), Self::AwaitingResponse { .. }) {
273            return self.poll_awaiting_response(cx);
274        }
275
276        panic!("Called poll on BatchFuture in invalid state")
277    }
278}
alloy_rpc_client/batch.rs

alloy_rpc_client/
batch.rs
