From 012a30ce61f6e5a5fb46842756cd7335d1b652ab Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 2 Nov 2021 13:30:36 +0000 Subject: [PATCH] feat(task_queue): ack task closes #111 and closes #86 --- src/database/entities/message.rs | 53 +++------- src/routes/channels/channel_ack.rs | 34 ++----- src/task_queue/mod.rs | 2 + src/task_queue/task_ack.rs | 136 +++++++++++++++++++++++++ src/task_queue/task_last_message_id.rs | 10 +- src/task_queue/task_process_embeds.rs | 10 +- src/task_queue/task_web_push.rs | 9 +- 7 files changed, 178 insertions(+), 76 deletions(-) create mode 100644 src/task_queue/task_ack.rs diff --git a/src/database/entities/message.rs b/src/database/entities/message.rs index 50fb417f..c6a07182 100644 --- a/src/database/entities/message.rs +++ b/src/database/entities/message.rs @@ -1,20 +1,16 @@ -use crate::util::variables::{USE_JANUARY, VAPID_PRIVATE_KEY, PUBLIC_URL}; +use crate::util::variables::{USE_JANUARY, PUBLIC_URL}; use crate::{ database::*, notifications::{events::ClientboundNotification, websocket::is_online}, util::result::{Error, Result}, }; -use futures::StreamExt; -use mongodb::options::UpdateOptions; use mongodb::{ bson::{doc, to_bson, DateTime, Document}, }; -use rauth::entities::{Model, Session}; use rocket::serde::json::Value; use serde::{Deserialize, Serialize}; use ulid::Ulid; -use web_push::{ContentEncoding, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder, WebPushClient, WebPushMessageBuilder}; use std::time::SystemTime; #[derive(Serialize, Deserialize, Debug, Clone)] @@ -211,9 +207,17 @@ impl Message { // if mentions { // spawn task_queue ( update channel_unreads ) // } - /*if let Some(mentions) = &self.mentions { - - }*/ + if let Some(mentions) = &self.mentions { + for user in mentions { + crate::task_queue::task_ack::queue( + channel.id().into(), + user.clone(), + crate::task_queue::task_ack::AckEvent::AddMention { + ids: vec![ self.id.clone() ] + } + ).await; + } + } // if (channel => DM | Group) | mentions { // spawn task_queue ( web push ) @@ -241,39 +245,6 @@ impl Message { } } - /* - - // ! FIXME: also temp code - // ! THIS ADDS ANY MENTIONS - if let Some(mentions) = &self.mentions { - let message = self.id.clone(); - let channel = self.channel.clone(); - let mentions = mentions.clone(); - async_std::task::spawn(async move { - get_collection("channel_unreads") - .update_many( - doc! { - "_id.channel": channel, - "_id.user": { - "$in": mentions - } - }, - doc! { - "$push": { - "mentions": message - } - }, - UpdateOptions::builder().upsert(true).build(), - ) - .await - /*.map_err(|_| Error::DatabaseError { - operation: "update_many", - with: "channel_unreads", - })?;*/ - .unwrap(); - }); - }*/ - Ok(()) } diff --git a/src/routes/channels/channel_ack.rs b/src/routes/channels/channel_ack.rs index ce56a2d2..45e24467 100644 --- a/src/routes/channels/channel_ack.rs +++ b/src/routes/channels/channel_ack.rs @@ -2,9 +2,6 @@ use crate::database::*; use crate::notifications::events::ClientboundNotification; use crate::util::result::{Error, Result, EmptyResponse}; -use mongodb::bson::doc; -use mongodb::options::UpdateOptions; - #[put("//ack/")] pub async fn req(user: User, target: Ref, message: Ref) -> Result { if user.bot.is_some() { @@ -22,31 +19,16 @@ pub async fn req(user: User, target: Ref, message: Ref) -> Result Err(Error::MissingPermission)? } - let id = target.id(); - get_collection("channel_unreads") - .update_one( - doc! { - "_id.channel": id, - "_id.user": &user.id - }, - doc! { - "$unset": { - "mentions": 1 - }, - "$set": { - "last_id": &message.id - } - }, - UpdateOptions::builder().upsert(true).build(), - ) - .await - .map_err(|_| Error::DatabaseError { - operation: "update_one", - with: "channel_unreads", - })?; + crate::task_queue::task_ack::queue( + target.id().into(), + user.id.clone(), + crate::task_queue::task_ack::AckEvent::AckMessage { + id: message.id.clone() + } + ).await; ClientboundNotification::ChannelAck { - id: id.to_string(), + id: target.id().into(), user: user.id.clone(), message_id: message.id, } diff --git a/src/task_queue/mod.rs b/src/task_queue/mod.rs index 5f27dd65..ff995ceb 100644 --- a/src/task_queue/mod.rs +++ b/src/task_queue/mod.rs @@ -1,9 +1,11 @@ pub mod task_last_message_id; pub mod task_process_embeds; pub mod task_web_push; +pub mod task_ack; pub async fn start_queues() { async_std::task::spawn(task_last_message_id::run()); async_std::task::spawn(task_process_embeds::run()); async_std::task::spawn(task_web_push::run()); + async_std::task::spawn(task_ack::run()); } diff --git a/src/task_queue/task_ack.rs b/src/task_queue/task_ack.rs new file mode 100644 index 00000000..b6ba265f --- /dev/null +++ b/src/task_queue/task_ack.rs @@ -0,0 +1,136 @@ +// Queue Type: Debounced +use std::{collections::HashMap, time::{Duration, Instant}}; + +use async_channel::{ Sender, Receiver, bounded }; +use mongodb::{bson::doc, options::UpdateOptions}; + +use crate::database::*; + +// Commit to database every 30 seconds if the task is particularly active. +static EXPIRE_CONSTANT: u64 = 30; + +// Otherwise, commit to database after 5 seconds. +static SAVE_CONSTANT: u64 = 5; + +struct Message { + channel: String, + user: String, + event: AckEvent, +} + +#[derive(Debug, Eq, PartialEq)] +pub enum AckEvent { + AddMention { ids: Vec }, + AckMessage { id: String } +} + +#[derive(Debug)] +struct Task { + event: AckEvent, + last_updated: Instant, + first_seen: Instant, +} + +lazy_static! { + static ref CHANNEL: (Sender, Receiver) = bounded(100); +} + +pub async fn queue(channel: String, user: String, event: AckEvent) { + CHANNEL.0.send(Message { channel, user, event }).await.ok(); +} + +pub async fn run() { + let unreads = get_collection("channel_unreads"); + let mut tasks = HashMap::<(String, String), Task>::new(); + let mut keys = vec![]; + + loop { + // Find due tasks. + for (key, Task { first_seen, last_updated, .. }) in &tasks { + if first_seen.elapsed().as_secs() > EXPIRE_CONSTANT || + last_updated.elapsed().as_secs() > SAVE_CONSTANT { + keys.push(key.clone()); + } + } + + // Commit any due tasks to the database. + for key in &keys { + if let Some(Task { event, .. }) = tasks.remove(key) { + let (user, channel) = key; + + match event { + AckEvent::AddMention { ids } => { + unreads.update_one( + doc! { + "_id.channel": channel, + "_id.user": user, + }, + doc! { + "$push": { + "mentions": { + "$each": ids + } + } + }, + UpdateOptions::builder().upsert(true).build(), + ) + .await + .ok(); + }, + AckEvent::AckMessage { id } => { + unreads.update_one( + doc! { + "_id.channel": channel, + "_id.user": user, + }, + doc! { + "$unset": { + "mentions": 1 + }, + "$set": { + "last_id": id + } + }, + UpdateOptions::builder().upsert(true).build(), + ) + .await + .ok(); + } + } + } + } + + // Clear keys + keys.clear(); + + // Queue incoming tasks. + while let Ok(Message { channel, user, mut event }) = CHANNEL.1.try_recv() { + let key = (user, channel); + if let Some(mut existing_task) = tasks.get_mut(&key) { + existing_task.last_updated = Instant::now(); + + match &mut event { + AckEvent::AddMention { ids } => { + if let AckEvent::AddMention { ids: existing } = &mut existing_task.event { + existing.append(ids); + } else { + existing_task.event = event; + } + } + AckEvent::AckMessage { .. } => { + existing_task.event = event; + } + } + } else { + tasks.insert(key, Task { + event, + last_updated: Instant::now(), + first_seen: Instant::now() + }); + } + } + + // Sleep for an arbitrary amount of time. + async_std::task::sleep(Duration::from_secs(1)).await; + } +} diff --git a/src/task_queue/task_last_message_id.rs b/src/task_queue/task_last_message_id.rs index 31d1b67d..54fc2e09 100644 --- a/src/task_queue/task_last_message_id.rs +++ b/src/task_queue/task_last_message_id.rs @@ -12,7 +12,11 @@ static EXPIRE_CONSTANT: u64 = 30; // Otherwise, commit to database after 5 seconds. static SAVE_CONSTANT: u64 = 5; -type Message = (String, String, bool); +struct Message { + channel: String, + id: String, + is_dm: bool +} #[derive(Debug)] struct Task { @@ -27,7 +31,7 @@ lazy_static! { } pub async fn queue(channel: String, id: String, is_dm: bool) { - CHANNEL.0.send((channel, id, is_dm)).await.ok(); + CHANNEL.0.send(Message { channel, id, is_dm }).await.ok(); } pub async fn run() { @@ -68,7 +72,7 @@ pub async fn run() { keys.clear(); // Queue incoming tasks. - while let Ok((channel, id, is_dm)) = CHANNEL.1.try_recv() { + while let Ok(Message { channel, id, is_dm }) = CHANNEL.1.try_recv() { if let Some(mut existing_task) = tasks.get_mut(&channel) { existing_task.id = id; existing_task.last_updated = Instant::now(); diff --git a/src/task_queue/task_process_embeds.rs b/src/task_queue/task_process_embeds.rs index ff40d097..86512f18 100644 --- a/src/task_queue/task_process_embeds.rs +++ b/src/task_queue/task_process_embeds.rs @@ -4,20 +4,24 @@ use mongodb::bson::{doc, to_bson}; use crate::{database::*, notifications::events::ClientboundNotification}; -type Message = (String, String, String); +struct Message { + channel: String, + id: String, + content: String +} lazy_static! { static ref CHANNEL: (Sender, Receiver) = bounded(100); } pub async fn queue(channel: String, id: String, content: String) { - CHANNEL.0.send((channel, id, content)).await.ok(); + CHANNEL.0.send(Message { channel, id, content }).await.ok(); } pub async fn run() { let messages = get_collection("messages"); - while let Ok((channel, id, content)) = CHANNEL.1.recv().await { + while let Ok(Message { channel, id, content }) = CHANNEL.1.recv().await { if let Ok(embeds) = Embed::generate(content).await { if let Ok(bson) = to_bson(&embeds) { if let Ok(_) = messages diff --git a/src/task_queue/task_web_push.rs b/src/task_queue/task_web_push.rs index 8110b0a2..2b1628fa 100644 --- a/src/task_queue/task_web_push.rs +++ b/src/task_queue/task_web_push.rs @@ -8,14 +8,17 @@ use futures::StreamExt; use crate::util::variables::VAPID_PRIVATE_KEY; use crate::database::*; -type Message = (Vec, String); +struct Message { + recipients: Vec, + payload: String +} lazy_static! { static ref CHANNEL: (Sender, Receiver) = bounded(100); } pub async fn queue(recipients: Vec, payload: String) { - CHANNEL.0.send((recipients, payload)).await.ok(); + CHANNEL.0.send(Message { recipients, payload }).await.ok(); } pub async fn run() { @@ -23,7 +26,7 @@ pub async fn run() { let key = base64::decode_config(VAPID_PRIVATE_KEY.clone(), base64::URL_SAFE) .expect("valid `VAPID_PRIVATE_KEY`"); - while let Ok((recipients, payload)) = CHANNEL.1.recv().await { + while let Ok(Message { recipients, payload }) = CHANNEL.1.recv().await { if let Ok(mut cursor) = Session::find( &get_db(), doc! {