Batch hot rank updates

Cargo.toml

@@ -140,3 +140,4 @@ opentelemetry-otlp = { version = "0.10.0", optional = true }
 pict-rs = { version = "0.4.0-rc.3", optional = true }
 tokio.workspace = true
 actix-cors = "0.6.4"
+chrono = { workspace = true }

migrations/2023-06-24-185942_aggegates_published_indexes/down.sql

@@ -0,0 +1,2 @@
+drop index idx_comment_aggregates_published;
+drop index idx_community_aggregates_published;

migrations/2023-06-24-185942_aggegates_published_indexes/up.sql

@@ -0,0 +1,4 @@
+-- Add indexes on published column (needed for hot_rank updates)
+
+create index idx_community_aggregates_published on community_aggregates (published desc);
+create index idx_comment_aggregates_published on comment_aggregates (published desc);

src/scheduled_tasks.rs

@@ -1,25 +1,21 @@
+use chrono::NaiveDateTime;
 use clokwerk::{Scheduler, TimeUnits as CTimeUnits};
 use diesel::{
   dsl::{now, IntervalDsl},
+  result::Error,
+  sql_types::{Integer, Timestamp},
   Connection,
   ExpressionMethods,
   QueryDsl,
+  QueryableByName,
 };
 // Import week days and WeekDay
 use diesel::{sql_query, PgConnection, RunQueryDsl};
 use lemmy_api_common::context::LemmyContext;
 use lemmy_db_schema::{
-  schema::{
-    activity,
-    comment_aggregates,
-    community_aggregates,
-    community_person_ban,
-    instance,
-    person,
-    post_aggregates,
-  },
+  schema::{activity, community_person_ban, instance, person},
   source::instance::{Instance, InstanceForm},
-  utils::{functions::hot_rank, naive_now},
+  utils::naive_now,
 };
 use lemmy_routes::nodeinfo::NodeInfo;
 use lemmy_utils::{error::LemmyError, REQWEST_TIMEOUT};
@@ -46,9 +42,9 @@ pub fn setup(
     update_banned_when_expired(&mut conn);
   });
 
-  // Update hot ranks every 5 minutes
+  // Update hot ranks every 15 minutes
   let url = db_url.clone();
-  scheduler.every(CTimeUnits::minutes(5)).run(move || {
+  scheduler.every(CTimeUnits::minutes(15)).run(move || {
     let mut conn = PgConnection::establish(&url).expect("could not establish connection");
     update_hot_ranks(&mut conn, true);
   });
@@ -89,64 +85,93 @@ fn startup_jobs(db_url: &str) {
 }
 
 /// Update the hot_rank columns for the aggregates tables
+/// Runs in batches until all necessary rows are updated once
 fn update_hot_ranks(conn: &mut PgConnection, last_week_only: bool) {
-  let mut post_update = diesel::update(post_aggregates::table).into_boxed();
-  let mut comment_update = diesel::update(comment_aggregates::table).into_boxed();
-  let mut community_update = diesel::update(community_aggregates::table).into_boxed();
-
-  // Only update for the last week of content
-  if last_week_only {
+  let process_start_time = if last_week_only {
     info!("Updating hot ranks for last week...");
-    let last_week = now - diesel::dsl::IntervalDsl::weeks(1);
-
-    post_update = post_update.filter(post_aggregates::published.gt(last_week));
-    comment_update = comment_update.filter(comment_aggregates::published.gt(last_week));
-    community_update = community_update.filter(community_aggregates::published.gt(last_week));
+    naive_now() - chrono::Duration::days(7)
   } else {
     info!("Updating hot ranks for all history...");
-  }
+    NaiveDateTime::from_timestamp_opt(0, 0).expect("0 timestamp creation")
+  };
 
-  match post_update
-    .set((
-      post_aggregates::hot_rank.eq(hot_rank(post_aggregates::score, post_aggregates::published)),
-      post_aggregates::hot_rank_active.eq(hot_rank(
-        post_aggregates::score,
-        post_aggregates::newest_comment_time_necro,
-      )),
-    ))
-    .execute(conn)
-  {
-    Ok(_) => {}
-    Err(e) => {
-      error!("Failed to update post_aggregates hot_ranks: {}", e)
-    }
-  }
+  process_hot_ranks_in_batches(
+    conn,
+    "post_aggregates",
+    "SET hot_rank = hot_rank(a.score, a.published), 
+         hot_rank_active = hot_rank(a.score, a.newest_comment_time_necro)",
+    process_start_time,
+  );
+
+  process_hot_ranks_in_batches(
+    conn,
+    "comment_aggregates",
+    "SET hot_rank = hot_rank(a.score, a.published)",
+    process_start_time,
+  );
+
+  process_hot_ranks_in_batches(
+    conn,
+    "community_aggregates",
+    "SET hot_rank = hot_rank(a.subscribers, a.published)",
+    process_start_time,
+  );
+
+  info!("Finished hot ranks update!");
+}
 
-  match comment_update
-    .set(comment_aggregates::hot_rank.eq(hot_rank(
-      comment_aggregates::score,
-      comment_aggregates::published,
-    )))
-    .execute(conn)
-  {
-    Ok(_) => {}
-    Err(e) => {
-      error!("Failed to update comment_aggregates hot_ranks: {}", e)
-    }
-  }
+#[derive(QueryableByName)]
+struct HotRanksUpdateResult {
+  #[diesel(sql_type = Timestamp)]
+  published: NaiveDateTime,
+}
 
-  match community_update
-    .set(community_aggregates::hot_rank.eq(hot_rank(
-      community_aggregates::subscribers,
-      community_aggregates::published,
-    )))
-    .execute(conn)
-  {
-    Ok(_) => {
-      info!("Done.");
-    }
-    Err(e) => {
-      error!("Failed to update community_aggregates hot_ranks: {}", e)
+/// Runs the hot rank update query in batches until all ids have been processed.
+/// In `set_clause`, "a" will refer to the current aggregates table.
+/// Locked rows are skipped in order to prevent deadlocks (they will likely get updated on the next
+/// run)
+fn process_hot_ranks_in_batches(
+  conn: &mut PgConnection,
+  table_name: &str,
+  set_clause: &str,
+  process_start_time: NaiveDateTime,
+) {
+  let update_batch_size = 1000; // Bigger batches than this tend to cause seq scans
+  let mut previous_batch_last_published = process_start_time;
+  loop {
+    // Raw `sql_query` is used as a performance optimization - Diesel does not support doing this
+    // in a single query (neither as a CTE, nor using a subquery)
+    let result = sql_query(format!(
+      r#"WITH batch AS (SELECT a.id
+               FROM {aggregates_table} a
+               WHERE a.published > $1
+               ORDER BY a.published
+               LIMIT $2
+               FOR UPDATE SKIP LOCKED)
+         UPDATE {aggregates_table} a {set_clause}
+             FROM batch WHERE a.id = batch.id RETURNING a.published;
+    "#,
+      aggregates_table = table_name,
+      set_clause = set_clause
+    ))
+    .bind::<Timestamp, _>(previous_batch_last_published)
+    .bind::<Integer, _>(update_batch_size)
+    .get_results::<HotRanksUpdateResult>(conn);
+
+    match result {
+      Ok(updated_rows) => match updated_rows.last() {
+        Some(row) => {
+          previous_batch_last_published = row.published;
+        }
+        None => {
+          info!("Finished processing {} hot ranks!", table_name);
+          break;
+        }
+      },
+      Err(e) => {
+        error!("Failed to update {} hot_ranks: {}", table_name, e);
+        break;
+      }
     }
   }
 }