From 995b62607127d32cff34fb152f40ae76549d5e82 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Fri, 28 Feb 2020 11:55:34 -0500
Subject: [PATCH] Fix (and test) net-tokio outbound conns without a threaded
 env

---
 lightning-net-tokio/src/lib.rs | 42 ++++++++++++++++++++++++++--------
 1 file changed, 33 insertions(+), 9 deletions(-)

diff --git a/lightning-net-tokio/src/lib.rs b/lightning-net-tokio/src/lib.rs
index b04e93f74..c1f66d3fa 100644
--- a/lightning-net-tokio/src/lib.rs
+++ b/lightning-net-tokio/src/lib.rs
@@ -271,16 +271,32 @@ pub fn setup_inbound<CMH: ChannelMessageHandler + 'static>(peer_manager: Arc<pee
 ///
 /// See the module-level documentation for how to handle the event_notify mpsc::Sender.
 pub fn setup_outbound<CMH: ChannelMessageHandler + 'static>(peer_manager: Arc<peer_handler::PeerManager<SocketDescriptor, Arc<CMH>>>, event_notify: mpsc::Sender<()>, their_node_id: PublicKey, stream: TcpStream) -> impl std::future::Future<Output=()> {
-	let (reader, write_receiver, read_receiver, us) = Connection::new(event_notify, stream);
+	let (reader, mut write_receiver, read_receiver, us) = Connection::new(event_notify, stream);
 
 	let handle_opt = if let Ok(initial_send) = peer_manager.new_outbound_connection(their_node_id, SocketDescriptor::new(us.clone())) {
 		Some(tokio::spawn(async move {
-			if SocketDescriptor::new(us.clone()).send_data(&initial_send, true) != initial_send.len() {
-				// We should essentially always have enough room in a TCP socket buffer to send the
-				// initial 10s of bytes, if not, just give up as hopeless.
-				eprintln!("Failed to write first full message to socket!");
-				peer_manager.socket_disconnected(&SocketDescriptor::new(Arc::clone(&us)));
-			} else {
+			// We should essentially always have enough room in a TCP socket buffer to send the
+			// initial 10s of bytes, however, tokio running in single-threaded mode will always
+			// fail writes and wake us back up later to write, so we handle a Pending, but still
+			// expect to write the full set of bytes at once and use a relatively tight timeout.
+			if let Ok(Ok(())) = tokio::time::timeout(Duration::from_millis(100), async {
+				loop {
+					match SocketDescriptor::new(us.clone()).send_data(&initial_send, true) {
+						v if v == initial_send.len() => break Ok(()),
+						0 => {
+							write_receiver.recv().await;
+							// In theory we could check for if we've been instructed to disconnect
+							// the peer here, but its OK to just skip it - we'll check for it in
+							// schedule_read prior to any relevant calls into RL.
+						},
+						_ => {
+							eprintln!("Failed to write first full message to socket!");
+							peer_manager.socket_disconnected(&SocketDescriptor::new(Arc::clone(&us)));
+							break Err(());
+						}
+					}
+				}
+			}).await {
 				Connection::schedule_read(peer_manager, us, reader, read_receiver, write_receiver).await;
 			}
 		}))
@@ -521,8 +537,7 @@ mod tests {
 		}
 	}
 
-	#[tokio::test(threaded_scheduler)]
-	async fn basic_connection_test() {
+	async fn do_basic_connection_test() {
 		let secp_ctx = Secp256k1::new();
 		let a_key = SecretKey::from_slice(&[1; 32]).unwrap();
 		let b_key = SecretKey::from_slice(&[1; 32]).unwrap();
@@ -587,4 +602,13 @@ mod tests {
 		fut_a.await;
 		fut_b.await;
 	}
+
+	#[tokio::test(threaded_scheduler)]
+	async fn basic_threaded_connection_test() {
+		do_basic_connection_test().await;
+	}
+	#[tokio::test]
+	async fn basic_unthreaded_connection_test() {
+		do_basic_connection_test().await;
+	}
 }
-- 
2.39.5