f much simpler design from @andrei-21

Author: TheBlueMatt

lightning/src/ln/peer_handler.rs

@@ -36,7 +36,7 @@ use crate::prelude::*;
 use crate::io;
 use alloc::collections::LinkedList;
 use crate::sync::{Arc, Mutex, MutexGuard, FairRwLock};
-use core::sync::atomic::{AtomicBool, AtomicU32, AtomicU8, Ordering};
+use core::sync::atomic::{AtomicBool, AtomicU32, AtomicI32, Ordering};
 use core::{cmp, hash, fmt, mem};
 use core::ops::Deref;
 use core::convert::Infallible;
@@ -682,12 +682,12 @@ pub struct PeerManager<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: D
 	///
 	/// Because the event handler calls into user code which may block, we don't want to block a
 	/// second thread waiting for another thread to handle events which is then blocked on user
-	/// code, so we store an atomic enum here with three states -
+	/// code, so we store an atomic counter here:
 	///  * 0 indicates no event processor is running
 	///  * 1 indicates an event processor is running
-	///  * 2 indicates an event processor is running but needs to start again from the top once it
-	///    finishes as another thread tried to start processing events but returned early.
-	event_processing_state: AtomicU8,
+	///  * > 1 indicates an event processor is running but needs to start again from the top once
+	///        it finishes as another thread tried to start processing events but returned early.
+	event_processing_state: AtomicI32,
 
 	/// Used to track the last value sent in a node_announcement "timestamp" field. We ensure this
 	/// value increases strictly since we don't assume access to a time source.
@@ -857,7 +857,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 			message_handler,
 			peers: FairRwLock::new(HashMap::new()),
 			node_id_to_descriptor: Mutex::new(HashMap::new()),
-			event_processing_state: AtomicU8::new(0),
+			event_processing_state: AtomicI32::new(0),
 			ephemeral_key_midstate,
 			peer_counter: AtomicCounter::new(),
 			gossip_processing_backlogged: AtomicBool::new(false),
@@ -1786,38 +1786,10 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 	/// [`ChannelManager::process_pending_htlc_forwards`]: crate::ln::channelmanager::ChannelManager::process_pending_htlc_forwards
 	/// [`send_data`]: SocketDescriptor::send_data
 	pub fn process_events(&self) {
-		if self.event_processing_state.compare_exchange(0, 1, Ordering::AcqRel, Ordering::Relaxed).is_err() {
-			// If we failed to take the happy path, there may be another thread running. Try to
-			// either swap up to tell the existing thread to go again or claim the running lock.
-			loop {
-				match self.event_processing_state.load(Ordering::Acquire) {
-					0 =>
-						if self.event_processing_state.compare_exchange(
-							0, 1, Ordering::AcqRel, Ordering::Relaxed
-						).is_ok() {
-							// The other thread finished, we now have the lock and should do a normal
-							// event processing run.
-							break;
-						},
-					1 =>
-						if self.event_processing_state.compare_exchange(
-							1, 2, Ordering::AcqRel, Ordering::Relaxed
-						).is_ok() {
-							// The running thread will go again, we're done
-							return;
-						},
-					2 => {
-						// There's already a thread going that will start over once it finishes, we
-						// don't need to do anything
-						return;
-					}
-					x => {
-						debug_assert!(false, "event_processing state was invalid: {}", x);
-						self.event_processing_state.store(1, Ordering::Release);
-						break;
-					}
-				}
-			}
+		if self.event_processing_state.fetch_add(1, Ordering::AcqRel) > 0 {
+			// If we're not the first event processor to get here, just return early, the increment
+			// we just did will be treated as "go around again" at the end.
+			return;
 		}
 
 		loop {
@@ -2152,12 +2124,10 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 				}
 			}
 
-			if self.event_processing_state.compare_exchange(1, 0, Ordering::AcqRel, Ordering::Relaxed).is_err() {
-				// If we're running we must either be in state 1 or 2, and if not state 1 we must
-				// now be in state 2 (go again)!
-				let go_again_swap =
-					self.event_processing_state.compare_exchange(2, 1, Ordering::AcqRel, Ordering::Relaxed);
-				debug_assert!(go_again_swap.is_ok());
+			if self.event_processing_state.fetch_add(-1, Ordering::AcqRel) != 1 {
+				// If another thread incremented the state while we were running we should go
+				// around again, but only once.
+				self.event_processing_state.store(1, Ordering::Release);
 				continue;
 			}
 			break;