From 83789a3076e986782af60ba32b0398414c1c82d7 Mon Sep 17 00:00:00 2001 From: Alex Auvolat Date: Mon, 7 Dec 2020 12:39:19 +0100 Subject: [PATCH] Call on_disconnected_handler earlier (as soon as .disconnect() is called) --- src/netapp.rs | 41 +++++++++++++++++++++++------------------ 1 file 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, 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