chore: more robust followee_head_checker

dpc · dpc · commit ea01b3aab88a · 2025-03-08T21:17:30.000-08:00
diff --git a/.aider.conf.yml b/.aider.conf.yml
@@ -1,4 +1,2 @@
-lint-cmd: cargo check -q
 test-cmd: cargo nextest run --cargo-quiet
 cache-prompts: true
-stream: false
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,7 +16,7 @@ resolver = "2"
 
 [workspace.package]
 version = "0.1.0"
-edition = "2021"
+edition = "2024"
 license = "MIT OR APACHE-2.0 OR MPL-2.0"
 description = "A p2p (f2f) social network."
 
diff --git a/crates/rostra-client/src/task.rs b/crates/rostra-client/src/task.rs
@@ -1,3 +1,4 @@
+pub(crate) mod connection_cache;
 pub(crate) mod followee_head_checker;
 pub(crate) mod head_merger;
 pub(crate) mod head_update_broadcaster;
diff --git a/crates/rostra-client/src/task/connection_cache.rs b/crates/rostra-client/src/task/connection_cache.rs
@@ -0,0 +1,53 @@
+use std::collections::BTreeMap;
+use std::collections::btree_map::Entry;
+
+use rostra_core::id::RostraId;
+use rostra_p2p::Connection;
+
+use crate::ClientRef;
+
+#[derive(Debug)]
+pub enum ConnectionState {
+    Connected(Connection),
+    Failed,
+}
+
+#[derive(Default)]
+pub struct ConnectionCache {
+    connections: BTreeMap<RostraId, ConnectionState>,
+}
+
+impl ConnectionCache {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub async fn get_or_connect(
+        &mut self,
+        client: &ClientRef<'_>,
+        id: RostraId,
+    ) -> Option<&mut Connection> {
+        match self.connections.entry(id) {
+            Entry::Occupied(entry) => match entry.get() {
+                ConnectionState::Connected(_) => {}
+                ConnectionState::Failed => return None,
+            },
+            Entry::Vacant(entry) => match client.connect(id).await {
+                Ok(conn) => {
+                    entry.insert(ConnectionState::Connected(conn));
+                }
+                Err(_) => {
+                    entry.insert(ConnectionState::Failed);
+                    return None;
+                }
+            },
+        }
+
+        let ConnectionState::Connected(conn) =
+            self.connections.get_mut(&id).expect("Just inserted")
+        else {
+            unreachable!()
+        };
+        Some(conn)
+    }
+}
diff --git a/crates/rostra-client/src/task/followee_head_checker.rs b/crates/rostra-client/src/task/followee_head_checker.rs
@@ -1,11 +1,12 @@
-use std::collections::{BinaryHeap, HashMap};
+use std::collections::{BTreeMap, BinaryHeap, HashMap};
 use std::sync::Arc;
 use std::time::Duration;
 
 use rostra_client_db::{Database, IdsFolloweesRecord, InsertEventOutcome};
+use rostra_core::ShortEventId;
 use rostra_core::event::PersonaId;
 use rostra_core::id::RostraId;
-use rostra_core::ShortEventId;
+use rostra_p2p::Connection;
 use rostra_p2p::connection::GetHeadRequest;
 use rostra_util::is_rostra_dev_mode_set;
 use rostra_util_error::{BoxedErrorResult, FmtCompact, WhateverResult};
@@ -14,8 +15,9 @@ use snafu::ResultExt as _;
 use tokio::sync::watch;
 use tracing::{debug, info, instrument, trace};
 
-use crate::client::Client;
+use super::connection_cache::ConnectionCache;
 use crate::ClientRef;
+use crate::client::Client;
 const LOG_TARGET: &str = "rostra::head_checker";
 
 pub struct FolloweeHeadChecker {
@@ -33,7 +35,6 @@ impl FolloweeHeadChecker {
             client: client.handle(),
             db: client.db().to_owned(),
             self_id: client.rostra_id(),
-
             followee_updated: client.self_followees_subscribe(),
             check_for_updates_rx: client.check_for_updates_tx_subscribe(),
         }
@@ -70,6 +71,9 @@ impl FolloweeHeadChecker {
                 break;
             };
 
+            let mut connections = ConnectionCache::new();
+            let mut followers_by_followee = BTreeMap::new();
+
             let self_followees = storage.get_self_followees().await;
 
             for (followee, _persona_id) in [&(self.self_id, PersonaId::default())]
@@ -98,9 +102,16 @@ impl FolloweeHeadChecker {
                         }
                         Ok(Some(head)) => {
                             info!(target: LOG_TARGET, id = %followee, %source, "Has updates");
-                            if let Err(err) = self.download_new_data(&client, *followee, head).await
+                            if let Err(err) = self
+                                .download_new_data(
+                                    *followee,
+                                    head,
+                                    &mut connections,
+                                    &mut followers_by_followee,
+                                )
+                                .await
                             {
-                                info!(target: LOG_TARGET, err = %err.fmt_compact(), id = %followee, "Failed to download new data");
+                                info!(target: LOG_TARGET, err = %(&*err).fmt_compact(), id = %followee, "Failed to download new data");
                             }
                         }
                     }
@@ -147,20 +158,73 @@ impl FolloweeHeadChecker {
     }
 
     async fn download_new_data(
+        &self,
+        rostra_id: RostraId,
+        head: ShortEventId,
+        connections: &mut ConnectionCache,
+        followers_by_followee: &mut BTreeMap<RostraId, Vec<RostraId>>,
+    ) -> BoxedErrorResult<()> {
+        let followers = if let Some(followers) = followers_by_followee.get(&rostra_id) {
+            followers
+        } else {
+            let client = self.client.client_ref().boxed()?;
+            let storage = client.db();
+            let followers = storage.get_followers(rostra_id).await;
+            followers_by_followee.insert(rostra_id, followers);
+
+            followers_by_followee
+                .get(&rostra_id)
+                .expect("Just inserted")
+        };
+
+        for follower_id in followers.iter().chain([rostra_id, self.self_id].iter()) {
+            let Ok(client) = self.client.client_ref().boxed() else {
+                break;
+            };
+            let Some(conn) = connections.get_or_connect(&client, *follower_id).await else {
+                continue;
+            };
+
+            debug!(target: LOG_TARGET,
+                rostra_id = %rostra_id,
+                head = %head,
+                follower_id = %follower_id,
+                "Getting event data from a peer"
+            );
+
+            match self
+                .download_new_data_from(&client, rostra_id, conn, head)
+                .await
+            {
+                Ok(true) => {
+                    return Ok(());
+                }
+                Ok(false) => {}
+                Err(err) => {
+                    debug!(target: LOG_TARGET,
+                        rostra_id = %rostra_id,
+                        head = %head,
+                        follower_id = %follower_id,
+                        err = %err.fmt_compact(),
+                        "Error getting event from a peer"
+                    );
+                }
+            }
+        }
+        Ok(())
+    }
+    async fn download_new_data_from(
         &self,
         client: &ClientRef<'_>,
         rostra_id: RostraId,
+        conn: &mut Connection,
         head: ShortEventId,
-    ) -> WhateverResult<()> {
+    ) -> WhateverResult<bool> {
         let mut events = BinaryHeap::from([(0, head)]);
+        let mut downloaded_anything = false;
 
         let storage = client.db();
 
-        let conn = client
-            .connect(rostra_id)
-            .await
-            .whatever_context("Failed to connect")?;
-
         let peer_id = conn.remote_node_id();
 
         while let Some((depth, event_id)) = events.pop() {
@@ -187,6 +251,7 @@ impl FolloweeHeadChecker {
                 );
                 continue;
             };
+            downloaded_anything = true;
             let (insert_outcome, process_state) = storage.process_event(&event).await;
 
             if let InsertEventOutcome::Inserted {
@@ -226,6 +291,6 @@ impl FolloweeHeadChecker {
             }
         }
 
-        Ok(())
+        Ok(downloaded_anything)
     }
 }
diff --git a/crates/rostra-client/src/task/missing_event_content_fetcher.rs b/crates/rostra-client/src/task/missing_event_content_fetcher.rs
@@ -1,17 +1,17 @@
 use std::collections::BTreeMap;
 use std::time::Duration;
 
+use rostra_core::ShortEventId;
 use rostra_core::event::VerifiedEvent;
 use rostra_core::id::RostraId;
-use rostra_core::ShortEventId;
-use rostra_p2p::Connection;
 use rostra_util::is_rostra_dev_mode_set;
 use rostra_util_error::{BoxedErrorResult, FmtCompact, WhateverResult};
 use snafu::ResultExt as _;
 use tracing::{debug, instrument, trace};
 
+use super::connection_cache::ConnectionCache;
+use crate::LOG_TARGET;
 use crate::client::Client;
-use crate::{ClientHandle, LOG_TARGET};
 
 #[derive(Clone)]
 pub struct MissingEventContentFetcher {
@@ -50,16 +50,15 @@ impl MissingEventContentFetcher {
 
             let mut cursor: Option<ShortEventId> = None;
 
-            let mut connections = BTreeMap::new();
+            let mut connections = ConnectionCache::new();
             let mut followers_by_followee = BTreeMap::new();
 
             loop {
                 let (events, new_cursor) = db.paginate_missing_events_contents(cursor, 100).await;
 
                 for (author_id, event_id) in events {
                     let _ = self
-                        .get_event_from_followers(
-                            &self.client,
+                        .get_event_content_from_followers(
                             author_id,
                             event_id,
                             &mut connections,
@@ -78,12 +77,11 @@ impl MissingEventContentFetcher {
         }
     }
 
-    async fn get_event_from_followers(
+    async fn get_event_content_from_followers(
         &self,
-        client: &ClientHandle,
         author_id: RostraId,
         event_id: ShortEventId,
-        connections: &mut BTreeMap<RostraId, Connection>,
+        connections: &mut ConnectionCache,
         followers_by_followee: &mut BTreeMap<RostraId, Vec<RostraId>>,
         db: &rostra_client_db::Database,
     ) -> BoxedErrorResult<()> {
@@ -99,16 +97,12 @@ impl MissingEventContentFetcher {
         };
 
         for follower_id in followers.iter().chain([author_id, self.self_id].iter()) {
-            if connections.get(follower_id).is_none() {
-                let conn = client
-                    .client_ref()
-                    .boxed()?
-                    .connect(*follower_id)
-                    .await
-                    .boxed()?;
-                connections.insert(*follower_id, conn);
-            }
-            let conn = connections.get_mut(follower_id).expect("Must exist");
+            let Ok(client) = self.client.client_ref().boxed() else {
+                break;
+            };
+            let Some(conn) = connections.get_or_connect(&client, *follower_id).await else {
+                continue;
+            };
 
             debug!(target:  LOG_TARGET,
                 author_id = %author_id,
diff --git a/crates/rostra-client/src/task/missing_event_fetcher.rs b/crates/rostra-client/src/task/missing_event_fetcher.rs
@@ -1,15 +1,14 @@
-use std::collections::BTreeMap;
-
+use rostra_core::ShortEventId;
 use rostra_core::event::{EventExt as _, SignedEventExt as _, VerifiedEvent};
 use rostra_core::id::RostraId;
-use rostra_core::ShortEventId;
 use rostra_p2p::Connection;
 use rostra_util_error::{BoxedErrorResult, FmtCompact, WhateverResult};
 use snafu::ResultExt as _;
 use tracing::{debug, instrument, trace, warn};
 
+use super::connection_cache::ConnectionCache;
+use crate::LOG_TARGET;
 use crate::client::Client;
-use crate::{ClientHandle, LOG_TARGET};
 
 #[derive(Clone)]
 pub struct MissingEventFetcher {
@@ -69,22 +68,22 @@ impl MissingEventFetcher {
                 continue;
             }
 
-            let mut connections = BTreeMap::new();
+            let mut connections = ConnectionCache::new();
 
             for follower_id in followers.iter().chain([self.self_id].iter()) {
+                let Ok(client) = self.client.client_ref().boxed() else {
+                    break;
+                };
+                let Some(conn) = connections.get_or_connect(&client, *follower_id).await else {
+                    continue;
+                };
+
                 for missing_event in &missing_events {
                     if db.has_event(*missing_event).await {
                         continue;
                     }
                     match self
-                        .get_event_from(
-                            &self.client,
-                            author_id,
-                            *missing_event,
-                            *follower_id,
-                            &mut connections,
-                            &db,
-                        )
+                        .get_event_from(author_id, *missing_event, *follower_id, conn, &db)
                         .await
                     {
                         Ok(_) => {}
@@ -97,7 +96,6 @@ impl MissingEventFetcher {
                                 err = %(&*err).fmt_compact(),
                                 "Error while getting id from a peer"
                             );
-                            connections.remove(follower_id);
                         }
                     }
                 }
@@ -107,24 +105,12 @@ impl MissingEventFetcher {
 
     async fn get_event_from(
         &self,
-        client: &ClientHandle,
         author_id: RostraId,
         event_id: ShortEventId,
         follower_id: RostraId,
-        connections: &mut BTreeMap<RostraId, Connection>,
+        conn: &mut Connection,
         db: &rostra_client_db::Database,
     ) -> BoxedErrorResult<()> {
-        if connections.get(&follower_id).is_none() {
-            let conn = client
-                .client_ref()
-                .boxed()?
-                .connect(follower_id)
-                .await
-                .boxed()?;
-            connections.insert(follower_id, conn);
-        }
-        let conn = connections.get_mut(&follower_id).expect("Must exist");
-
         debug!(target:  LOG_TARGET,
             author_id = %author_id,
             event_id = %event_id,
diff --git a/crates/rostra-p2p/src/connection.rs b/crates/rostra-p2p/src/connection.rs

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+pub(crate) mod connection_cache;`
`1`	`2`	`pub(crate) mod followee_head_checker;`
`2`	`3`	`pub(crate) mod head_merger;`
`3`	`4`	`pub(crate) mod head_update_broadcaster;`