// Queue Type: Debounced use crate::{Database, Message, AMQP}; use deadqueue::limited::Queue; use once_cell::sync::Lazy; use revolt_config::capture_message; use revolt_models::v0::PushNotification; use std::{ collections::{HashMap, HashSet}, time::Duration, }; use validator::HasLen; use revolt_result::Result; use super::DelayedTask; use crate::Channel::TextChannel; /// Enumeration of possible events #[derive(Debug, Eq, PartialEq)] pub enum AckEvent { /// Add mentions for a channel ProcessMessage { /// push notification, message, recipients, push silenced messages: Vec<(Option, Message, Vec, bool)>, }, /// Acknowledge message in a channel for a user AckMessage { /// Message ID id: String, }, } /// Task information struct Data { /// Channel to ack channel: String, /// User to ack for user: Option, /// Event event: AckEvent, } #[derive(Debug)] struct Task { event: AckEvent, } static Q: Lazy> = Lazy::new(|| Queue::new(10_000)); /// Queue a new task for a worker pub async fn queue_ack(channel: String, user: String, event: AckEvent) { Q.try_push(Data { channel, user: Some(user), event, }) .ok(); info!( "Queue is using {} slots from {}. Queued type: ACK", Q.len(), Q.capacity() ); } /// Do not add more than one message per event. pub async fn queue_message(channel: String, event: AckEvent) { Q.try_push(Data { channel, user: None, event, }) .ok(); info!( "Queue is using {} slots from {}. Queued type: MENTION", Q.len(), Q.capacity() ); } pub async fn handle_ack_event( event: &AckEvent, db: &Database, amqp: &AMQP, user: &Option, channel: &str, ) -> Result<()> { match &event { #[allow(clippy::disallowed_methods)] // event is sent by higher level function AckEvent::AckMessage { id } => { let user = user.as_ref().unwrap(); let user: &str = user.as_str(); 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 Err(err) = amqp .ack_message(user.to_string(), channel.to_string(), id.to_owned()) .await { revolt_config::capture_error(&err); } }; } } AckEvent::ProcessMessage { messages } => { let mut users: HashSet<&String> = HashSet::new(); info!( "Processing {} messages from channel {}", messages.len(), messages[0].1.channel ); // find all the users we'll be notifying messages.iter().for_each(|(_, _, recipents, _)| { users.extend(recipents.iter()); }); info!("Found {} users to notify.", users.len()); for user in users { let message_ids: Vec = messages .iter() .filter_map(|(_, message, recipients, _)| { if recipients.contains(user) { Some(message.id.clone()) } else { None } }) .collect(); if !message_ids.is_empty() { db.add_mention_to_unread(channel, user, &message_ids) .await?; } info!("Added {} mentions for user {}", message_ids.len(), &user); } let mut mass_mentions = vec![]; for (push, message, recipients, silenced) in messages { if *silenced || push.is_none() || (recipients.is_empty() && !message.contains_mass_push_mention()) { debug!( "Rejecting push: silenced: {}, recipient count: {}, push exists: {:?}", *silenced, recipients.length(), push.is_some() ); continue; } debug!( "Sending push event to AMQP; message {} for {} users", push.as_ref().unwrap().message.id, recipients.len() ); if let Err(err) = amqp .message_sent(recipients.clone(), push.clone().unwrap()) .await { revolt_config::capture_error(&err); } if message.contains_mass_push_mention() { mass_mentions.push(push.clone().unwrap()); } } if !mass_mentions.is_empty() { debug!( "Sending mass mention push event to AMQP; channel {}", &mass_mentions[0].message.channel ); let channel = db .fetch_channel(&mass_mentions[0].message.channel) .await .expect("Failed to fetch channel from db"); if let TextChannel { server, .. } = channel { if let Err(err) = amqp.mass_mention_message_sent(server, mass_mentions).await { revolt_config::capture_error(&err); } } else { panic!("Unknown channel type when sending mass mention event"); } } } }; Ok(()) } /// Start a new worker pub async fn worker(db: Database, amqp: AMQP) { let mut tasks = HashMap::<(Option, String, u8), DelayedTask>::new(); let mut keys: Vec<(Option, String, u8)> = vec![]; loop { // Find due tasks. for (key, task) in &tasks { if task.should_run() { keys.push(key.clone()); } } // Commit any due tasks to the database. for key in &keys { if let Some(task) = tasks.remove(key) { let Task { event } = task.data; let (user, channel, _) = key; if let Err(err) = handle_ack_event(&event, &db, &amqp, user, channel).await { revolt_config::capture_error(&err); error!("{err:?} for {event:?}. ({user:?}, {channel})"); } else { info!("User {user:?} ack in {channel} with {event:?}"); } } } // Clear keys keys.clear(); // Queue incoming tasks. while let Some(Data { channel, user, mut event, }) = Q.try_pop() { info!("Took next ack from queue, now {} remaining", Q.len()); let key: (Option, String, u8) = ( user, channel, match &event { AckEvent::AckMessage { .. } => 0, AckEvent::ProcessMessage { .. } => 1, }, ); if let Some(task) = tasks.get_mut(&key) { match &mut event { AckEvent::ProcessMessage { messages: new_data } => { if let AckEvent::ProcessMessage { messages: existing } = &mut task.data.event { if let Some(new_event) = new_data.pop() { // if the message contains a mass mention, do not delay it any further. if new_event.1.contains_mass_push_mention() { // add the new message to the list of messages to be processed. existing.push(new_event); task.run_immediately(); continue; } existing.push(new_event); // put a cap on the amount of messages that can be queued, for particularly active channels if (existing.length() as u16) < revolt_config::config() .await .features .advanced .process_message_delay_limit { task.delay(); } } else { let err_msg = format!("Got zero-length message event: {event:?}"); capture_message(&err_msg, revolt_config::Level::Warning); info!("{err_msg}") } } else { panic!("Somehow got an ack message in the add mention arm"); } } AckEvent::AckMessage { .. } => { // replace the last acked message with the new acked message task.data.event = event; task.delay(); } } } else { tasks.insert(key, DelayedTask::new(Task { event })); } } // Sleep for an arbitrary amount of time. async_std::task::sleep(Duration::from_secs(1)).await; } }