diff --git a/Cargo.lock b/Cargo.lock index 53ec3d10..0b89fe1f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2794,6 +2794,7 @@ dependencies = [ name = "revolt" version = "0.0.0" dependencies = [ + "async-channel", "async-std", "async-tungstenite", "base64 0.13.0", diff --git a/Cargo.toml b/Cargo.toml index 1ed32c1e..9f87eafd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,6 +42,7 @@ rmp-serde = { git = "https://github.com/3Hren/msgpack-rust", rev = "5bf2c24203ad # async futures = "0.3.8" chrono = "0.4.15" +async-channel = "1.6.1" reqwest = { version = "0.11.4", features = ["json"] } async-std = { version = "1.8.0", features = ["tokio1", "tokio02", "attributes"] } diff --git a/src/database/entities/message.rs b/src/database/entities/message.rs index 4a0c4bfc..8b972005 100644 --- a/src/database/entities/message.rs +++ b/src/database/entities/message.rs @@ -181,10 +181,28 @@ impl Message { } pub async fn publish(self, channel: &Channel, process_embeds: bool) -> Result<()> { - // construct message and publish - // commit message to database + // Publish message event + ClientboundNotification::Message(self.clone()) + .publish(channel.id().to_string()); + + // Commit message to database + get_collection("messages") + .insert_one(to_bson(&self).unwrap().as_document().unwrap().clone(), None) + .await + .map_err(|_| Error::DatabaseError { + operation: "insert_one", + with: "message", + })?; // spawn task_queue ( update last_message_id ) + match channel { + Channel::DirectMessage { id, .. } => + crate::task_queue::task_last_message_id::queue(id.clone(), self.id.clone(), true).await, + Channel::Group { id, .. } | Channel::TextChannel { id, .. } => + crate::task_queue::task_last_message_id::queue(id.clone(), self.id.clone(), false).await, + _ => {} + } + // spawn task_queue ( process embeds ) // if mentions { @@ -195,14 +213,7 @@ impl Message { // spawn task_queue ( web push ) // } - get_collection("messages") - .insert_one(to_bson(&self).unwrap().as_document().unwrap().clone(), None) - .await - .map_err(|_| Error::DatabaseError { - operation: "insert_one", - with: "message", - })?; - + /* // ! FIXME: all this code is legitimately crap // ! rewrite when can be asked @@ -263,7 +274,6 @@ impl Message { } let mentions = self.mentions.clone(); - ClientboundNotification::Message(self.clone()).publish(channel.id().to_string()); /* Web Push Test Code @@ -333,7 +343,7 @@ impl Message { } } } - }); + });*/ Ok(()) } diff --git a/src/main.rs b/src/main.rs index 2beb3f59..00c616c6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -20,6 +20,7 @@ pub mod routes; pub mod redis; pub mod util; pub mod version; +pub mod task_queue; use async_std::task; use futures::join; @@ -50,6 +51,7 @@ async fn main() { database::connect().await; redis::connect().await; notifications::hive::init_hive().await; + task_queue::start_queues().await; ctrlc::set_handler(move || { // Force ungraceful exit to avoid hang. diff --git a/src/task_queue/mod.rs b/src/task_queue/mod.rs new file mode 100644 index 00000000..d013d5a8 --- /dev/null +++ b/src/task_queue/mod.rs @@ -0,0 +1,5 @@ +pub mod task_last_message_id; + +pub async fn start_queues() { + async_std::task::spawn(task_last_message_id::run()); +} diff --git a/src/task_queue/task_last_message_id.rs b/src/task_queue/task_last_message_id.rs new file mode 100644 index 00000000..db13fbb3 --- /dev/null +++ b/src/task_queue/task_last_message_id.rs @@ -0,0 +1,87 @@ +use std::{collections::HashMap, time::{Duration, Instant}}; + +use async_channel::{ Sender, Receiver, bounded }; +use mongodb::bson::doc; + +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; + +type Message = (String, String, bool); + +#[derive(Debug)] +struct Task { + id: String, + is_dm: bool, + last_updated: Instant, + first_seen: Instant, +} + +lazy_static! { + static ref CHANNEL: (Sender, Receiver) = bounded(100); +} + +pub async fn queue(channel: String, id: String, is_dm: bool) { + CHANNEL.0.send((channel, id, is_dm)).await.ok(); +} + +pub async fn run() { + let channels = get_collection("channels"); + let mut tasks = HashMap::::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 { id, is_dm, .. }) = tasks.remove(key) { + let mut set = doc! { "last_message_id": id.clone() }; + + if is_dm { + set.insert("active", true); + } + + channels + .update_one( + doc! { "_id": key }, + doc! { "$set": set }, + None, + ) + .await + .ok(); + } + } + + // Clear keys + keys.clear(); + + // Queue incoming tasks. + while let Ok((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(); + } else { + tasks.insert(channel, Task { + id, + is_dm, + last_updated: Instant::now(), + first_seen: Instant::now() + }); + } + } + + // Sleep for an arbitrary amount of time. + async_std::task::sleep(Duration::from_secs(1)).await; + } +}