close unneeded relays

Author: ksedgwic

crates/notedeck/src/subman.rs

@@ -1,5 +1,5 @@
 use futures::{channel::mpsc, FutureExt, StreamExt};
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::fmt;
 use std::{cell::RefCell, cmp::Ordering, rc::Rc};
 use thiserror::Error;
@@ -24,6 +24,7 @@ use crate::RelaySpec;
 /// async fn main() -> Result<(), Box<dyn Error>> {
 ///     let mut ndb = Ndb::new("the/db/path/", &Config::new())?;
 ///     let mut subman = SubMan::new(ndb.clone(), RelayPool::new());
+///     let default_relays = vec![];
 ///
 ///     // Define a filter and build the subscription specification
 ///     let filter = Filter::new().kinds(vec![1, 2, 3]).build();
@@ -33,7 +34,7 @@ use crate::RelaySpec;
 ///         .build();
 ///
 ///     // Subscribe and obtain a SubReceiver
-///     let mut receiver = subman.subscribe(spec)?;
+///     let mut receiver = subman.subscribe(spec, &default_relays)?;
 ///
 ///     // Process incoming note keys
 ///     loop {
@@ -514,6 +515,7 @@ impl SubMan {
     pub fn process_relays<H: LegacyRelayHandler>(
         &mut self,
         legacy_relay_handler: &mut H,
+        default_relays: &[RelaySpec],
     ) -> SubResult<()> {
         let wakeup = move || {
             // ignore
@@ -580,6 +582,9 @@ impl SubMan {
                 }
             }
         }
+
+        self.close_unneeded_relays(default_relays);
+
         Ok(())
     }
 
@@ -646,6 +651,41 @@ impl SubMan {
             }
         }
     }
+
+    fn close_unneeded_relays(&mut self, default_relays: &[RelaySpec]) {
+        let current_relays: BTreeSet<String> = self.pool.urls();
+        let needed_relays: BTreeSet<String> = self.needed_relays(default_relays);
+        let unneeded_relays: BTreeSet<_> =
+            current_relays.difference(&needed_relays).cloned().collect();
+        if !unneeded_relays.is_empty() {
+            debug!("closing unneeded relays: {:?}", unneeded_relays);
+            self.pool.remove_urls(&unneeded_relays);
+        }
+    }
+
+    fn needed_relays(&self, default_relays: &[RelaySpec]) -> BTreeSet<String> {
+        let mut needed: BTreeSet<String> = default_relays.iter().map(|rs| rs.url.clone()).collect();
+        // for every remote subscription
+        for ssr in self.rmt.values() {
+            // that has remote substate (all will)
+            if let Some(ref rmtsubstate) = ssr.borrow().rmt {
+                // for each subscription remote relay
+                for (relay, state) in &rmtsubstate.relays {
+                    // include any that are in-play
+                    match state {
+                        RelaySubState::Error(_) | RelaySubState::Closed => {
+                            // these are terminal and we don't need this relay
+                        }
+                        _ => {
+                            // relays in all other states are needed
+                            _ = needed.insert(relay.clone());
+                        }
+                    }
+                }
+            }
+        }
+        needed
+    }
 }
 
 pub trait LegacyRelayHandler {
@@ -755,13 +795,15 @@ mod tests {
         // setup an ndb and subman to test
         let (_mndb, ndb) = ManagedNdb::setup(&testdbs_path_async!());
         let mut subman = SubMan::new(ndb.clone(), RelayPool::new());
+        let default_relays = vec![];
 
         // subscribe to some stuff
         let mut receiver = subman.subscribe(
             SubSpecBuilder::new()
                 .filters(vec![Filter::new().kinds(vec![1]).build()])
                 .constraint(SubConstraint::OnlyLocal)
                 .build(),
+            &default_relays,
         )?;
         let lclid = receiver.lclid().unwrap();
 
@@ -796,13 +838,15 @@ mod tests {
         // setup an ndb and subman to test
         let (_mndb, ndb) = ManagedNdb::setup(&testdbs_path_async!());
         let mut subman = SubMan::new(ndb.clone(), RelayPool::new());
+        let default_relays = vec![];
 
         // subscribe to some stuff
         let mut receiver = subman.subscribe(
             SubSpecBuilder::new()
                 .filters(vec![Filter::new().kinds(vec![1]).build()])
                 .constraint(SubConstraint::OnlyLocal)
                 .build(),
+            &default_relays,
         )?;
         let lclid = receiver.lclid().unwrap();
 
@@ -842,13 +886,15 @@ mod tests {
         // setup an ndb and subman to test
         let (_mndb, ndb) = ManagedNdb::setup(&testdbs_path_async!());
         let mut subman = SubMan::new(ndb.clone(), RelayPool::new());
+        let default_relays = vec![];
 
         // subscribe to some stuff
         let mut receiver = subman.subscribe(
             SubSpecBuilder::new()
                 .filters(vec![Filter::new().kinds(vec![1]).build()])
                 .constraint(SubConstraint::OnlyLocal)
                 .build(),
+            &default_relays,
         )?;
         let lclid = receiver.lclid().unwrap();
 

crates/notedeck_columns/src/app.rs

@@ -155,6 +155,7 @@ fn try_process_event<'a>(
     ctx.input(|i| handle_key_events(i, current_columns));
 
     {
+        let default_relays = app_ctx.accounts.get_all_selected_account_relays();
         let mut relay_handler = RelayHandler::new(
             app_ctx.unknown_ids,
             app_ctx.note_cache,
@@ -163,7 +164,10 @@ fn try_process_event<'a>(
             &mut damus.timeline_cache,
             damus.since_optimize,
         );
-        app_ctx.subman.process_relays(&mut relay_handler).ok();
+        app_ctx
+            .subman
+            .process_relays(&mut relay_handler, &default_relays)
+            .ok();
     }
 
     for (_kind, timeline) in damus.timeline_cache.timelines.iter_mut() {