Implement (v1 of) rate limiting for onion messages.

In this commit, we add business logic for checking if a peer's outbound buffer
has room for onion messages, and if so pulls them from an implementer of a new
trait, OnionMessageProvider.

Because other steps need to be taken before the onion_message module is made
public and the PeerManager is parameterized by an OnionMessageHandler (which
causes a lot of cascading boilerplate changes and affects the public api), we
implement OnionMessageProvider on () and store that in PeerManager as a
stopgap.

Makes sure channel messages are prioritized over OMs and we only go to write
them when there's sufficient space in the peer's buffer.

Author: valentinewallace

lightning/src/ln/peer_handler.rs

@@ -27,7 +27,7 @@ use ln::wire;
 use ln::wire::Encode;
 use routing::gossip::{NetworkGraph, P2PGossipSync};
 use util::atomic_counter::AtomicCounter;
-use util::events::{MessageSendEvent, MessageSendEventsProvider};
+use util::events::{MessageSendEvent, MessageSendEventsProvider, OnionMessageProvider};
 use util::logger::Logger;
 
 use prelude::*;
@@ -81,6 +81,11 @@ impl Deref for IgnoringMessageHandler {
 	fn deref(&self) -> &Self { self }
 }
 
+// TODO: implement `OnionMessageProvider` for `IgnoringMessageHandler` instead
+impl OnionMessageProvider for () {
+	fn next_onion_messages_for_peer(&self, _peer_node_id: PublicKey, _max_messages: usize) -> Vec<msgs::OnionMessage> { Vec::new() }
+}
+
 // Implement Type for Infallible, note that it cannot be constructed, and thus you can never call a
 // method that takes self for it.
 impl wire::Type for Infallible {
@@ -294,15 +299,23 @@ enum InitSyncTracker{
 /// forwarding gossip messages to peers altogether.
 const FORWARD_INIT_SYNC_BUFFER_LIMIT_RATIO: usize = 2;
 
+/// The ratio between buffer sizes at which we stop sending initial sync messages vs when we pause
+/// forwarding onion messages to peers altogether.
+const OM_BUFFER_LIMIT_RATIO: usize = 2;
+
 /// When the outbound buffer has this many messages, we'll stop reading bytes from the peer until
 /// we have fewer than this many messages in the outbound buffer again.
-/// We also use this as the target number of outbound gossip messages to keep in the write buffer,
-/// refilled as we send bytes.
+/// We also use this as the target number of outbound gossip and onion messages to keep in the write
+/// buffer, refilled as we send bytes.
 const OUTBOUND_BUFFER_LIMIT_READ_PAUSE: usize = 10;
 /// When the outbound buffer has this many messages, we'll simply skip relaying gossip messages to
 /// the peer.
 const OUTBOUND_BUFFER_LIMIT_DROP_GOSSIP: usize = OUTBOUND_BUFFER_LIMIT_READ_PAUSE * FORWARD_INIT_SYNC_BUFFER_LIMIT_RATIO;
 
+/// When the outbound buffer has this many messages, we won't poll for new onion messages for this
+/// peer.
+const OUTBOUND_BUFFER_LIMIT_PAUSE_OMS: usize = OUTBOUND_BUFFER_LIMIT_READ_PAUSE * OM_BUFFER_LIMIT_RATIO;
+
 /// If we've sent a ping, and are still awaiting a response, we may need to churn our way through
 /// the socket receive buffer before receiving the ping.
 ///
@@ -378,6 +391,14 @@ impl Peer {
 			InitSyncTracker::NodesSyncing(pk) => pk < node_id,
 		}
 	}
+
+	/// Returns the number of onion messages we can fit in this peer's buffer.
+	fn onion_message_buffer_slots_available(&self) -> usize {
+		cmp::min(
+			OUTBOUND_BUFFER_LIMIT_PAUSE_OMS.saturating_sub(self.pending_outbound_buffer.len()),
+			(BUFFER_DRAIN_MSGS_PER_TICK * OM_BUFFER_LIMIT_RATIO).saturating_sub(self.msgs_sent_since_pong))
+	}
+
 	/// Returns whether this peer's buffer is full and we should drop gossip messages.
 	fn buffer_full_drop_gossip(&self) -> bool {
 		if self.pending_outbound_buffer.len() > OUTBOUND_BUFFER_LIMIT_DROP_GOSSIP
@@ -432,6 +453,8 @@ pub struct PeerManager<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: De
 		L::Target: Logger,
 		CMH::Target: CustomMessageHandler {
 	message_handler: MessageHandler<CM, RM>,
+	onion_message_handler: (), // TODO: add `OnionMessageHandler` trait to
+	                           // `MessageHandler` and get rid of this
 	/// Connection state for each connected peer - we have an outer read-write lock which is taken
 	/// as read while we're doing processing for a peer and taken write when a peer is being added
 	/// or removed.
@@ -587,6 +610,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
 
 		PeerManager {
 			message_handler,
+			onion_message_handler: (),
 			peers: FairRwLock::new(HashMap::new()),
 			node_id_to_descriptor: Mutex::new(HashMap::new()),
 			event_processing_lock: Mutex::new(()),
@@ -797,8 +821,12 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
 	/// ready to call `[write_buffer_space_avail`] again if a write call generated here isn't
 	/// sufficient!
 	///
+	/// If any bytes are written, [`process_events`] should be called afterwards.
+	//  TODO: why?
+	///
 	/// [`send_data`]: SocketDescriptor::send_data
 	/// [`write_buffer_space_avail`]: PeerManager::write_buffer_space_avail
+	/// [`process_events`]: PeerManager::process_events
 	pub fn write_buffer_space_avail(&self, descriptor: &mut Descriptor) -> Result<(), PeerHandleError> {
 		let peers = self.peers.read().unwrap();
 		match peers.get(descriptor) {
@@ -1387,6 +1415,8 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
 	/// You don't have to call this function explicitly if you are using [`lightning-net-tokio`]
 	/// or one of the other clients provided in our language bindings.
 	///
+	/// Note that this method should be called again if any bytes are written.
+	///
 	/// Note that if there are any other calls to this function waiting on lock(s) this may return
 	/// without doing any work. All available events that need handling will be handled before the
 	/// other calls return.
@@ -1637,6 +1667,22 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
 
 			for (descriptor, peer_mutex) in peers.iter() {
 				self.do_attempt_write_data(&mut (*descriptor).clone(), &mut *peer_mutex.lock().unwrap());
+
+				// Only see if we have room for onion messages after we've written all channel messages, to
+				// ensure they take priority.
+				let (peer_node_id, om_buffer_slots_avail) = {
+					let peer = peer_mutex.lock().unwrap();
+					if let Some(peer_node_id) = peer.their_node_id {
+						(Some(peer_node_id.clone()), peer.onion_message_buffer_slots_available())
+					} else { (None, 0) }
+				};
+				if peer_node_id.is_some() && om_buffer_slots_avail > 0 {
+					for event in self.onion_message_handler.next_onion_messages_for_peer(
+						peer_node_id.unwrap(), om_buffer_slots_avail)
+					{
+						// TODO: forward onion message to peer
+					}
+				}
 			}
 		}
 		if !peers_to_disconnect.is_empty() {

lightning/src/onion_message/messenger.rs

@@ -21,8 +21,10 @@ use ln::onion_utils;
 use super::blinded_route::{BlindedRoute, ForwardTlvs, ReceiveTlvs};
 use super::packet::{BIG_PACKET_HOP_DATA_LEN, ForwardControlTlvs, Packet, Payload, ReceiveControlTlvs, SMALL_PACKET_HOP_DATA_LEN};
 use super::utils;
+use util::events::OnionMessageProvider;
 use util::logger::Logger;
 
+use core::mem;
 use core::ops::Deref;
 use sync::{Arc, Mutex};
 use prelude::*;
@@ -278,6 +280,25 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
 	}
 }
 
+impl<Signer: Sign, K: Deref, L: Deref> OnionMessageProvider for OnionMessenger<Signer, K, L>
+	where K::Target: KeysInterface<Signer = Signer>,
+	      L::Target: Logger,
+{
+	fn next_onion_messages_for_peer(&self, peer_node_id: PublicKey, max_messages: usize) -> Vec<msgs::OnionMessage> {
+		let mut pending_msgs = self.pending_messages.lock().unwrap();
+		if let Some(msgs) = pending_msgs.get_mut(&peer_node_id) {
+			if max_messages >= msgs.len() {
+				let mut peer_pending_msgs = Vec::new();
+				mem::swap(msgs, &mut peer_pending_msgs);
+				return peer_pending_msgs
+			} else {
+				return msgs.split_off(max_messages)
+			}
+		}
+		Vec::new()
+	}
+}
+
 // TODO: parameterize the below Simple* types with OnionMessenger and handle the messages it
 // produces
 /// Useful for simplifying the parameters of [`SimpleArcChannelManager`] and

lightning/src/util/events.rs

@@ -1195,6 +1195,14 @@ pub trait MessageSendEventsProvider {
 	fn get_and_clear_pending_msg_events(&self) -> Vec<MessageSendEvent>;
 }
 
+/// A trait indicating an object may generate onion message send events
+//  TODO: make this `pub` when we make onion message functionality public
+pub(crate) trait OnionMessageProvider {
+	/// Gets up to `max_messages` pending onion message events.
+	//  TODO: make this return `MessageSendEvent`s rather than `msgs::OnionMessage`s
+	fn next_onion_messages_for_peer(&self, peer_node_id: PublicKey, max_messages: usize) -> Vec<msgs::OnionMessage>;
+}
+
 /// A trait indicating an object may generate events.
 ///
 /// Events are processed by passing an [`EventHandler`] to [`process_pending_events`].