linera_rpc/simple/
transport.rs

1// Copyright (c) Facebook, Inc. and its affiliates.
2// Copyright (c) Zefchain Labs, Inc.
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::HashMap,
7    io, mem,
8    net::SocketAddr,
9    pin::{pin, Pin},
10    sync::Arc,
11};
12
13use async_trait::async_trait;
14use futures::{
15    future,
16    stream::{self, FuturesUnordered, SplitSink, SplitStream},
17    Sink, SinkExt, Stream, StreamExt, TryStreamExt,
18};
19use linera_base::identifiers::{BlobId, ChainId};
20use linera_core::{JoinSetExt as _, TaskHandle};
21use serde::{Deserialize, Serialize};
22use tokio::{
23    io::AsyncWriteExt,
24    net::{lookup_host, TcpListener, TcpStream, ToSocketAddrs, UdpSocket},
25    sync::Mutex,
26    task::JoinSet,
27};
28use tokio_util::{codec::Framed, sync::CancellationToken, udp::UdpFramed};
29use tracing::{error, warn};
30
31use crate::{
32    simple::{codec, codec::Codec},
33    RpcMessage,
34};
35
36/// Suggested buffer size
37pub const DEFAULT_MAX_DATAGRAM_SIZE: &str = "65507";
38
39/// Number of tasks to spawn before attempting to reap some finished tasks to prevent memory leaks.
40const REAP_TASKS_THRESHOLD: usize = 100;
41
42/// The transport protocols supported by the simple network.
43#[derive(clap::ValueEnum, Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
44pub enum TransportProtocol {
45    /// The UDP transport protocol.
46    Udp,
47    /// The TCP transport protocol.
48    Tcp,
49}
50
51impl std::str::FromStr for TransportProtocol {
52    type Err = String;
53
54    fn from_str(s: &str) -> Result<Self, Self::Err> {
55        clap::ValueEnum::from_str(s, true)
56    }
57}
58
59impl std::fmt::Display for TransportProtocol {
60    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
61        write!(f, "{self:?}")
62    }
63}
64
65impl TransportProtocol {
66    /// Returns the URL scheme name for this transport protocol.
67    pub fn scheme(&self) -> &'static str {
68        match self {
69            TransportProtocol::Udp => "udp",
70            TransportProtocol::Tcp => "tcp",
71        }
72    }
73}
74
75/// A pool of (outgoing) data streams.
76pub trait ConnectionPool: Send {
77    /// Sends a message to the given address, opening a connection if necessary.
78    fn send_message_to<'a>(
79        &'a mut self,
80        message: RpcMessage,
81        address: &'a str,
82    ) -> future::BoxFuture<'a, Result<(), codec::Error>>;
83}
84
85/// The handler required to create a service.
86///
87/// The implementation needs to implement [`Clone`] because a seed instance is used to generate
88/// cloned instances, where each cloned instance handles a single request. Multiple cloned instances
89/// may exist at the same time and handle separate requests concurrently.
90#[async_trait]
91pub trait MessageHandler: Clone {
92    /// Handles a single request message, returning an optional response.
93    async fn handle_message(&mut self, message: RpcMessage) -> Option<RpcMessage>;
94
95    /// Handle a notification subscription request. Returns a stream of notification
96    /// messages if supported, or `None` if subscriptions are not supported.
97    async fn handle_subscribe(
98        &mut self,
99        _chains: Vec<ChainId>,
100    ) -> Option<Pin<Box<dyn Stream<Item = RpcMessage> + Send>>> {
101        None
102    }
103
104    /// Handle a batch blob download request by streaming one
105    /// `RpcMessage::DownloadBlobResponse` per requested blob ID.
106    /// Returns `None` if not supported.
107    async fn handle_download_blobs(
108        &mut self,
109        _blob_ids: Vec<BlobId>,
110    ) -> Option<Pin<Box<dyn Stream<Item = RpcMessage> + Send>>> {
111        None
112    }
113}
114
115/// The result of spawning a server is oneshot channel to track completion, and the set of
116/// executing tasks.
117pub struct ServerHandle {
118    /// The handle tracking completion of the server task.
119    pub handle: TaskHandle<Result<(), std::io::Error>>,
120}
121
122impl ServerHandle {
123    /// Waits for the server task to finish.
124    pub async fn join(self) -> Result<(), std::io::Error> {
125        self.handle.await.map_err(|_| {
126            std::io::Error::new(
127                std::io::ErrorKind::Interrupted,
128                "Server task did not finish successfully",
129            )
130        })?
131    }
132}
133
134/// A trait alias for a protocol transport.
135///
136/// A transport is an active connection that can be used to send and receive
137/// [`RpcMessage`]s.
138pub trait Transport:
139    Stream<Item = Result<RpcMessage, codec::Error>> + Sink<RpcMessage, Error = codec::Error>
140{
141}
142
143impl<T> Transport for T where
144    T: Stream<Item = Result<RpcMessage, codec::Error>> + Sink<RpcMessage, Error = codec::Error>
145{
146}
147
148impl TransportProtocol {
149    /// Creates a transport for this protocol.
150    pub async fn connect(
151        self,
152        address: impl ToSocketAddrs,
153    ) -> Result<impl Transport, std::io::Error> {
154        let mut addresses = lookup_host(address)
155            .await
156            .expect("Invalid address to connect to");
157        let address = addresses
158            .next()
159            .expect("Couldn't resolve address to connect to");
160
161        let stream: futures::future::Either<_, _> = match self {
162            TransportProtocol::Udp => {
163                let socket = UdpSocket::bind(&"0.0.0.0:0").await?;
164
165                UdpFramed::new(socket, Codec)
166                    .with(move |message| future::ready(Ok((message, address))))
167                    .map_ok(|(message, _address)| message)
168                    .left_stream()
169            }
170            TransportProtocol::Tcp => {
171                let stream = TcpStream::connect(address).await?;
172
173                Framed::new(stream, Codec).right_stream()
174            }
175        };
176
177        Ok(stream)
178    }
179
180    /// Creates a [`ConnectionPool`] for this protocol.
181    pub async fn make_outgoing_connection_pool(
182        self,
183    ) -> Result<Box<dyn ConnectionPool>, std::io::Error> {
184        let pool: Box<dyn ConnectionPool> = match self {
185            Self::Udp => Box::new(UdpConnectionPool::new().await?),
186            Self::Tcp => Box::new(TcpConnectionPool::new()),
187        };
188        Ok(pool)
189    }
190
191    /// Runs a server for this protocol and the given message handler.
192    pub fn spawn_server<S>(
193        self,
194        address: impl ToSocketAddrs + Send + 'static,
195        state: S,
196        shutdown_signal: CancellationToken,
197        join_set: &mut JoinSet<()>,
198    ) -> ServerHandle
199    where
200        S: MessageHandler + Send + 'static,
201    {
202        let handle = match self {
203            Self::Udp => join_set.spawn_task(UdpServer::run(address, state, shutdown_signal)),
204            Self::Tcp => join_set.spawn_task(TcpServer::run(address, state, shutdown_signal)),
205        };
206        ServerHandle { handle }
207    }
208}
209
210/// An implementation of [`ConnectionPool`] based on UDP.
211struct UdpConnectionPool {
212    transport: UdpFramed<Codec>,
213}
214
215impl UdpConnectionPool {
216    async fn new() -> Result<Self, std::io::Error> {
217        let socket = UdpSocket::bind(&"0.0.0.0:0").await?;
218        let transport = UdpFramed::new(socket, Codec);
219        Ok(Self { transport })
220    }
221}
222
223impl ConnectionPool for UdpConnectionPool {
224    fn send_message_to<'a>(
225        &'a mut self,
226        message: RpcMessage,
227        address: &'a str,
228    ) -> future::BoxFuture<'a, Result<(), codec::Error>> {
229        Box::pin(async move {
230            let address = address.parse().map_err(std::io::Error::other)?;
231            self.transport.send((message, address)).await
232        })
233    }
234}
235
236/// Server implementation for UDP.
237pub struct UdpServer<State> {
238    handler: State,
239    udp_sink: SharedUdpSink,
240    udp_stream: SplitStream<UdpFramed<Codec>>,
241    active_handlers: HashMap<SocketAddr, TaskHandle<()>>,
242    join_set: JoinSet<()>,
243}
244
245/// Type alias for the outgoing endpoint of UDP messages.
246type SharedUdpSink = Arc<Mutex<SplitSink<UdpFramed<Codec>, (RpcMessage, SocketAddr)>>>;
247
248impl<State> UdpServer<State>
249where
250    State: MessageHandler + Send + 'static,
251{
252    /// Runs the UDP server implementation.
253    pub async fn run(
254        address: impl ToSocketAddrs,
255        state: State,
256        shutdown_signal: CancellationToken,
257    ) -> Result<(), std::io::Error> {
258        let mut server = Self::bind(address, state).await?;
259
260        loop {
261            tokio::select! { biased;
262                _ = shutdown_signal.cancelled() => {
263                    server.shutdown().await;
264                    return Ok(());
265                }
266                result = server.udp_stream.next() => match result {
267                    Some(Ok((message, peer))) => server.handle_message(message, peer),
268                    Some(Err(error)) => server.handle_error(error).await?,
269                    None => unreachable!("`UdpFramed` should never return `None`"),
270                },
271            }
272        }
273    }
274
275    /// Creates a [`UdpServer`] bound to the provided `address`, handling messages using the
276    /// provided `handler`.
277    async fn bind(address: impl ToSocketAddrs, handler: State) -> Result<Self, std::io::Error> {
278        let socket = UdpSocket::bind(address).await?;
279        let (udp_sink, udp_stream) = UdpFramed::new(socket, Codec).split();
280
281        Ok(UdpServer {
282            handler,
283            udp_sink: Arc::new(Mutex::new(udp_sink)),
284            udp_stream,
285            active_handlers: HashMap::new(),
286            join_set: JoinSet::new(),
287        })
288    }
289
290    /// Spawns a task to handle a single incoming message.
291    fn handle_message(&mut self, message: RpcMessage, peer: SocketAddr) {
292        let previous_task = self.active_handlers.remove(&peer);
293        let mut state = self.handler.clone();
294        let udp_sink = self.udp_sink.clone();
295
296        let new_task = self.join_set.spawn_task(async move {
297            if let Some(reply) = state.handle_message(message).await {
298                if let Some(task) = previous_task {
299                    if let Err(error) = task.await {
300                        warn!("Message handler task panicked: {}", error);
301                    }
302                }
303                let status = udp_sink.lock().await.send((reply, peer)).await;
304                if let Err(error) = status {
305                    error!("Failed to send query response: {}", error);
306                }
307            }
308        });
309
310        self.active_handlers.insert(peer, new_task);
311
312        if self.active_handlers.len() >= REAP_TASKS_THRESHOLD {
313            // Collect finished tasks to avoid leaking memory.
314            self.active_handlers.retain(|_, task| task.is_running());
315            self.join_set.reap_finished_tasks();
316        }
317    }
318
319    /// Handles an error while receiving a message.
320    async fn handle_error(&mut self, error: codec::Error) -> Result<(), std::io::Error> {
321        match error {
322            codec::Error::IoError(io_error) => {
323                error!("I/O error in UDP server: {io_error}");
324                self.shutdown().await;
325                Err(io_error)
326            }
327            other_error => {
328                warn!("Received an invalid message: {other_error}");
329                Ok(())
330            }
331        }
332    }
333
334    /// Gracefully shuts down the server, waiting for existing tasks to finish.
335    async fn shutdown(&mut self) {
336        let handlers = mem::take(&mut self.active_handlers);
337        let mut handler_results = handlers.into_values().collect::<FuturesUnordered<_>>();
338
339        while let Some(result) = handler_results.next().await {
340            if let Err(error) = result {
341                warn!("Message handler panicked: {}", error);
342            }
343        }
344
345        self.join_set.await_all_tasks().await;
346    }
347}
348
349/// An implementation of [`ConnectionPool`] based on TCP.
350struct TcpConnectionPool {
351    streams: HashMap<String, Framed<TcpStream, Codec>>,
352}
353
354impl TcpConnectionPool {
355    fn new() -> Self {
356        let streams = HashMap::new();
357        Self { streams }
358    }
359
360    async fn get_stream(
361        &mut self,
362        address: &str,
363    ) -> Result<&mut Framed<TcpStream, Codec>, io::Error> {
364        if !self.streams.contains_key(address) {
365            match TcpStream::connect(address).await {
366                Ok(s) => {
367                    self.streams
368                        .insert(address.to_string(), Framed::new(s, Codec));
369                }
370                Err(error) => {
371                    error!("Failed to open connection to {}: {}", address, error);
372                    return Err(error);
373                }
374            };
375        };
376        Ok(self.streams.get_mut(address).unwrap())
377    }
378}
379
380impl ConnectionPool for TcpConnectionPool {
381    fn send_message_to<'a>(
382        &'a mut self,
383        message: RpcMessage,
384        address: &'a str,
385    ) -> future::BoxFuture<'a, Result<(), codec::Error>> {
386        Box::pin(async move {
387            let stream = self.get_stream(address).await?;
388            let result = stream.send(message).await;
389            if result.is_err() {
390                self.streams.remove(address);
391            }
392            result
393        })
394    }
395}
396
397/// Server implementation for TCP.
398pub struct TcpServer<State> {
399    connection: Framed<TcpStream, Codec>,
400    handler: State,
401    shutdown_signal: CancellationToken,
402}
403
404impl<State> TcpServer<State>
405where
406    State: MessageHandler + Send + 'static,
407{
408    /// Runs the TCP server implementation.
409    ///
410    /// Listens for connections and spawns a task with a new [`TcpServer`] instance to serve that
411    /// client.
412    pub async fn run(
413        address: impl ToSocketAddrs,
414        handler: State,
415        shutdown_signal: CancellationToken,
416    ) -> Result<(), std::io::Error> {
417        let listener = TcpListener::bind(address).await?;
418
419        let accept_stream = stream::try_unfold(listener, |listener| async move {
420            let (socket, _) = listener.accept().await?;
421            Ok::<_, io::Error>(Some((socket, listener)))
422        });
423        let mut accept_stream = pin!(accept_stream);
424
425        let connection_shutdown_signal = shutdown_signal.child_token();
426        let mut join_set = JoinSet::new();
427        let mut reap_countdown = REAP_TASKS_THRESHOLD;
428
429        loop {
430            tokio::select! { biased;
431                _ = shutdown_signal.cancelled() => {
432                    join_set.await_all_tasks().await;
433                    return Ok(());
434                }
435                maybe_socket = accept_stream.next() => match maybe_socket {
436                    Some(Ok(socket)) => {
437                        let server = TcpServer::new_connection(
438                            socket,
439                            handler.clone(),
440                            connection_shutdown_signal.clone(),
441                        );
442                        join_set.spawn_task(server.serve());
443                        reap_countdown -= 1;
444                    }
445                    Some(Err(error)) => {
446                        join_set.await_all_tasks().await;
447                        return Err(error);
448                    }
449                    None => unreachable!(
450                        "The `accept_stream` should never finish unless there's an error",
451                    ),
452                },
453            }
454
455            if reap_countdown == 0 {
456                join_set.reap_finished_tasks();
457                reap_countdown = REAP_TASKS_THRESHOLD;
458            }
459        }
460    }
461
462    /// Creates a new [`TcpServer`] to serve a single connection established on the provided
463    /// [`TcpStream`].
464    fn new_connection(
465        tcp_stream: TcpStream,
466        handler: State,
467        shutdown_signal: CancellationToken,
468    ) -> Self {
469        TcpServer {
470            connection: Framed::new(tcp_stream, Codec),
471            handler,
472            shutdown_signal,
473        }
474    }
475
476    /// Serves a client through a single connection.
477    async fn serve(mut self) {
478        loop {
479            tokio::select! { biased;
480                _ = self.shutdown_signal.cancelled() => {
481                    let mut tcp_stream = self.connection.into_inner();
482                    if let Err(error) = tcp_stream.shutdown().await {
483                        let peer = tcp_stream
484                            .peer_addr()
485                            .map_or_else(|_| "an unknown peer".to_owned(), |address| address.to_string());
486                        warn!("Failed to close connection to {peer}: {error:?}");
487                    }
488                    return;
489                }
490                result = self.connection.next() => match result {
491                    Some(Ok(RpcMessage::SubscribeNotifications(chains))) => {
492                        self.handle_subscription(chains).await;
493                        return;
494                    }
495                    Some(Ok(RpcMessage::DownloadBlobs(blob_ids))) => {
496                        self.handle_download_blobs(blob_ids).await;
497                        return;
498                    }
499                    Some(Ok(message)) => self.handle_message(message).await,
500                    Some(Err(error)) => {
501                        Self::handle_error(&error);
502                        return;
503                    }
504                    None => break,
505                },
506            }
507        }
508    }
509
510    /// Handles a single request message from a client.
511    async fn handle_message(&mut self, message: RpcMessage) {
512        if let Some(reply) = self.handler.handle_message(message).await {
513            if let Err(error) = self.connection.send(reply).await {
514                error!("Failed to send query response: {error}");
515            }
516        }
517    }
518
519    /// Handles a notification subscription request by switching to streaming mode.
520    async fn handle_subscription(&mut self, chains: Vec<ChainId>) {
521        let Some(mut stream) = self.handler.handle_subscribe(chains).await else {
522            return;
523        };
524        loop {
525            tokio::select! { biased;
526                _ = self.shutdown_signal.cancelled() => break,
527                msg = stream.next() => match msg {
528                    Some(notification) => {
529                        if let Err(error) = self.connection.send(notification).await {
530                            error!("Failed to send notification: {error}");
531                            break;
532                        }
533                    }
534                    None => break,
535                }
536            }
537        }
538    }
539
540    /// Handles a batch blob download request by streaming one response per blob.
541    async fn handle_download_blobs(&mut self, blob_ids: Vec<BlobId>) {
542        let Some(mut stream) = self.handler.handle_download_blobs(blob_ids).await else {
543            return;
544        };
545        loop {
546            tokio::select! { biased;
547                _ = self.shutdown_signal.cancelled() => break,
548                msg = stream.next() => match msg {
549                    Some(message) => {
550                        if let Err(error) = self.connection.send(message).await {
551                            error!("Failed to send blob response: {error}");
552                            break;
553                        }
554                    }
555                    None => break,
556                }
557            }
558        }
559    }
560
561    /// Handles an error received while attempting to receive from the connection.
562    ///
563    /// Ignores a successful connection termination, while logging an unexpected connection
564    /// termination or any other error.
565    fn handle_error(error: &codec::Error) {
566        if !matches!(
567            error,
568            codec::Error::IoError(error)
569                if error.kind() == io::ErrorKind::UnexpectedEof
570                || error.kind() == io::ErrorKind::ConnectionReset
571        ) {
572            error!("Error while reading TCP stream: {error}");
573        }
574    }
575}
linera_rpc/simple/transport.rs

linera_rpc/simple/
transport.rs
