From 76402184887f2b2b34baa27de395e3cd806a3565 Mon Sep 17 00:00:00 2001
From: Ivan <ivan.zderadicka@gmail.com>
Date: Sun, 29 Sep 2024 22:28:36 +0200
Subject: [PATCH] Make outgoing addr Option

---
 crates/librqbit/src/torrent_state/live/mod.rs |  3 +-
 .../src/torrent_state/live/peer/mod.rs        | 30 ++++++++-----------
 .../src/torrent_state/live/peers/mod.rs       |  2 +-
 3 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/crates/librqbit/src/torrent_state/live/mod.rs b/crates/librqbit/src/torrent_state/live/mod.rs
index fced756..95a4681 100644
--- a/crates/librqbit/src/torrent_state/live/mod.rs
+++ b/crates/librqbit/src/torrent_state/live/mod.rs
@@ -66,7 +66,6 @@ use librqbit_core::{
     torrent_metainfo::TorrentMetaV1Info,
 };
 use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
-use peer::OutgoingAddressType;
 use peer_binary_protocol::{
     extended::{
         handshake::ExtendedHandshake, ut_metadata::UtMetadata, ut_pex::UtPex, ExtendedMessage,
@@ -932,7 +931,7 @@ impl<'a> PeerConnectionHandler for &'a PeerHandler {
                 self.state
                     .peers
                     .with_peer_mut(self.addr, "update outgoing addr", |peer| {
-                        peer.outgoing_address = OutgoingAddressType::Known(outgoing_addr)
+                        peer.outgoing_address = Some(outgoing_addr)
                     });
             }
         }
diff --git a/crates/librqbit/src/torrent_state/live/peer/mod.rs b/crates/librqbit/src/torrent_state/live/peer/mod.rs
index 90d32fe..b0895f8 100644
--- a/crates/librqbit/src/torrent_state/live/peer/mod.rs
+++ b/crates/librqbit/src/torrent_state/live/peer/mod.rs
@@ -18,19 +18,11 @@ pub(crate) type InflightRequest = ChunkInfo;
 pub(crate) type PeerRx = UnboundedReceiver<WriterRequest>;
 pub(crate) type PeerTx = UnboundedSender<WriterRequest>;
 
-#[derive(Debug, Default)]
-pub(crate) enum OutgoingAddressType {
-    #[default]
-    Default,
-    None,
-    Known(SocketAddr),
-}
-
 #[derive(Debug, Default)]
 pub(crate) struct Peer {
     pub state: PeerStateNoMut,
     pub stats: stats::atomic::PeerStats,
-    pub outgoing_address: OutgoingAddressType,
+    pub outgoing_address: Option<SocketAddr>,
 }
 
 impl Peer {
@@ -46,7 +38,14 @@ impl Peer {
         Self {
             state,
             stats: Default::default(),
-            outgoing_address: OutgoingAddressType::None,
+            outgoing_address: None,
+        }
+    }
+
+    pub fn new_with_outgoing_address(addr: SocketAddr) -> Self {
+        Self {
+            outgoing_address: Some(addr),
+            ..Default::default()
         }
     }
 
@@ -57,18 +56,15 @@ impl Peer {
     ) -> Option<SocketAddr> {
         if let PeerState::NotNeeded = self.state.get() {
             match self.outgoing_address {
-                OutgoingAddressType::Default => {
-                    self.state.set(PeerState::Queued, counters);
-                    Some(known_address)
-                }
-                OutgoingAddressType::None => None,
-                OutgoingAddressType::Known(socket_addr) => {
+                None => None,
+                Some(socket_addr) => {
                     if known_address == socket_addr {
                         self.state.set(PeerState::Queued, counters);
                     } else {
                         debug!(
                             peer = known_address.to_string(),
-                            "peer will by retried on  different address {}", socket_addr
+                            outgoing_addr = socket_addr.to_string(),
+                            "peer will by retried on different address",
                         );
                     }
                     Some(socket_addr)
diff --git a/crates/librqbit/src/torrent_state/live/peers/mod.rs b/crates/librqbit/src/torrent_state/live/peers/mod.rs
index 03c6c10..110df73 100644
--- a/crates/librqbit/src/torrent_state/live/peers/mod.rs
+++ b/crates/librqbit/src/torrent_state/live/peers/mod.rs
@@ -43,7 +43,7 @@ impl PeerStates {
         match self.states.entry(addr) {
             Entry::Occupied(_) => None,
             Entry::Vacant(vac) => {
-                vac.insert(Default::default());
+                vac.insert(Peer::new_with_outgoing_address(addr));
                 atomic_inc(&self.stats.queued);
                 atomic_inc(&self.session_stats.peers.queued);