From 8c18481d1dab411bd9270ed56e2b6b5c1f465d3f Mon Sep 17 00:00:00 2001
From: strawberry <strawberry@puppygock.gay>
Date: Fri, 10 Jan 2025 23:51:08 -0500
Subject: [PATCH] optimise resetting all user presences

Signed-off-by: strawberry <strawberry@puppygock.gay>
---
 src/admin/query/presence.rs  |  2 +-
 src/service/presence/data.rs |  4 +--
 src/service/presence/mod.rs  | 53 +++++++++++++++++++++++-------------
 src/service/services.rs      |  9 ++----
 4 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/src/admin/query/presence.rs b/src/admin/query/presence.rs
index 45bb6bd9..0de6b696 100644
--- a/src/admin/query/presence.rs
+++ b/src/admin/query/presence.rs
@@ -32,7 +32,7 @@ pub(super) async fn process(
 	match subcommand {
 		| PresenceCommand::GetPresence { user_id } => {
 			let timer = tokio::time::Instant::now();
-			let results = services.presence.db.get_presence(&user_id).await;
+			let results = services.presence.get_presence(&user_id).await;
 			let query_time = timer.elapsed();
 
 			Ok(RoomMessageEventContent::notice_markdown(format!(
diff --git a/src/service/presence/data.rs b/src/service/presence/data.rs
index 3d614333..4ec0a7ee 100644
--- a/src/service/presence/data.rs
+++ b/src/service/presence/data.rs
@@ -12,7 +12,7 @@ use ruma::{events::presence::PresenceEvent, presence::PresenceState, UInt, UserI
 use super::Presence;
 use crate::{globals, users, Dep};
 
-pub struct Data {
+pub(crate) struct Data {
 	presenceid_presence: Arc<Map>,
 	userid_presenceid: Arc<Map>,
 	services: Services,
@@ -36,7 +36,7 @@ impl Data {
 		}
 	}
 
-	pub async fn get_presence(&self, user_id: &UserId) -> Result<(u64, PresenceEvent)> {
+	pub(super) async fn get_presence(&self, user_id: &UserId) -> Result<(u64, PresenceEvent)> {
 		let count = self
 			.userid_presenceid
 			.get(user_id)
diff --git a/src/service/presence/mod.rs b/src/service/presence/mod.rs
index ab7c76a1..eb4105e5 100644
--- a/src/service/presence/mod.rs
+++ b/src/service/presence/mod.rs
@@ -4,7 +4,10 @@ mod presence;
 use std::{sync::Arc, time::Duration};
 
 use async_trait::async_trait;
-use conduwuit::{checked, debug, error, result::LogErr, Error, Result, Server};
+use conduwuit::{
+	checked, debug, debug_warn, error, result::LogErr, trace, Error, Result, Server,
+};
+use database::Database;
 use futures::{stream::FuturesUnordered, Stream, StreamExt, TryFutureExt};
 use loole::{Receiver, Sender};
 use ruma::{events::presence::PresenceEvent, presence::PresenceState, OwnedUserId, UInt, UserId};
@@ -18,12 +21,13 @@ pub struct Service {
 	timeout_remote_users: bool,
 	idle_timeout: u64,
 	offline_timeout: u64,
-	pub db: Data,
+	db: Data,
 	services: Services,
 }
 
 struct Services {
 	server: Arc<Server>,
+	db: Arc<Database>,
 	globals: Dep<globals::Service>,
 	users: Dep<users::Service>,
 }
@@ -44,6 +48,7 @@ impl crate::Service for Service {
 			db: Data::new(&args),
 			services: Services {
 				server: args.server.clone(),
+				db: args.db.clone(),
 				globals: args.depend::<globals::Service>("globals"),
 				users: args.depend::<users::Service>("users"),
 			},
@@ -171,7 +176,9 @@ impl Service {
 	}
 
 	// Unset online/unavailable presence to offline on startup
-	pub async fn unset_all_presence(&self) -> Result<()> {
+	pub async fn unset_all_presence(&self) {
+		let _cork = self.services.db.cork();
+
 		for user_id in &self
 			.services
 			.users
@@ -184,28 +191,36 @@ impl Service {
 
 			let presence = match presence {
 				| Ok((_, ref presence)) => &presence.content,
-				| _ => return Ok(()),
+				| _ => continue,
 			};
 
-			let need_reset = match presence.presence {
-				| PresenceState::Unavailable | PresenceState::Online => true,
-				| _ => false,
-			};
-
-			if !need_reset {
+			if !matches!(
+				presence.presence,
+				PresenceState::Unavailable | PresenceState::Online | PresenceState::Busy
+			) {
+				trace!(?user_id, ?presence, "Skipping user");
 				continue;
 			}
 
-			self.set_presence(
-				user_id,
-				&PresenceState::Offline,
-				Some(false),
-				presence.last_active_ago,
-				presence.status_msg.clone(),
-			)
-			.await?;
+			trace!(?user_id, ?presence, "Resetting presence to offline");
+
+			_ = self
+				.set_presence(
+					user_id,
+					&PresenceState::Offline,
+					Some(false),
+					presence.last_active_ago,
+					presence.status_msg.clone(),
+				)
+				.await
+				.inspect_err(|e| {
+					debug_warn!(
+						?presence,
+						"{user_id} has invalid presence in database and failed to reset it to \
+						 offline: {e}"
+					);
+				});
 		}
-		Ok(())
 	}
 
 	/// Returns the most recent presence updates that happened after the event
diff --git a/src/service/services.rs b/src/service/services.rs
index 9e099759..1aa87f58 100644
--- a/src/service/services.rs
+++ b/src/service/services.rs
@@ -123,13 +123,10 @@ impl Services {
 			.start()
 			.await?;
 
-		// clear online statuses
-		if self.server.config.allow_local_presence {
-			_ = self.presence.unset_all_presence().await;
-		}
-
-		// set the server user as online
+		// reset dormant online/away statuses to offline, and set the server user as
+		// online
 		if self.server.config.allow_local_presence && !self.db.is_read_only() {
+			self.presence.unset_all_presence().await;
 			_ = self
 				.presence
 				.ping_presence(&self.globals.server_user, &ruma::presence::PresenceState::Online)