apollo_propeller: implement outbound substream with message batching

Author: sirandreww-starkware

Cargo.lock

@@ -14,6 +14,7 @@ apollo_protobuf.workspace = true
 asynchronous-codec.workspace = true
 futures.workspace = true
 libp2p.workspace = true
+prost.workspace = true
 reed-solomon-simd.workspace = true
 sha2.workspace = true
 thiserror.workspace = true

crates/apollo_propeller/Cargo.toml

@@ -16,6 +16,8 @@ use libp2p::swarm::handler::{
     FullyNegotiatedOutbound,
 };
 use libp2p::swarm::{Stream, StreamProtocol, SubstreamProtocol};
+use prost::Message;
+use tracing::warn;
 
 use crate::protocol::{PropellerCodec, PropellerProtocol};
 use crate::PropellerUnit;
@@ -37,34 +39,62 @@ pub enum HandlerIn {
 }
 
 /// Protocol Handler that manages substreams with a peer.
+///
+/// We use separate unidirectional substreams: outbound for sending and inbound for receiving.
+// TODO(AndrewL): Add this to the specs.
 pub struct Handler {
     /// Upgrade configuration for the propeller protocol.
     listen_protocol: PropellerProtocol,
     /// The single long-lived inbound substream.
     inbound_substream: Option<InboundSubstreamState>,
+    /// The single long-lived outbound substream.
+    outbound_substream: OutboundSubstreamState,
     /// Queue of messages to send.
     send_queue: VecDeque<ProtoUnit>,
     /// Queue of received messages to emit.
     receive_queue: VecDeque<PropellerUnit>,
+    /// Maximum wire message size for batching.
+    max_wire_message_size: usize,
 }
 
-/// State of the inbound substream, opened either by us or by the remote.
+/// State of the inbound substream, opened by the remote peer.
 enum InboundSubstreamState {
     /// Waiting for a message from the remote. The idle state for an inbound substream.
     WaitingInput(Framed<Stream, PropellerCodec>),
     /// The substream is being closed.
     Closing(Framed<Stream, PropellerCodec>),
 }
 
+/// State of the outbound substream, opened by us.
+enum OutboundSubstreamState {
+    /// No substream exists and no request is pending.
+    Idle,
+    /// A substream request has been sent but not yet negotiated.
+    Pending,
+    /// The substream is active and ready to send messages.
+    Active {
+        substream: Framed<Stream, PropellerCodec>,
+        /// True if we've sent data and need to flush the stream.
+        ///
+        /// In this implementation we use flush to ensure that the data was actually sent. Without
+        /// flushing periodically we cannot be certain what was sent and what failed when a failure
+        /// occurs.
+        // TODO(AndrewL): Discuss not flushing
+        should_flush: bool,
+    },
+}
+
 impl Handler {
     /// Builds a new [`Handler`].
     pub fn new(stream_protocol: StreamProtocol, max_wire_message_size: usize) -> Self {
         let protocol = PropellerProtocol::new(stream_protocol, max_wire_message_size);
         Handler {
             listen_protocol: protocol,
             inbound_substream: None,
+            outbound_substream: OutboundSubstreamState::Idle,
             send_queue: VecDeque::new(),
             receive_queue: VecDeque::new(),
+            max_wire_message_size,
         }
     }
 
@@ -153,6 +183,33 @@ impl Handler {
         }
     }
 
+    /// Create a batch of messages from the send queue that fits within max_wire_message_size.
+    fn create_message_batch(
+        send_queue: &mut VecDeque<ProtoUnit>,
+        max_wire_message_size: usize,
+    ) -> ProtoBatch {
+        if send_queue.is_empty() {
+            return ProtoBatch { batch: Vec::new() };
+        }
+
+        let mut batch = ProtoBatch { batch: vec![send_queue.pop_front().unwrap()] };
+        if batch.encoded_len() >= max_wire_message_size {
+            warn!("Propeller unit size exceeds max wire message size, sending will fail");
+        }
+
+        while let Some(msg) = send_queue.front() {
+            batch.batch.push(msg.clone());
+            if batch.encoded_len() < max_wire_message_size {
+                send_queue.pop_front();
+            } else {
+                batch.batch.pop();
+                break;
+            }
+        }
+
+        batch
+    }
+
     fn on_fully_negotiated_inbound(&mut self, substream: Framed<Stream, PropellerCodec>) {
         if self.inbound_substream.is_some() {
             // TODO(AndrewL): Either remove this warning or make it once every N ms.
@@ -165,25 +222,125 @@ impl Handler {
 
     fn on_fully_negotiated_outbound(
         &mut self,
-        _fully_negotiated_outbound: FullyNegotiatedOutbound<
+        fully_negotiated_outbound: FullyNegotiatedOutbound<
             <Handler as ConnectionHandler>::OutboundProtocol,
         >,
     ) {
-        // TODO(AndrewL): Implement outbound substream handling
-        todo!("Outbound substream handling not yet implemented")
+        if let OutboundSubstreamState::Active { should_flush, .. } = &self.outbound_substream {
+            if *should_flush {
+                tracing::warn!(
+                    "New outbound substream while existing substream has pending data, data may \
+                     be lost"
+                );
+            }
+        }
+
+        let substream = fully_negotiated_outbound.protocol;
+        self.outbound_substream = OutboundSubstreamState::Active { substream, should_flush: false };
     }
 
     fn poll_send(
         &mut self,
-        _cx: &mut Context<'_>,
+        cx: &mut Context<'_>,
     ) -> Poll<
         ConnectionHandlerEvent<
             <Handler as ConnectionHandler>::OutboundProtocol,
             (),
             <Handler as ConnectionHandler>::ToBehaviour,
         >,
     > {
-        // TODO(AndrewL): Implement outbound message sending
+        // No need to optimize requesting a sub-stream only when there are units to send since the
+        // streams are long-lived (We only open a sub-stream twice per connection assuming no errors
+        // occur).
+        // If we don't have an outbound substream and haven't requested one, request one
+        match &self.outbound_substream {
+            OutboundSubstreamState::Idle => {
+                self.outbound_substream = OutboundSubstreamState::Pending;
+                return Poll::Ready(ConnectionHandlerEvent::OutboundSubstreamRequest {
+                    protocol: SubstreamProtocol::new(self.listen_protocol.clone(), ()),
+                });
+            }
+            OutboundSubstreamState::Pending => {
+                return Poll::Pending;
+            }
+            _ => {}
+        }
+
+        loop {
+            let OutboundSubstreamState::Active { mut substream, mut should_flush } =
+                std::mem::replace(&mut self.outbound_substream, OutboundSubstreamState::Idle)
+            else {
+                unreachable!("outbound_substream is Active at the start of this loop");
+            };
+
+            if self.send_queue.is_empty() {
+                // Queue is empty, maybe we just need to flush the stream
+                if should_flush {
+                    match Sink::poll_flush(Pin::new(&mut substream), cx) {
+                        Poll::Ready(Ok(())) => {
+                            should_flush = false;
+                            self.outbound_substream =
+                                OutboundSubstreamState::Active { substream, should_flush };
+                            continue;
+                        }
+                        Poll::Ready(Err(e)) => {
+                            tracing::error!("Failed to flush outbound stream: {e}");
+                            return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(
+                                HandlerOut::SendError(e.to_string()),
+                            ));
+                        }
+                        Poll::Pending => {
+                            self.outbound_substream =
+                                OutboundSubstreamState::Active { substream, should_flush };
+                            break;
+                        }
+                    }
+                } else {
+                    self.outbound_substream =
+                        OutboundSubstreamState::Active { substream, should_flush };
+                    break;
+                }
+            } else {
+                match Sink::poll_ready(Pin::new(&mut substream), cx) {
+                    Poll::Ready(Ok(())) => {
+                        let message = Self::create_message_batch(
+                            &mut self.send_queue,
+                            self.max_wire_message_size,
+                        );
+                        match Sink::start_send(Pin::new(&mut substream), message) {
+                            Ok(()) => {
+                                // Try sending more messages if there are any
+                                should_flush = true;
+                                self.outbound_substream =
+                                    OutboundSubstreamState::Active { substream, should_flush };
+                                continue;
+                            }
+                            Err(e) => {
+                                // TODO(AndrewL): Units were lost, consider a re-try mechanism.
+                                tracing::error!("Failed to send message on outbound stream: {e}");
+                                return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(
+                                    HandlerOut::SendError(e.to_string()),
+                                ));
+                            }
+                        }
+                    }
+                    Poll::Ready(Err(e)) => {
+                        // TODO(AndrewL): Units were lost, consider a re-try mechanism.
+                        tracing::error!("Failed to send message on outbound stream: {e}");
+                        return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(
+                            HandlerOut::SendError(e.to_string()),
+                        ));
+                    }
+                    Poll::Pending => {
+                        // Not ready to send more messages yet
+                        self.outbound_substream =
+                            OutboundSubstreamState::Active { substream, should_flush };
+                        break;
+                    }
+                }
+            }
+        }
+
         Poll::Pending
     }
 
@@ -259,9 +416,18 @@ impl ConnectionHandler for Handler {
             ConnectionEvent::FullyNegotiatedOutbound(fully_negotiated_outbound) => {
                 self.on_fully_negotiated_outbound(fully_negotiated_outbound)
             }
-            _ => {
-                // TODO(AndrewL): Handle DialUpgradeError variants
+            ConnectionEvent::DialUpgradeError(_) => {
+                // Reset to Idle so we can request a new substream
+                if matches!(self.outbound_substream, OutboundSubstreamState::Pending) {
+                    self.outbound_substream = OutboundSubstreamState::Idle;
+                } else {
+                    tracing::error!(
+                        "Dial upgrade error but no pending substream found. (File a bug report if \
+                         you see this)"
+                    );
+                }
             }
+            _ => {}
         }
     }
 }