feat: Send badge updates from message acks

crates/core/database/src/models/channel_unreads/ops.rs

@@ -7,13 +7,13 @@ mod reference;
 
 #[async_trait]
 pub trait AbstractChannelUnreads: Sync + Send {
-    /// Acknowledge a message.
+    /// Acknowledge a message, and returns updated channel unread.
     async fn acknowledge_message(
         &self,
         channel_id: &str,
         user_id: &str,
         message_id: &str,
-    ) -> Result<()>;
+    ) -> Result<Option<ChannelUnread>>;
 
     /// Acknowledge many channels.
     async fn acknowledge_channels(&self, user_id: &str, channel_ids: &[String]) -> Result<()>;
@@ -28,4 +28,7 @@ pub trait AbstractChannelUnreads: Sync + Send {
 
     /// Fetch all channel unreads for a user.
     async fn fetch_unreads(&self, user_id: &str) -> Result<Vec<ChannelUnread>>;
+
+    /// Fetch unread for a specific user in a channel.
+    async fn fetch_unread(&self, user_id: &str, channel_id: &str) -> Result<Option<ChannelUnread>>;
 }

crates/core/database/src/models/channel_unreads/ops/mongodb.rs

@@ -1,4 +1,6 @@
 use bson::Document;
+use mongodb::options::FindOneAndUpdateOptions;
+use mongodb::options::ReturnDocument;
 use mongodb::options::UpdateOptions;
 use revolt_result::Result;
 use ulid::Ulid;
@@ -12,31 +14,35 @@ static COL: &str = "channel_unreads";
 
 #[async_trait]
 impl AbstractChannelUnreads for MongoDb {
-    /// Acknowledge a message.
+    /// Acknowledge a message, and returns updated channel unread.
     async fn acknowledge_message(
         &self,
         channel_id: &str,
         user_id: &str,
         message_id: &str,
-    ) -> Result<()> {
+    ) -> Result<Option<ChannelUnread>> {
-        self.col::<Document>(COL)
+        self.col::<ChannelUnread>(COL)
-            .update_one(
+            .find_one_and_update(
                 doc! {
                     "_id.channel": channel_id,
                     "_id.user": user_id,
                 },
                 doc! {
-                    "$unset": {
+                    "$pull": {
-                        "mentions": 1_i32
+                        "mentions": {
+                            "$lt": message_id
+                        }
                     },
                     "$set": {
                         "last_id": message_id
                     }
                 },
-                UpdateOptions::builder().upsert(true).build(),
+                FindOneAndUpdateOptions::builder()
+                    .upsert(true)
+                    .return_document(ReturnDocument::After)
+                    .build(),
             )
             .await
-            .map(|_| ())
             .map_err(|_| create_database_error!("update_one", COL))
     }
 
@@ -116,4 +122,18 @@ impl AbstractChannelUnreads for MongoDb {
             }
         )
     }
+
+    /// Fetch unread for a specific user in a channel.
+    async fn fetch_unread(&self, user_id: &str, channel_id: &str) -> Result<Option<ChannelUnread>> {
+        query!(
+            self,
+            find_one,
+            COL,
+            doc! {
+                "_id.user": user_id,
+                "_id.channel": channel_id
+            }
+        )
+    }
+
 }

crates/core/database/src/models/channel_unreads/ops/reference.rs

@@ -13,7 +13,7 @@ impl AbstractChannelUnreads for ReferenceDb {
         channel_id: &str,
         user_id: &str,
         message_id: &str,
-    ) -> Result<()> {
+    ) -> Result<Option<ChannelUnread>> {
         let mut unreads = self.channel_unreads.lock().await;
         let key = ChannelCompositeKey {
             channel: channel_id.to_string(),
@@ -27,14 +27,14 @@ impl AbstractChannelUnreads for ReferenceDb {
             unreads.insert(
                 key.clone(),
                 ChannelUnread {
-                    id: key,
+                    id: key.clone(),
                     last_id: Some(message_id.to_string()),
                     mentions: None,
                 },
             );
         }
 
-        Ok(())
+        Ok(unreads.get(&key).cloned())
     }
 
     /// Acknowledge many channels.
@@ -87,4 +87,14 @@ impl AbstractChannelUnreads for ReferenceDb {
             .cloned()
             .collect())
     }
+
+    /// Fetch unread for a specific user in a channel.
+    async fn fetch_unread(&self, user_id: &str, channel_id: &str) -> Result<Option<ChannelUnread>> {
+        let unreads = self.channel_unreads.lock().await;
+
+        Ok(unreads.get(&ChannelCompositeKey {
+            channel: channel_id.to_string(),
+            user: user_id.to_string()
+        }).cloned())
+    }
 }

crates/core/database/src/tasks/ack.rs

@@ -5,7 +5,9 @@ use deadqueue::limited::Queue;
 use once_cell::sync::Lazy;
 use std::{collections::HashMap, time::Duration};
 
-use super::DelayedTask;
+use revolt_result::Result;
+
+use super::{apple_notifications::{self, ApnJob}, DelayedTask};
 
 /// Enumeration of possible events
 #[derive(Debug, Eq, PartialEq)]
@@ -52,8 +54,43 @@ pub async fn queue(channel: String, user: String, event: AckEvent) {
     info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
 }
 
+pub async fn handle_ack_event(event: &AckEvent, db: &Database, authifier_db: &authifier::Database, user: &str, channel: &str) -> Result<()> {
+    match &event {
+        #[allow(clippy::disallowed_methods)] // event is sent by higher level function
+        AckEvent::AckMessage { id } => {
+            let unread = db.fetch_unread(user, channel).await?;
+            let updated = db.acknowledge_message(channel, user, id).await?;
+
+            if let (Some(before), Some(after)) = (unread, updated) {
+                let before_mentions = before.mentions.unwrap_or_default().len();
+                let after_mentions = after.mentions.unwrap_or_default().len();
+
+                let mentions_acked = before_mentions - after_mentions;
+
+                if mentions_acked > 0 {
+                    if let Ok(sessions) = authifier_db.find_sessions(user).await {
+                        for session in sessions {
+                            if let Some(sub) = session.subscription {
+                                if sub.endpoint == "apn" {
+                                    apple_notifications::queue(ApnJob::from_ack(session.id, user.to_string(), sub.auth)).await;
+                                }
+                            }
+                        }
+                    }
+                };
+
+            }
+        },
+        AckEvent::AddMention { ids } => {
+            db.add_mention_to_unread(channel, user, ids).await?;
+        }
+    };
+
+    Ok(())
+}
+
 /// Start a new worker
-pub async fn worker(db: Database) {
+pub async fn worker(db: Database, authifier_db: authifier::Database) {
     let mut tasks = HashMap::<(String, String), DelayedTask<Task>>::new();
     let mut keys = vec![];
 
@@ -71,13 +108,7 @@ pub async fn worker(db: Database) {
                 let Task { event } = task.data;
                 let (user, channel) = key;
 
-                if let Err(err) = match &event {
+                if let Err(err) = handle_ack_event(&event, &db, &authifier_db, user, channel).await {
-                    #[allow(clippy::disallowed_methods)] // event is sent by higher level function
-                    AckEvent::AckMessage { id } => db.acknowledge_message(channel, user, id).await,
-                    AckEvent::AddMention { ids } => {
-                        db.add_mention_to_unread(channel, user, ids).await
-                    }
-                } {
                     error!("{err:?} for {event:?}. ({user}, {channel})");
                 } else {
                     info!("User {user} ack in {channel} with {event:?}");

crates/core/database/src/tasks/mod.rs

@@ -18,7 +18,7 @@ pub async fn start_workers(db: Database, authifier_db: authifier::Database) {
     task::spawn(apple_notifications::worker(db.clone()));
 
     for _ in 0..WORKER_COUNT {
-        task::spawn(ack::worker(db.clone()));
+        task::spawn(ack::worker(db.clone(), authifier_db.clone()));
         task::spawn(last_message_id::worker(db.clone()));
         task::spawn(process_embeds::worker(db.clone()));
         task::spawn(web_push::worker(authifier_db.clone()));