jsonrpsee_server/transport/
ws.rs

1use std::{sync::Arc, time::Duration};
2
3use crate::{
4    server::{handle_rpc_call, ConnectionState, ServerConfig},
5    LOG_TARGET,
6};
7
8use futures_util::{
9    future::{self, Either},
10    io::{BufReader, BufWriter},
11    Future, StreamExt, TryStreamExt,
12};
13use hyper::upgrade::Upgraded;
14use hyper_util::rt::TokioIo;
15use jsonrpsee_core::{classify_message, middleware::RpcServiceT, server::MethodSink};
16use jsonrpsee_types::{
17    error::{reject_too_big_request, ErrorCode},
18    Id,
19};
20use soketto::connection::Error as SokettoError;
21use tokio::{
22    sync::{mpsc, oneshot, Semaphore},
23    time::timeout,
24};
25use tokio_stream::wrappers::ReceiverStream;
26use tokio_util::compat::Compat;
27
28pub(crate) type Sender = soketto::Sender<BufReader<BufWriter<Compat<TokioIo<Upgraded>>>>>;
29pub(crate) type Receiver = soketto::Receiver<BufReader<BufWriter<Compat<TokioIo<Upgraded>>>>>;
30
31pub use soketto::handshake::http::is_upgrade_request;
32
33/// Most queued messages written before a flush, which bounds how long written messages wait in
34/// the buffer.
35const MAX_MESSAGES_PER_FLUSH: usize = 64;
36
37pub(crate) struct BackgroundTaskParams<S> {
38    pub(crate) server_cfg: ServerConfig,
39    pub(crate) conn: ConnectionState,
40    pub(crate) ws_sender: Sender,
41    pub(crate) ws_receiver: Receiver,
42    pub(crate) rpc_service: S,
43    pub(crate) sink: MethodSink,
44    pub(crate) rx: mpsc::Receiver<String>,
45    /// Held by every running call, so that a graceful shutdown can wait for them.
46    pub(crate) pending_calls: mpsc::Sender<()>,
47    pub(crate) pending_calls_completed: mpsc::Receiver<()>,
48}
49
50pub(crate) async fn background_task<S>(params: BackgroundTaskParams<S>)
51where
52    S: RpcServiceT + Send + Sync + 'static,
53{
54    let BackgroundTaskParams {
55        server_cfg,
56        conn,
57        ws_sender,
58        ws_receiver,
59        rpc_service,
60        sink,
61        rx,
62        pending_calls,
63        pending_calls_completed,
64    } = params;
65    let ServerConfig {
66        batch_requests_config,
67        max_request_body_size,
68        write_timeout,
69        message_buffer_capacity,
70        ..
71    } = server_cfg;
72
73    let (conn_tx, conn_rx) = oneshot::channel();
74
75    // Spawn another task that sends out the responses on the Websocket.
76    let send_task_handle = tokio::spawn(send_task(rx, ws_sender, write_timeout, conn_rx));
77
78    let stopped = conn.stop_handle.clone().shutdown();
79    let rpc_service = Arc::new(rpc_service);
80
81    tokio::pin!(stopped);
82
83    let ws_stream = futures_util::stream::unfold(ws_receiver, |mut receiver| async {
84        let mut data = Vec::new();
85        loop {
86            match receiver.receive(&mut data).await {
87                Ok(soketto::Incoming::Data(_)) => return Some((Ok(data), receiver)),
88                // The server sends no pings, so pongs are unsolicited and ignored.
89                Ok(soketto::Incoming::Pong(_)) => data.clear(),
90                Ok(soketto::Incoming::Closed(_)) | Err(SokettoError::Closed) => return None,
91                // The closing reason is already logged by `soketto` trace log level.
92                // Return the `Closed` error to avoid logging unnecessary warnings on clean
93                // shutdown.
94                Err(e) => return Some((Err(e), receiver)),
95            }
96        }
97    })
98    .fuse();
99
100    tokio::pin!(ws_stream);
101
102    let call_permits = Arc::new(Semaphore::new(message_buffer_capacity as usize));
103
104    let result = loop {
105        let receive = tokio::select! {
106            receive = try_recv(&mut ws_stream, stopped) => receive,
107            // The send task has terminated, e.g. because a write timed out, so nothing can be sent
108            // back to the client anymore.
109            _ = sink.closed() => break Ok(Shutdown::ConnectionClosed),
110        };
111        let data = match receive {
112            Receive::ConnectionClosed => break Ok(Shutdown::ConnectionClosed),
113            Receive::Stopped => break Ok(Shutdown::Stopped),
114            Receive::Ok(data, stop) => {
115                stopped = stop;
116                data
117            }
118            Receive::Err(err, stop) => {
119                stopped = stop;
120
121                match err {
122                    SokettoError::Closed => {
123                        break Ok(Shutdown::ConnectionClosed);
124                    }
125                    SokettoError::MessageTooLarge { current, maximum } => {
126                        tracing::debug!(
127                            target: LOG_TARGET,
128                            "WS recv error: message too large current={}/max={}",
129                            current,
130                            maximum
131                        );
132                        if sink
133                            .send_error(Id::Null, reject_too_big_request(max_request_body_size))
134                            .await
135                            .is_err()
136                        {
137                            break Ok(Shutdown::ConnectionClosed);
138                        }
139
140                        continue;
141                    }
142                    err => {
143                        tracing::debug!(target: LOG_TARGET, "WS error: {}; terminate connection: {}", err, conn.conn_id);
144                        break Err(err);
145                    }
146                };
147            }
148        };
149
150        // A connection runs at most `message_buffer_capacity` calls at once. While waiting for a
151        // permit no further messages are read, which pushes back on the client, and the message
152        // that is waiting is the only one held. The permit is acquired after the read so that a
153        // connection at the limit that sends nothing more is still read and its close is noticed.
154        // Every call ends once the sink is closed, which frees the permits, so this wait ends
155        // when the connection closes.
156        let permit = tokio::select! {
157            permit = call_permits.clone().acquire_owned() => permit,
158            _ = &mut stopped => break Ok(Shutdown::Stopped),
159        };
160        // The semaphore is never closed.
161        let Ok(permit) = permit else { break Ok(Shutdown::ConnectionClosed) };
162
163        let rpc_service = rpc_service.clone();
164        let sink = sink.clone();
165        let pending_calls = pending_calls.clone();
166
167        tokio::spawn(async move {
168            let _pending_calls = pending_calls;
169            let _permit = permit;
170            let Some((data, is_single)) = classify_message(&data) else {
171                _ = sink.send_error(Id::Null, ErrorCode::ParseError.into()).await;
172                return;
173            };
174
175            // Calls are not tied to the connection task, so stop them once the connection is closed
176            // instead of running them to completion without anyone to send the response to.
177            let rp = tokio::select! {
178                // The call is polled first, so one that answers at once doesn't have to register
179                // for the close of the connection.
180                biased;
181                rp = handle_rpc_call(data, is_single, batch_requests_config, &*rpc_service) => rp,
182                _ = sink.closed() => return,
183            };
184
185            // Subscriptions are handled by the subscription callback and
186            // "ordinary notifications" should not be sent back to the client.
187            if rp.is_method_call() || rp.is_batch() {
188                // An error means the connection is closed, so there is nothing left to do.
189                let _ = sink.send(rp.into_json()).await;
190            }
191        });
192    };
193
194    // Drive all running methods to completion.
195    // **NOTE** Do not return early in this function. This `await` needs to run to guarantee
196    // proper drop behaviour.
197    drop(rpc_service);
198    drop(pending_calls);
199    graceful_shutdown(result, pending_calls_completed, ws_stream, conn_tx, send_task_handle).await;
200
201    drop(conn);
202}
203
204/// A task that waits for new messages via the `rx channel` and sends them out on the `WebSocket`.
205async fn send_task(
206    mut rx: mpsc::Receiver<String>,
207    mut ws_sender: Sender,
208    write_timeout: Duration,
209    mut stop: oneshot::Receiver<()>,
210) {
211    let mut write_timed_out = false;
212
213    loop {
214        let response = tokio::select! {
215            // Stop is polled first: a subscription that queues messages faster than the client
216            // reads them keeps `rx.recv()` ready, which would never let a stopped server end the
217            // session.
218            biased;
219            _ = &mut stop, if !rx.is_closed() => {
220                // Take nothing new but write out what is queued: answers to calls this server has
221                // run, which a graceful shutdown owes its clients. `rx.recv()` returns `None` once
222                // the closed queue is drained.
223                rx.close();
224                continue;
225            }
226            response = rx.recv() => response,
227        };
228        // Nothing else to receive.
229        let Some(response) = response else { break };
230
231        let mut result = timeout(write_timeout, ws_sender.send_text_owned(response)).await;
232        // Write the messages that are already queued too, and flush them together.
233        let mut written = 1;
234        while matches!(result, Ok(Ok(()))) && written < MAX_MESSAGES_PER_FLUSH {
235            let Ok(response) = rx.try_recv() else { break };
236            result = timeout(write_timeout, ws_sender.send_text_owned(response)).await;
237            written += 1;
238        }
239        if matches!(result, Ok(Ok(()))) {
240            result = timeout(write_timeout, ws_sender.flush()).await;
241        }
242
243        // If websocket message send fail then terminate the connection.
244        match result {
245            Ok(Ok(())) => {}
246            Ok(Err(err)) => {
247                tracing::debug!(target: LOG_TARGET, "WS send error: {}", err);
248                break;
249            }
250            Err(_) => {
251                tracing::debug!(target: LOG_TARGET, "WS send timed out after {write_timeout:?}");
252                write_timed_out = true;
253                break;
254            }
255        }
256    }
257
258    // Closing `rx` makes `MethodSink::closed` resolve, which stops the receive loop, pending calls
259    // and subscriptions of this connection. Messages still queued are dropped with this task.
260    rx.close();
261
262    // Send the close message, unless a write already timed out: the client isn't reading, so the
263    // close frame would block on the same full socket. The receive half is dropped by the
264    // connection task, which closes the socket.
265    if !write_timed_out {
266        let _ = timeout(write_timeout, ws_sender.close()).await;
267    }
268}
269
270enum Receive<S> {
271    ConnectionClosed,
272    Stopped,
273    Err(SokettoError, S),
274    Ok(Vec<u8>, S),
275}
276
277/// Attempts to read data from WebSocket fails if the server was stopped.
278async fn try_recv<T, S>(ws_stream: &mut T, stopped: S) -> Receive<S>
279where
280    S: Future<Output = ()> + Unpin,
281    T: StreamExt<Item = Result<Vec<u8>, SokettoError>> + Unpin,
282{
283    // `select` polls its first argument first and returns as soon as that one is ready, so the stop
284    // signal has to come first. Polling the stream first starves it: a peer that always has another
285    // message ready never lets the read be pending, and the session would then never stop.
286    match future::select(stopped, ws_stream.next()).await {
287        // Server has been stopped.
288        Either::Left(_) => Receive::Stopped,
289        // The connection is closed.
290        Either::Right((None, _)) => Receive::ConnectionClosed,
291        // The message has been received, we are done
292        Either::Right((Some(Ok(data)), stopped)) => Receive::Ok(data, stopped),
293        // Received an error, terminate the connection.
294        Either::Right((Some(Err(e)), stopped)) => Receive::Err(e, stopped),
295    }
296}
297
298#[derive(Debug, Copy, Clone)]
299pub(crate) enum Shutdown {
300    Stopped,
301    ConnectionClosed,
302}
303
304/// Enforce a graceful shutdown.
305///
306/// This will return once the connection has been terminated or all pending calls have been
307/// executed.
308async fn graceful_shutdown<S>(
309    result: Result<Shutdown, SokettoError>,
310    pending_calls: mpsc::Receiver<()>,
311    ws_stream: S,
312    mut conn_tx: oneshot::Sender<()>,
313    send_task_handle: tokio::task::JoinHandle<()>,
314) where
315    S: StreamExt<Item = Result<Vec<u8>, SokettoError>> + Unpin,
316{
317    let pending_calls = ReceiverStream::new(pending_calls);
318
319    if let Ok(Shutdown::Stopped) = result {
320        let graceful_shutdown = pending_calls.for_each(|_| async {});
321        let disconnect = ws_stream.try_for_each(|_| async { Ok(()) });
322
323        tokio::select! {
324            _ = graceful_shutdown => {}
325            res = disconnect => {
326                if let Err(err) = res {
327                    tracing::warn!(target: LOG_TARGET, "Graceful shutdown terminated because of error: `{err}`");
328                }
329            }
330            _ = conn_tx.closed() => {}
331        }
332    }
333
334    // Send a message to close down the "send task".
335    _ = conn_tx.send(());
336    // Ensure that send task has been closed.
337    _ = send_task_handle.await;
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343    use futures_util::stream;
344
345    #[tokio::test]
346    async fn try_recv_observes_stop_while_input_stays_ready() {
347        // A peer that always has another message ready must not keep the session from stopping.
348        let mut always_ready = stream::repeat_with(|| Ok::<_, SokettoError>(b"{}".to_vec()));
349
350        let received = try_recv(&mut always_ready, future::ready(())).await;
351        assert!(matches!(received, Receive::Stopped));
352    }
353}
jsonrpsee_server/transport/ws.rs

jsonrpsee_server/transport/
ws.rs
