From e945e9f30835f5907becbfbf4095a0ad4d0553a3 Mon Sep 17 00:00:00 2001
From: phiresky <phireskyde+git@gmail.com>
Date: Fri, 1 Sep 2023 12:30:35 +0000
Subject: [PATCH] dead instances in one query

---
 crates/db_schema/src/impls/instance.rs | 24 +++++++++++-------------
 crates/federate/src/lib.rs             |  4 +---
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/crates/db_schema/src/impls/instance.rs b/crates/db_schema/src/impls/instance.rs
index 9ac1b2312..7e162717c 100644
--- a/crates/db_schema/src/impls/instance.rs
+++ b/crates/db_schema/src/impls/instance.rs
@@ -64,15 +64,6 @@ impl Instance {
       .await
   }
 
-  pub async fn dead_instances(pool: &mut DbPool<'_>) -> Result<Vec<String>, Error> {
-    let conn = &mut get_conn(pool).await?;
-    instance::table
-      .select(instance::domain)
-      .filter(coalesce(instance::updated, instance::published).lt(now() - 3.days()))
-      .get_results(conn)
-      .await
-  }
-
   #[cfg(test)]
   pub async fn delete_all(pool: &mut DbPool<'_>) -> Result<usize, Error> {
     let conn = &mut get_conn(pool).await?;
@@ -96,10 +87,15 @@ impl Instance {
       .await
   }
 
-  /// returns a list of all instances, each with a flag of whether the instance is allowed or not
+  /// returns a list of all instances, each with a flag of whether the instance is allowed or not and dead or not
   /// ordered by id
-  pub async fn read_all_with_blocked(pool: &mut DbPool<'_>) -> Result<Vec<(Self, bool)>, Error> {
+  pub async fn read_all_with_blocked_and_dead(
+    pool: &mut DbPool<'_>,
+  ) -> Result<Vec<(Self, bool, bool)>, Error> {
     let conn = &mut get_conn(pool).await?;
+    let is_dead_expr = coalesce(instance::updated, instance::published).lt(now() - 3.days());
+    // this needs to be done in two steps because the meaning of the "blocked" column depends on the existence
+    // of any value at all in the allowlist. (so a normal join wouldn't work)
     let use_allowlist = federation_allowlist::table
       .select(count_star().gt(0))
       .get_result::<bool>(conn)
@@ -110,9 +106,10 @@ impl Instance {
         .select((
           Self::as_select(),
           federation_allowlist::id.nullable().is_not_null(),
+          is_dead_expr,
         ))
         .order_by(instance::id)
-        .get_results::<(Self, bool)>(conn)
+        .get_results::<(Self, bool, bool)>(conn)
         .await
     } else {
       instance::table
@@ -120,9 +117,10 @@ impl Instance {
         .select((
           Self::as_select(),
           federation_blocklist::id.nullable().is_null(),
+          is_dead_expr,
         ))
         .order_by(instance::id)
-        .get_results::<(Self, bool)>(conn)
+        .get_results::<(Self, bool, bool)>(conn)
         .await
     }
   }
diff --git a/crates/federate/src/lib.rs b/crates/federate/src/lib.rs
index 5d5b2617f..7e948c83a 100644
--- a/crates/federate/src/lib.rs
+++ b/crates/federate/src/lib.rs
@@ -46,11 +46,10 @@ async fn start_stop_federation_workers(
   let pool2 = &mut DbPool::Pool(&pool);
   let process_index = opts.process_index - 1;
   loop {
-    let dead: HashSet<String> = HashSet::from_iter(Instance::dead_instances(pool2).await?);
     let mut total_count = 0;
     let mut dead_count = 0;
     let mut disallowed_count = 0;
-    for (instance, allowed) in Instance::read_all_with_blocked(pool2).await? {
+    for (instance, allowed, is_dead) in Instance::read_all_with_blocked_and_dead(pool2).await? {
       if instance.id.inner() % opts.process_count != process_index {
         continue;
       }
@@ -58,7 +57,6 @@ async fn start_stop_federation_workers(
       if !allowed {
         disallowed_count += 1;
       }
-      let is_dead = dead.contains(&instance.domain);
       if is_dead {
         dead_count += 1;
       }