Do not close connections immediately on close signal, await for remaining responses

1 files changed, 22 insertions, 14 deletions

diff --git a/src/netapp.rs b/src/netapp.rs
index bf9a3f0..967105e 100644
--- a/src/netapp.rs
+++ b/src/netapp.rs
@@ -53,7 +53,7 @@ pub struct NetApp {
 
 	server_conns: RwLock<HashMap<ed25519::PublicKey, Arc<ServerConn>>>,
 	client_conns: RwLock<HashMap<ed25519::PublicKey, Arc<ClientConn>>>,
-	
+
 	pub(crate) msg_handlers: ArcSwap<HashMap<MessageKind, Arc<Handler>>>,
 	on_connected_handler:
 		ArcSwapOption<Box<dyn Fn(ed25519::PublicKey, SocketAddr, bool) + Send + Sync>>,
@@ -133,18 +133,22 @@ impl NetApp {
 	/// been successfully established. Do not set this if using a peering strategy,
 	/// as the peering strategy will need to set this itself.
 	pub fn on_connected<F>(&self, handler: F)
-		where F: Fn(ed25519::PublicKey, SocketAddr, bool) + Sized + Send + Sync + 'static
-		{
-		self.on_connected_handler.store(Some(Arc::new(Box::new(handler))));
+	where
+		F: Fn(ed25519::PublicKey, SocketAddr, bool) + Sized + Send + Sync + 'static,
+	{
+		self.on_connected_handler
+			.store(Some(Arc::new(Box::new(handler))));
 	}
 
 	/// Set the handler to be called when an existing connection (incoming or outgoing) has
 	/// been closed by either party. Do not set this if using a peering strategy,
 	/// as the peering strategy will need to set this itself.
 	pub fn on_disconnected<F>(&self, handler: F)
-		where F: Fn(ed25519::PublicKey, bool) + Sized + Send + Sync + 'static
-		{
-		self.on_disconnected_handler.store(Some(Arc::new(Box::new(handler))));
+	where
+		F: Fn(ed25519::PublicKey, bool) + Sized + Send + Sync + 'static,
+	{
+		self.on_disconnected_handler
+			.store(Some(Arc::new(Box::new(handler))));
 	}
 
 	/// Add a handler for a certain message type. Note that only one handler
@@ -240,11 +244,13 @@ impl NetApp {
 	pub fn disconnect(self: &Arc<Self>, pk: &ed25519::PublicKey) {
 		// 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);
+			let conn = self.client_conns.write().unwrap().remove(pk);
 			if let Some(c) = conn {
-				debug!("Closing connection to {} ({})",
-					   hex::encode(c.peer_pk),
-					   c.remote_addr);
+				debug!(
+					"Closing connection to {} ({})",
+					hex::encode(c.peer_pk),
+					c.remote_addr
+				);
 				c.close();
 			} else {
 				return;
@@ -268,9 +274,11 @@ impl NetApp {
 	pub fn server_disconnect(self: &Arc<Self>, pk: &ed25519::PublicKey) {
 		let conn = self.server_conns.read().unwrap().get(pk).cloned();
 		if let Some(c) = conn {
-			debug!("Closing incoming connection from {} ({})",
-				   hex::encode(c.peer_pk),
-				   c.remote_addr);
+			debug!(
+				"Closing incoming connection from {} ({})",
+				hex::encode(c.peer_pk),
+				c.remote_addr
+			);
 			c.close();
 		}
 	}