improve notification clearing

Author: ThatsNoMoon

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "highlights"
-version = "2.1.5"
+version = "2.1.6"
 authors = ["ThatsNoMoon <[email protected]>"]
 repository = "https://github.com/ThatsNoMoon/highlights"
 license = "OSL-3.0"

Cargo.toml

@@ -4,7 +4,11 @@
 //! Functions for sending, editing, and deleting notifications.
 
 use std::{
-	cmp::min, collections::HashMap, fmt::Write as _, ops::Range, time::Duration,
+	cmp::min,
+	collections::HashMap,
+	fmt::Write as _,
+	ops::Range,
+	time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
 use anyhow::{anyhow, bail, Context as _, Error, Result};
@@ -35,7 +39,7 @@ use tracing::{debug, error, info_span};
 use crate::{
 	bot::util::{followup_eph, user_can_read_channel},
 	db::{Ignore, Keyword, Notification, UserState, UserStateKind},
-	global::{EMBED_COLOR, ERROR_COLOR, NOTIFICATION_RETRIES},
+	global::{DISCORD_EPOCH, EMBED_COLOR, ERROR_COLOR, NOTIFICATION_RETRIES},
 	settings::settings,
 };
 
@@ -76,6 +80,16 @@ pub(crate) async fn should_notify_keyword(
 	keyword: &Keyword,
 	ignores: &[Ignore],
 ) -> Result<bool> {
+	if let Some(lifetime) = settings().behavior.notification_lifetime {
+		let creation = (message.id.0 >> 22) + DISCORD_EPOCH;
+		let now =
+			SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64;
+		let age = Duration::from_millis(now.saturating_sub(creation));
+		if age > lifetime {
+			return Ok(false);
+		}
+	}
+
 	if message
 		.mentions
 		.iter()
@@ -743,37 +757,58 @@ async fn clear_old_notifications(
 	lifetime: Duration,
 ) -> Result<()> {
 	debug!("Clearing old notifications");
-	Notification::old_notifications(lifetime)
-		.await?
-		.into_iter()
-		.map(|notification| {
-			clear_sent_notification(
-				ctx,
-				notification.user_id,
-				notification.notification_message,
-				"*Notification expired*",
-			)
-			.or_else(|e| async move {
-				match e.downcast_ref::<SerenityError>() {
-					Some(SerenityError::Http(inner)) => match &**inner {
-						HttpError::UnsuccessfulRequest(ErrorResponse {
-							status_code: StatusCode::NOT_FOUND,
-							..
-						}) => Ok(()),
+	let cutoff_time = SystemTime::now() - lifetime;
 
+	loop {
+		let notifications =
+			Notification::notifications_before(5, cutoff_time).await?;
+
+		if notifications.is_empty() {
+			break Ok(());
+		}
+
+		let sent_ids = notifications
+			.iter()
+			.map(|n| n.notification_message)
+			.collect::<Vec<_>>();
+
+		debug!("Clearing {} notifications", notifications.len());
+
+		let wait_cycle = sleep(Duration::from_secs(2));
+
+		notifications
+			.iter()
+			.map(|notification| {
+				clear_sent_notification(
+					ctx,
+					notification.user_id,
+					notification.notification_message,
+					"*Notification expired*",
+				)
+				.or_else(|e| async move {
+					match e.downcast_ref::<SerenityError>() {
+						Some(SerenityError::Http(inner)) => match &**inner {
+							HttpError::UnsuccessfulRequest(ErrorResponse {
+								status_code: StatusCode::NOT_FOUND,
+								..
+							}) => Ok(()),
+
+							_ => Err(e),
+						},
 						_ => Err(e),
-					},
-					_ => Err(e),
-				}
+					}
+				})
 			})
-		})
-		.collect::<FuturesUnordered<_>>()
-		.try_for_each(|_| async { Ok(()) })
-		.await?;
+			.collect::<FuturesUnordered<_>>()
+			.try_for_each(|_| async { Ok(()) })
+			.await?;
 
-	Notification::delete_old_notifications(lifetime).await?;
+		Notification::delete_notifications(sent_ids).await?;
 
-	Ok(())
+		debug!("Waiting before clearing more notifications");
+
+		wait_cycle.await;
+	}
 }
 
 #[cfg(test)]

src/bot/highlighting.rs

@@ -3,7 +3,7 @@
 
 //! Handling for sent notification messages.
 
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
+use std::time::{SystemTime, UNIX_EPOCH};
 
 use anyhow::Result;
 use futures_util::TryStreamExt;
@@ -12,13 +12,13 @@ use sea_orm::{
 		DeriveActiveModelBehavior, DeriveEntityModel, DerivePrimaryKey,
 		DeriveRelation, EntityTrait, EnumIter, PrimaryKeyTrait,
 	},
-	ColumnTrait, IntoActiveModel, QueryFilter,
+	ColumnTrait, Condition, IntoActiveModel, QueryFilter, QueryOrder,
+	QuerySelect,
 };
 use serenity::model::id::{MessageId, UserId};
 
 use super::{connection, DbInt, IdDbExt};
-
-const DISCORD_EPOCH: u64 = 1420070400000;
+use crate::global::DISCORD_EPOCH;
 
 #[derive(
 	Clone, Debug, PartialEq, Eq, DeriveEntityModel, DeriveActiveModelBehavior,
@@ -111,11 +111,14 @@ impl Notification {
 
 	/// Gets notifications older than a certain duration from the DB.
 	#[tracing::instrument]
-	pub(crate) async fn old_notifications(
-		age: Duration,
+	pub(crate) async fn notifications_before(
+		count: u64,
+		time: SystemTime,
 	) -> Result<Vec<Notification>> {
 		Entity::find()
-			.filter(Column::OriginalMessage.lte(age_to_oldest_snowflake(age)?))
+			.filter(Column::OriginalMessage.lte(time_to_max_snowflake(time)?))
+			.order_by_asc(Column::OriginalMessage)
+			.limit(count)
 			.stream(connection())
 			.await?
 			.map_err(Into::into)
@@ -124,26 +127,36 @@ impl Notification {
 			.await
 	}
 
-	/// Deletes notifications older than a certain duration from the DB.
-	#[tracing::instrument]
-	pub(crate) async fn delete_old_notifications(age: Duration) -> Result<()> {
+	/// Deletes a list of notifications from the DB.
+	#[tracing::instrument(skip_all)]
+	pub(crate) async fn delete_notifications(
+		message_ids: impl IntoIterator<Item = MessageId>,
+	) -> Result<()> {
 		Entity::delete_many()
-			.filter(Column::OriginalMessage.lte(age_to_oldest_snowflake(age)?))
+			.filter(message_ids.into_iter().fold(
+				Condition::any(),
+				|cond, id| {
+					cond.add(Column::NotificationMessage.eq(id.into_db()))
+				},
+			))
 			.exec(connection())
 			.await?;
 
 		Ok(())
 	}
 }
 
-fn age_to_oldest_snowflake(age: Duration) -> Result<u64> {
-	let millis = age.as_millis() as u64;
-	let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64;
-	let oldest_unix = now - millis;
-	let oldest_discord = oldest_unix - DISCORD_EPOCH;
+fn time_to_min_snowflake(time: SystemTime) -> Result<u64> {
+	let unix = time.duration_since(UNIX_EPOCH)?.as_millis() as u64;
+	let oldest_discord = unix - DISCORD_EPOCH;
 	Ok(oldest_discord << 22)
 }
 
+fn time_to_max_snowflake(time: SystemTime) -> Result<u64> {
+	let min = time_to_min_snowflake(time)?;
+	Ok(min | (!0 >> 22))
+}
+
 impl From<Model> for Notification {
 	fn from(model: Model) -> Self {
 		Self {

src/db/notification.rs

@@ -10,3 +10,5 @@ pub(crate) const NOTIFICATION_RETRIES: u8 = 5;
 pub(crate) const EMBED_COLOR: u32 = 0xefff47;
 /// Color of embeds reporting an error to the user.
 pub(crate) const ERROR_COLOR: u32 = 0xff4747;
+
+pub(crate) const DISCORD_EPOCH: u64 = 1420070400000;

src/global.rs

@@ -94,14 +94,10 @@ mod user_address {
 		where
 			E: de::Error,
 		{
-			let socket_addr = v
-				.to_socket_addrs()
-				.map_err(E::custom)?
-				.into_iter()
-				.next()
-				.ok_or_else(|| {
-					E::custom("provided host did not resolve to an address")
-				})?;
+			let socket_addr =
+				v.to_socket_addrs().map_err(E::custom)?.next().ok_or_else(
+					|| E::custom("provided host did not resolve to an address"),
+				)?;
 
 			Ok(UserAddress { socket_addr })
 		}