Clean up socket + durable IO (#7)

README.md

@@ -20,14 +20,20 @@ It was built because we needed a Rust-native messaging library like those above.
 
 - [ ] Multiple socket types
  - [x] Request/Reply
- - [ ] Channel
  - [ ] Publish/Subscribe
+ - [ ] Channel
  - [ ] Push/Pull
  - [ ] Survey/Respond
 - [ ] Stats (RTT, throughput, packet drops etc.)
-- [ ] Durable transports (built-in retries and reconnections)
+ - [x] Request/Reply basic stats
 - [ ] Queuing
-- [ ] Pluggable transport layer (TCP, UDP, QUIC etc.)
+- [ ] Pluggable transport layer
+ - [x] TCP
+ - [ ] TLS
+ - [ ] IPC
+ - [ ] UDP
+ - [ ] Inproc
+- [x] Durable IO abstraction (built-in retries and reconnections)
 - [ ] Simulation modes with [Turmoil](https:/tokio-rs/turmoil)
 
 ## Socket Types
@@ -65,6 +71,8 @@ async fn main() {
  println!("Response: {:?}", res);
 }
 ```
+## MSRV
+The minimum supported Rust version is 1.70.
 
 ## Contributions & Bug Reports
 

msg-socket/src/rep/driver.rs

@@ -0,0 +1,268 @@
+use bytes::Bytes;
+use futures::{Future, SinkExt, Stream, StreamExt};
+use std::{
+ collections::VecDeque,
+ net::SocketAddr,
+ pin::Pin,
+ sync::Arc,
+ task::{Context, Poll},
+};
+use tokio::{
+ io::{AsyncRead, AsyncWrite},
+ sync::{mpsc, oneshot},
+ task::JoinSet,
+};
+use tokio_stream::{StreamMap, StreamNotifyClose};
+use tokio_util::codec::Framed;
+
+use crate::{rep::SocketState, Authenticator, RepError, Request};
+use msg_transport::ServerTransport;
+use msg_wire::{auth, reqrep};
+
+pub(crate) struct PeerState<T: AsyncRead + AsyncWrite> {
+ pending_requests: JoinSet<Option<(u32, Bytes)>>,
+ conn: Framed<T, reqrep::Codec>,
+ addr: SocketAddr,
+ egress_queue: VecDeque<reqrep::Message>,
+ state: Arc<SocketState>,
+}
+
+pub(crate) struct RepDriver<T: ServerTransport> {
+ /// The server transport used to accept incoming connections.
+ pub(crate) transport: T,
+ /// The reply socket state, shared with the socket front-end.
+ pub(crate) state: Arc<SocketState>,
+ /// [`StreamMap`] of connected peers. The key is the peer's address.
+ /// Note that when the [`PeerState`] stream ends, it will be silently removed
+ /// from this map.
+ pub(crate) peer_states: StreamMap<SocketAddr, StreamNotifyClose<PeerState<T::Io>>>,
+ /// Sender to the socket front-end. Used to notify the socket of incoming requests.
+ pub(crate) to_socket: mpsc::Sender<Request>,
+ /// Optional connection authenticator.
+ pub(crate) auth: Option<Arc<dyn Authenticator>>,
+ /// A joinset of authentication tasks.
+ pub(crate) auth_tasks: JoinSet<Result<AuthResult<T::Io>, RepError>>,
+}
+
+pub(crate) struct AuthResult<S: AsyncRead + AsyncWrite> {
+ id: Bytes,
+ addr: SocketAddr,
+ stream: S,
+}
+
+impl<T: ServerTransport> Future for RepDriver<T> {
+ type Output = Result<(), RepError>;
+
+ fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ let this = self.get_mut();
+
+ loop {
+ if let Poll::Ready(Some((peer, msg))) = this.peer_states.poll_next_unpin(cx) {
+ match msg {
+ Some(Ok(request)) => {
+ tracing::debug!("Received request from peer {}", peer);
+ this.state.stats.increment_rx(request.msg().len());
+ let _ = this.to_socket.try_send(request);
+ }
+ Some(Err(e)) => {
+ tracing::error!("Error receiving message from peer {}: {:?}", peer, e);
+ }
+ None => {
+ tracing::debug!("Peer {} disconnected", peer);
+ this.state.stats.decrement_active_clients();
+ }
+ }
+
+ continue;
+ }
+
+ if let Poll::Ready(Some(Ok(auth))) = this.auth_tasks.poll_join_next(cx) {
+ match auth {
+ Ok(auth) => {
+ // Run custom authenticator
+ tracing::debug!("Authentication passed for {:?} ({})", auth.id, auth.addr);
+ this.state.stats.increment_active_clients();
+
+ this.peer_states.insert(
+ auth.addr,
+ StreamNotifyClose::new(PeerState {
+ pending_requests: JoinSet::new(),
+ conn: Framed::new(auth.stream, reqrep::Codec::new()),
+ addr: auth.addr,
+ // TODO: pre-allocate according to some options
+ egress_queue: VecDeque::with_capacity(64),
+ state: Arc::clone(&this.state),
+ }),
+ );
+ }
+ Err(e) => {
+ tracing::error!("Error authenticating client: {:?}", e);
+ }
+ }
+
+ continue;
+ }
+
+ // Poll the transport for new incoming connections
+ match this.transport.poll_accept(cx) {
+ Poll::Ready(Ok((stream, addr))) => {
+ // If authentication is enabled, start the authentication process
+ if let Some(ref auth) = this.auth {
+ let authenticator = Arc::clone(auth);
+ tracing::debug!("New connection from {}, authenticating", addr);
+ this.auth_tasks.spawn(async move {
+ let mut conn = Framed::new(stream, auth::Codec::new_server());
+
+ tracing::debug!("Waiting for auth");
+ // Wait for the response
+ let auth = conn
+ .next()
+ .await
+ .ok_or(RepError::SocketClosed)?
+ .map_err(|e| RepError::Auth(e.to_string()))?;
+
+ tracing::debug!("Auth received: {:?}", auth);
+
+ let auth::Message::Auth(id) = auth else {
+ conn.send(auth::Message::Reject).await?;
+ conn.flush().await?;
+ conn.close().await?;
+ return Err(RepError::Auth("Invalid auth message".to_string()));
+ };
+
+ // If authentication fails, send a reject message and close the connection
+ if !authenticator.authenticate(&id) {
+ conn.send(auth::Message::Reject).await?;
+ conn.flush().await?;
+ conn.close().await?;
+ return Err(RepError::Auth("Authentication failed".to_string()));
+ }
+
+ // Send ack
+ conn.send(auth::Message::Ack).await?;
+ conn.flush().await?;
+
+ Ok(AuthResult {
+ id,
+ addr,
+ stream: conn.into_inner(),
+ })
+ });
+ } else {
+ this.state.stats.increment_active_clients();
+ this.peer_states.insert(
+ addr,
+ StreamNotifyClose::new(PeerState {
+ pending_requests: JoinSet::new(),
+ conn: Framed::new(stream, reqrep::Codec::new()),
+ addr,
+ // TODO: pre-allocate according to some options
+ egress_queue: VecDeque::with_capacity(64),
+ state: Arc::clone(&this.state),
+ }),
+ );
+
+ tracing::debug!("New connection from {}", addr);
+ }
+
+ continue;
+ }
+ Poll::Ready(Err(e)) => {
+ // Errors here are usually about `WouldBlock`
+ tracing::error!("Error accepting connection: {:?}", e);
+
+ continue;
+ }
+ Poll::Pending => {}
+ }
+
+ return Poll::Pending;
+ }
+ }
+}
+
+impl<T: AsyncRead + AsyncWrite + Unpin> Stream for PeerState<T> {
+ type Item = Result<Request, RepError>;
+
+ /// Advances the state of the peer.
+ fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ let this = self.get_mut();
+
+ loop {
+ // Flush any messages on the outgoing buffer
+ let _ = this.conn.poll_flush_unpin(cx);
+
+ // Then, try to drain the egress queue.
+ if this.conn.poll_ready_unpin(cx).is_ready() {
+ if let Some(msg) = this.egress_queue.pop_front() {
+ let msg_len = msg.size();
+ match this.conn.start_send_unpin(msg) {
+ Ok(_) => {
+ this.state.stats.increment_tx(msg_len);
+ // We might be able to send more queued messages
+ continue;
+ }
+ Err(e) => {
+ tracing::error!("Failed to send message to socket: {:?}", e);
+ // End this stream as we can't send any more messages
+ return Poll::Ready(None);
+ }
+ }
+ }
+ }
+
+ // Then we check for completed requests, and push them onto the egress queue.
+ match this.pending_requests.poll_join_next(cx) {
+ Poll::Ready(Some(Ok(Some((id, payload))))) => {
+ let msg = reqrep::Message::new(id, payload);
+ this.egress_queue.push_back(msg);
+
+ continue;
+ }
+ Poll::Ready(Some(Ok(None))) => {
+ tracing::error!("Failed to respond to request");
+ this.state.stats.increment_failed_requests();
+
+ continue;
+ }
+ Poll::Ready(Some(Err(e))) => {
+ tracing::error!("Error receiving response: {:?}", e);
+ this.state.stats.increment_failed_requests();
+
+ continue;
+ }
+ _ => {}
+ }
+
+ // Finally we accept incoming requests from the peer.
+ match this.conn.poll_next_unpin(cx) {
+ Poll::Ready(Some(result)) => {
+ tracing::trace!("Received message from peer {}: {:?}", this.addr, result);
+ let msg = result?;
+ let msg_id = msg.id();
+
+ let (tx, rx) = oneshot::channel();
+
+ // Spawn a task to listen for the response. On success, return message ID and response.
+ this.pending_requests
+ .spawn(async move { rx.await.ok().map(|res| (msg_id, res)) });
+
+ let request = Request {
+ source: this.addr,
+ response: tx,
+ msg: msg.into_payload(),
+ };
+
+ return Poll::Ready(Some(Ok(request)));
+ }
+ Poll::Ready(None) => {
+ tracing::debug!("Connection closed");
+ return Poll::Ready(None);
+ }
+ Poll::Pending => {}
+ }
+
+ return Poll::Pending;
+ }
+ }
+}