From 439314372b3b9d45744f2844ca6a456506be0f6d Mon Sep 17 00:00:00 2001
From: Bo QIU <35757521+boqiu@users.noreply.github.com>
Date: Thu, 27 Jun 2024 01:51:37 +0800
Subject: [PATCH] Remove file location cache for connection timeout peer (#96)

---
 .../src/file_location_cache.rs                 | 18 ++++++++++++++++++
 node/sync/src/controllers/peers.rs             | 17 ++++++++++++++++-
 node/sync/src/controllers/serial.rs            |  2 +-
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/node/file_location_cache/src/file_location_cache.rs b/node/file_location_cache/src/file_location_cache.rs
index 0a7a6ad..dc85fad 100644
--- a/node/file_location_cache/src/file_location_cache.rs
+++ b/node/file_location_cache/src/file_location_cache.rs
@@ -126,6 +126,12 @@ impl AnnouncementCache {
         let result = self.items.values().cloned().collect();
         (result, collected)
     }
+
+    /// Removes announcement for the specified `peer_id` if any.
+    fn remove(&mut self, peer_id: &PeerId) -> Option<SignedAnnounceFile> {
+        self.priorities.remove(peer_id)?;
+        self.items.remove(peer_id)
+    }
 }
 
 /// Caches announcements for different files.
@@ -230,6 +236,14 @@ impl FileCache {
         self.update_on_announcement_cache_changed(&tx_id, collected);
         Some(result)
     }
+
+    /// Removes the announcement of specified file by `tx_id` and `peer_id`.
+    fn remove(&mut self, tx_id: &TxID, peer_id: &PeerId) -> Option<SignedAnnounceFile> {
+        let item = self.files.get_mut(tx_id)?;
+        let result = item.remove(peer_id)?;
+        self.update_on_announcement_cache_changed(tx_id, 1);
+        Some(result)
+    }
 }
 
 #[derive(Default)]
@@ -281,6 +295,10 @@ impl FileLocationCache {
         self.cache.lock().all(tx_id).unwrap_or_default()
     }
 
+    pub fn remove(&self, tx_id: &TxID, peer_id: &PeerId) -> Option<SignedAnnounceFile> {
+        self.cache.lock().remove(tx_id, peer_id)
+    }
+
     /// TODO: Trigger chunk_pool/sync to reconstruct if it changes?
     pub fn insert_peer_config(
         &self,
diff --git a/node/sync/src/controllers/peers.rs b/node/sync/src/controllers/peers.rs
index 645640d..efc6442 100644
--- a/node/sync/src/controllers/peers.rs
+++ b/node/sync/src/controllers/peers.rs
@@ -1,5 +1,7 @@
+use file_location_cache::FileLocationCache;
 use network::{Multiaddr, PeerAction, PeerId};
 use rand::seq::IteratorRandom;
+use shared_types::TxID;
 use std::cmp::Ordering;
 use std::collections::{BTreeSet, HashMap};
 use std::sync::Arc;
@@ -46,13 +48,19 @@ impl PeerInfo {
 pub struct SyncPeers {
     peers: HashMap<PeerId, PeerInfo>,
     ctx: Option<Arc<SyncNetworkContext>>,
+    file_location_cache: Option<(TxID, Arc<FileLocationCache>)>,
 }
 
 impl SyncPeers {
-    pub fn new(ctx: Arc<SyncNetworkContext>) -> Self {
+    pub fn new(
+        ctx: Arc<SyncNetworkContext>,
+        tx_id: TxID,
+        file_location_cache: Arc<FileLocationCache>,
+    ) -> Self {
         Self {
             peers: Default::default(),
             ctx: Some(ctx),
+            file_location_cache: Some((tx_id, file_location_cache)),
         }
     }
 
@@ -192,6 +200,8 @@ impl SyncPeers {
                     if info.since.elapsed() >= PEER_CONNECT_TIMEOUT {
                         info!(%peer_id, %info.addr, "Peer connection timeout");
                         bad_peers.push(*peer_id);
+
+                        // Ban peer in case of continuous connection timeout
                         if let Some(ctx) = &self.ctx {
                             ctx.report_peer(
                                 *peer_id,
@@ -199,6 +209,11 @@ impl SyncPeers {
                                 "Dail timeout",
                             );
                         }
+
+                        // Remove cached file announcement if connection timeout
+                        if let Some((tx_id, cache)) = &self.file_location_cache {
+                            cache.remove(tx_id, peer_id);
+                        }
                     }
                 }
 
diff --git a/node/sync/src/controllers/serial.rs b/node/sync/src/controllers/serial.rs
index b80c333..474acbc 100644
--- a/node/sync/src/controllers/serial.rs
+++ b/node/sync/src/controllers/serial.rs
@@ -110,7 +110,7 @@ impl SerialSyncController {
             next_chunk: goal.index_start,
             failures: 0,
             state: SyncState::Idle,
-            peers: SyncPeers::new(ctx.clone()),
+            peers: SyncPeers::new(ctx.clone(), tx_id, file_location_cache.clone()),
             ctx,
             store,
             file_location_cache,