chore(monorepo): delta, january, quark

crates/quark/src/tasks/ack.rs

@@ -0,0 +1,123 @@
+// Queue Type: Debounced
+use crate::Database;
+
+use deadqueue::limited::Queue;
+use mongodb::bson::doc;
+use std::{collections::HashMap, time::Duration};
+
+use super::DelayedTask;
+
+/// Enumeration of possible events
+#[derive(Debug, Eq, PartialEq)]
+pub enum AckEvent {
+    /// Add mentions for a user in a channel
+    AddMention {
+        /// Message IDs
+        ids: Vec<String>,
+    },
+
+    /// 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: String,
+    /// Event
+    event: AckEvent,
+}
+
+#[derive(Debug)]
+struct Task {
+    event: AckEvent,
+}
+
+lazy_static! {
+    static ref Q: Queue<Data> = Queue::new(10_000);
+}
+
+/// Queue a new task for a worker
+pub async fn queue(channel: String, user: String, event: AckEvent) {
+    Q.try_push(Data {
+        channel,
+        user,
+        event,
+    })
+    .ok();
+
+    info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
+}
+
+/// Start a new worker
+pub async fn worker(db: Database) {
+    let mut tasks = HashMap::<(String, String), DelayedTask<Task>>::new();
+    let mut keys = 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) = match &event {
+                    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:?}");
+                }
+            }
+        }
+
+        // Clear keys
+        keys.clear();
+
+        // Queue incoming tasks.
+        while let Some(Data {
+            channel,
+            user,
+            mut event,
+        }) = Q.try_pop()
+        {
+            let key = (user, channel);
+            if let Some(mut task) = tasks.get_mut(&key) {
+                task.delay();
+
+                match &mut event {
+                    AckEvent::AddMention { ids } => {
+                        if let AckEvent::AddMention { ids: existing } = &mut task.data.event {
+                            existing.append(ids);
+                        } else {
+                            task.data.event = event;
+                        }
+                    }
+                    AckEvent::AckMessage { .. } => {
+                        task.data.event = event;
+                    }
+                }
+            } else {
+                tasks.insert(key, DelayedTask::new(Task { event }));
+            }
+        }
+
+        // Sleep for an arbitrary amount of time.
+        async_std::task::sleep(Duration::from_secs(1)).await;
+    }
+}

crates/quark/src/tasks/last_message_id.rs

@@ -0,0 +1,90 @@
+// Queue Type: Debounced
+use crate::{models::channel::PartialChannel, Database};
+
+use deadqueue::limited::Queue;
+use log::info;
+use mongodb::bson::doc;
+use std::{collections::HashMap, time::Duration};
+
+use super::DelayedTask;
+
+/// Task information
+struct Data {
+    /// Channel to update
+    channel: String,
+    /// Latest message ID
+    id: String,
+    /// Whether the channel is a DM
+    is_dm: bool,
+}
+
+/// Task information
+#[derive(Debug)]
+struct Task {
+    /// Latest message ID
+    id: String,
+    /// Whether the channel is a DM
+    is_dm: bool,
+}
+
+lazy_static! {
+    static ref Q: Queue<Data> = Queue::new(10_000);
+}
+
+/// Queue a new task for a worker
+pub async fn queue(channel: String, id: String, is_dm: bool) {
+    Q.try_push(Data { channel, id, is_dm }).ok();
+    info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
+}
+
+/// Start a new worker
+pub async fn worker(db: Database) {
+    let mut tasks = HashMap::<String, DelayedTask<Task>>::new();
+    let mut keys = 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 { id, is_dm, .. } = task.data;
+
+                let mut channel = PartialChannel {
+                    last_message_id: Some(id.to_string()),
+                    ..Default::default()
+                };
+
+                if is_dm {
+                    channel.active = Some(true);
+                }
+
+                match db.update_channel(key, &channel, vec![]).await {
+                    Ok(_) => info!("Updated last_message_id for {key} to {id}."),
+                    Err(err) => error!("Failed to update last_message_id with {err:?}!"),
+                }
+            }
+        }
+
+        // Clear keys
+        keys.clear();
+
+        // Queue incoming tasks.
+        while let Some(Data { channel, id, is_dm }) = Q.try_pop() {
+            if let Some(mut task) = tasks.get_mut(&channel) {
+                task.data.id = id;
+                task.delay();
+            } else {
+                tasks.insert(channel, DelayedTask::new(Task { id, is_dm }));
+            }
+        }
+
+        // Sleep for an arbitrary amount of time.
+        async_std::task::sleep(Duration::from_secs(1)).await;
+    }
+}

crates/quark/src/tasks/mod.rs

@@ -0,0 +1,58 @@
+//! Semi-important background task management
+
+use crate::Database;
+
+use async_std::task;
+use std::time::Instant;
+
+const WORKER_COUNT: usize = 5;
+
+pub mod ack;
+pub mod last_message_id;
+pub mod process_embeds;
+pub mod web_push;
+
+/// Spawn background workers
+pub async fn start_workers(db: Database) {
+    for _ in 0..WORKER_COUNT {
+        task::spawn(ack::worker(db.clone()));
+        task::spawn(last_message_id::worker(db.clone()));
+        task::spawn(process_embeds::worker(db.clone()));
+        task::spawn(web_push::worker(db.clone()));
+    }
+}
+
+/// Task with additional information on when it should run
+pub struct DelayedTask<T> {
+    pub data: T,
+    last_updated: Instant,
+    first_seen: Instant,
+}
+
+/// 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;
+
+impl<T> DelayedTask<T> {
+    /// Create a new delayed task
+    pub fn new(data: T) -> Self {
+        DelayedTask {
+            data,
+            last_updated: Instant::now(),
+            first_seen: Instant::now(),
+        }
+    }
+
+    /// Push a task further back in time
+    pub fn delay(&mut self) {
+        self.last_updated = Instant::now()
+    }
+
+    /// Check if a task should run yet
+    pub fn should_run(&self) -> bool {
+        self.first_seen.elapsed().as_secs() > EXPIRE_CONSTANT
+            || self.last_updated.elapsed().as_secs() > SAVE_CONSTANT
+    }
+}

crates/quark/src/tasks/process_embeds.rs

@@ -0,0 +1,57 @@
+use crate::util::variables::delta::{JANUARY_URL, MAX_EMBED_COUNT};
+use crate::{
+    models::{message::AppendMessage, Message},
+    types::january::Embed,
+    Database,
+};
+
+use deadqueue::limited::Queue;
+use log::error;
+
+/// Task information
+#[derive(Debug)]
+struct EmbedTask {
+    /// Channel we're processing the event in
+    channel: String,
+    /// ID of the message we're processing
+    id: String,
+    /// Content of the message
+    content: String,
+}
+
+lazy_static! {
+    static ref Q: Queue<EmbedTask> = Queue::new(10_000);
+}
+
+/// Queue a new task for a worker
+pub async fn queue(channel: String, id: String, content: String) {
+    Q.try_push(EmbedTask {
+        channel,
+        id,
+        content,
+    })
+    .ok();
+
+    info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
+}
+
+/// Start a new worker
+pub async fn worker(db: Database) {
+    loop {
+        let task = Q.pop().await;
+        if let Ok(embeds) = Embed::generate(task.content, &*JANUARY_URL, *MAX_EMBED_COUNT).await {
+            if let Err(err) = Message::append(
+                &db,
+                task.id,
+                task.channel,
+                AppendMessage {
+                    embeds: Some(embeds),
+                },
+            )
+            .await
+            {
+                error!("Encountered an error appending to message: {:?}", err);
+            }
+        }
+    }
+}

crates/quark/src/tasks/web_push.rs

@@ -0,0 +1,135 @@
+use crate::util::variables::delta::VAPID_PRIVATE_KEY;
+use crate::{bson::doc, r#impl::mongo::MongoDb, Database};
+
+use deadqueue::limited::Queue;
+use futures::StreamExt;
+use rauth::entities::Session;
+use web_push::{
+    ContentEncoding, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder, WebPushClient,
+    WebPushMessageBuilder,
+};
+
+/// Task information
+#[derive(Debug)]
+struct PushTask {
+    /// User IDs of the targets that are to receive this notification
+    recipients: Vec<String>,
+    /// Raw JSON payload to send to clients
+    payload: String,
+}
+
+lazy_static! {
+    static ref Q: Queue<PushTask> = Queue::new(10_000);
+}
+
+/// Queue a new task for a worker
+pub async fn queue(recipients: Vec<String>, payload: String) {
+    if recipients.is_empty() {
+        return;
+    }
+
+    Q.try_push(PushTask {
+        recipients,
+        payload,
+    })
+    .ok();
+
+    info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
+}
+
+/// Start a new worker
+pub async fn worker(db: Database) {
+    let client = WebPushClient::new();
+    let key = base64::decode_config(VAPID_PRIVATE_KEY.clone(), base64::URL_SAFE)
+        .expect("valid `VAPID_PRIVATE_KEY`");
+
+    if let Database::MongoDb(MongoDb(db)) = db {
+        loop {
+            let task = Q.pop().await;
+
+            // ! FIXME: this is hard-coded until rauth is merged into quark
+            if let Ok(mut cursor) = db
+                .database("revolt")
+                .collection::<Session>("sessions")
+                .find(
+                    doc! {
+                        "user_id": {
+                            "$in": task.recipients
+                        },
+                        "subscription": {
+                            "$exists": true
+                        }
+                    },
+                    None,
+                )
+                .await
+            {
+                while let Some(Ok(session)) = cursor.next().await {
+                    if let Some(sub) = session.subscription {
+                        let subscription = SubscriptionInfo {
+                            endpoint: sub.endpoint,
+                            keys: SubscriptionKeys {
+                                auth: sub.auth,
+                                p256dh: sub.p256dh,
+                            },
+                        };
+
+                        match WebPushMessageBuilder::new(&subscription) {
+                            Ok(mut builder) => {
+                                match VapidSignatureBuilder::from_pem(
+                                    std::io::Cursor::new(&key),
+                                    &subscription,
+                                ) {
+                                    Ok(sig_builder) => match sig_builder.build() {
+                                        Ok(signature) => {
+                                            builder.set_vapid_signature(signature);
+                                            builder.set_payload(
+                                                ContentEncoding::AesGcm,
+                                                task.payload.as_bytes(),
+                                            );
+
+                                            match builder.build() {
+                                                Ok(msg) => match client.send(msg).await {
+                                                    Ok(_) => {
+                                                        info!(
+                                                            "Sent Web Push notification to {:?}.",
+                                                            session.id
+                                                        )
+                                                    }
+                                                    Err(err) => {
+                                                        error!(
+                                                            "Hit error sending Web Push! {:?}",
+                                                            err
+                                                        )
+                                                    }
+                                                },
+                                                Err(err) => {
+                                                    error!(
+                                                        "Failed to build message for {}! {:?}",
+                                                        session.user_id, err
+                                                    )
+                                                }
+                                            }
+                                        }
+                                        Err(err) => error!(
+                                            "Failed to build signature for {}! {:?}",
+                                            session.user_id, err
+                                        ),
+                                    },
+                                    Err(err) => error!(
+                                        "Failed to create signature builder for {}! {:?}",
+                                        session.user_id, err
+                                    ),
+                                }
+                            }
+                            Err(err) => error!(
+                                "Invalid subscription information for {}! {:?}",
+                                session.user_id, err
+                            ),
+                        }
+                    }
+                }
+            }
+        }
+    }
+}