From 24ec6fffd79d9f6d98120ec9985a17098ce0675c Mon Sep 17 00:00:00 2001
From: draoi <redacted@noreply.com>
Date: Tue, 2 Apr 2024 11:03:31 +0200
Subject: [PATCH] doc: add cautionary comments about unregister()

---
 src/net/hosts.rs                  | 7 +++++++
 src/net/session/manual_session.rs | 2 +-
 src/net/session/mod.rs            | 2 +-
 src/net/session/refine_session.rs | 8 ++++----
 4 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/net/hosts.rs b/src/net/hosts.rs
index 46d68d500..f293fa3e5 100644
--- a/src/net/hosts.rs
+++ b/src/net/hosts.rs
@@ -849,6 +849,9 @@ impl Hosts {
 
             addrs_len += i + 1;
             self.container.store_or_update(color.clone(), addr.clone(), *last_seen).await;
+
+            // Free up this peer for usage by other parts of the code base.
+            // This is a safe since the hostlist modification is now complete.
             self.unregister(addr).await;
         }
 
@@ -929,6 +932,10 @@ impl Hosts {
     /// Remove a host from the HostRegistry. Must be called after move(), when the refinery
     /// process fails, or when a channel stops. Prevents hosts from getting trapped in the
     /// HostState logical machinery.
+    ///
+    /// NOTE: Misuse of this call is dangerous since it frees up the peer to be used by
+    /// the refinery or outbound connect loop, and may result in invalid states. It should
+    /// only be called when it is completely safe to do so.
     pub async fn unregister(&self, addr: &Url) {
         self.registry.write().await.remove(addr);
         debug!(target: "net::hosts::unregister()", "Removed {} from HostRegistry", addr);
diff --git a/src/net/session/manual_session.rs b/src/net/session/manual_session.rs
index ef0b997c9..22a8c356a 100644
--- a/src/net/session/manual_session.rs
+++ b/src/net/session/manual_session.rs
@@ -153,7 +153,7 @@ impl ManualSession {
                             );
 
                             // Stop tracking this peer, to avoid it getting stuck in the Connect
-                            // state.
+                            // state. This is safe since we have failed to connect at this point.
                             self.p2p().hosts().unregister(&addr).await;
                         }
                     }
diff --git a/src/net/session/mod.rs b/src/net/session/mod.rs
index 01492cea6..afce7738e 100644
--- a/src/net/session/mod.rs
+++ b/src/net/session/mod.rs
@@ -78,7 +78,7 @@ pub async fn remove_sub_on_stop(p2p: P2pPtr, channel: ChannelPtr, type_id: Sessi
         p2p.hosts().move_host(addr, last_seen, HostColor::Grey).await.unwrap();
     }
 
-    // Remove channel from p2p
+    // Remove channel from the HostRegistry. Free up this addr for any future operation.
     p2p.hosts().unregister(channel.address()).await;
 
     debug!(target: "net::session::remove_sub_on_stop()", "[END]");
diff --git a/src/net/session/refine_session.rs b/src/net/session/refine_session.rs
index 0a5a0b01e..10351f364 100644
--- a/src/net/session/refine_session.rs
+++ b/src/net/session/refine_session.rs
@@ -258,10 +258,8 @@ impl GreylistRefinery {
                         );
 
                         // Remove this entry from HostRegistry to avoid this host getting
-                        // stuck in the Refining state.
-                        //
-                        // It is not necessary to call this when the refinery passes, since the
-                        // state will be changed to Connected.
+                        // stuck in the Refining state. This is a safe since the hostlist
+                        // modification is now complete.
                         hosts.unregister(url).await;
 
                         continue
@@ -275,6 +273,8 @@ impl GreylistRefinery {
 
                     // Add to the whitelist and remove from the greylist.
                     hosts.move_host(url, last_seen, HostColor::White).await.unwrap();
+
+                    // When move is complete we can safely stop tracking this peer.
                     hosts.unregister(url).await;
 
                     debug!(target: "net::refinery", "GreylistRefinery complete!");