Call on_disconnected_handler earlier (as soon as .disconnect() is called)

1 files changed, 23 insertions, 18 deletions

diff --git a/src/netapp.rs b/src/netapp.rs
index e903f44..bf9a3f0 100644
--- a/src/netapp.rs
+++ b/src/netapp.rs
@@ -238,26 +238,29 @@ impl NetApp {
 	/// Close the outgoing connection we have to a node specified by its public key,
 	/// if such a connection is currently open.
 	pub fn disconnect(self: &Arc<Self>, pk: &ed25519::PublicKey) {
-		// Don't disconnect from ourself (we aren't connected anyways)
-		// but pretend we did
-		if *pk == self.pubkey {
-			let pk = *pk;
-			let self2 = self.clone();
-			tokio::spawn(async move {
-				if let Some(h) = self2.on_disconnected_handler.load().as_ref() {
-					h(pk, false);
-				}
-			});
-			return;
+		// If pk is ourself, we're not supposed to have a connection open
+		if *pk != self.pubkey {
+			let conn = self.client_conns.read().unwrap().remove(pk);
+			if let Some(c) = conn {
+				debug!("Closing connection to {} ({})",
+					   hex::encode(c.peer_pk),
+					   c.remote_addr);
+				c.close();
+			} else {
+				return;
+			}
 		}
 
-		let conn = self.client_conns.read().unwrap().get(pk).cloned();
-		if let Some(c) = conn {
-			debug!("Closing connection to {} ({})",
-				   hex::encode(c.peer_pk),
-				   c.remote_addr);
-			c.close();
-		}
+		// call on_disconnected_handler immediately, since the connection
+		// was removed
+		// (if pk == self.pubkey, we pretend we disconnected)
+		let pk = *pk;
+		let self2 = self.clone();
+		tokio::spawn(async move {
+			if let Some(h) = self2.on_disconnected_handler.load().as_ref() {
+				h(pk, false);
+			}
+		});
 	}
 
 	/// Close the incoming connection from a certain client to us,
@@ -360,6 +363,8 @@ impl NetApp {
 				}
 			}
 		}
+		// else case: happens if connection was removed in .disconnect()
+		// in which case on_disconnected_handler was already called
 	}
 
 	/// Send a message to a remote host to which a client connection is already