feat(core): implement message tasks

includes:
- acknowledgements
- last_message_id
- embeds
- web push

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

@@ -0,0 +1,122 @@
+// Queue Type: Debounced
+use crate::Database;
+
+use deadqueue::limited::Queue;
+use once_cell::sync::Lazy;
+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,
+}
+
+static Q: Lazy<Queue<Data>> = Lazy::new(|| 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 {
+                    #[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:?}");
+                }
+            }
+        }
+
+        // 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(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/core/database/src/tasks/last_message_id.rs

@@ -0,0 +1,87 @@
+// Queue Type: Debounced
+use deadqueue::limited::Queue;
+use once_cell::sync::Lazy;
+use std::{collections::HashMap, time::Duration};
+
+use crate::{Database, PartialChannel};
+
+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,
+}
+
+static Q: Lazy<Queue<Data>> = Lazy::new(|| 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(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/core/database/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().into()));
+    }
+}
+
+/// 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/core/database/src/tasks/process_embeds.rs

@@ -0,0 +1,170 @@
+use crate::{models::Message, AppendMessage, Database};
+
+use futures::future::join_all;
+use linkify::{LinkFinder, LinkKind};
+use regex::Regex;
+use revolt_config::config;
+use revolt_result::Result;
+
+use async_lock::Semaphore;
+use async_std::task::spawn;
+use deadqueue::limited::Queue;
+use once_cell::sync::Lazy;
+use revolt_models::v0::Embed;
+use std::{collections::HashSet, sync::Arc};
+
+use isahc::prelude::*;
+
+/// 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,
+}
+
+static Q: Lazy<Queue<EmbedTask>> = Lazy::new(|| 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) {
+    let semaphore = Arc::new(Semaphore::new(
+        config().await.api.workers.max_concurrent_connections,
+    ));
+
+    loop {
+        let task = Q.pop().await;
+        let db = db.clone();
+        let semaphore = semaphore.clone();
+
+        spawn(async move {
+            let config = config().await;
+            let embeds = generate(
+                task.content,
+                &config.hosts.january,
+                config.features.limits.default.message_embeds,
+                semaphore,
+            )
+            .await;
+
+            if let Ok(embeds) = embeds {
+                if let Err(err) = Message::append(
+                    &db,
+                    task.id,
+                    task.channel,
+                    AppendMessage {
+                        embeds: Some(embeds),
+                    },
+                )
+                .await
+                {
+                    error!("Encountered an error appending to message: {:?}", err);
+                }
+            }
+        });
+    }
+}
+
+static RE_CODE: Lazy<Regex> = Lazy::new(|| Regex::new("```(?:.|\n)+?```|`(?:.|\n)+?`").unwrap());
+static RE_IGNORED: Lazy<Regex> = Lazy::new(|| Regex::new("(<http.+>)").unwrap());
+
+pub async fn generate(
+    content: String,
+    host: &str,
+    max_embeds: usize,
+    semaphore: Arc<Semaphore>,
+) -> Result<Vec<Embed>> {
+    // Ignore code blocks.
+    let content = RE_CODE.replace_all(&content, "");
+
+    // Ignore all content between angle brackets starting with http.
+    let content = RE_IGNORED.replace_all(&content, "");
+
+    let content = content
+        // Ignore quoted lines.
+        .split('\n')
+        .map(|v| {
+            if let Some(c) = v.chars().next() {
+                if c == '>' {
+                    return "";
+                }
+            }
+
+            v
+        })
+        .collect::<Vec<&str>>()
+        .join("\n");
+
+    let mut finder = LinkFinder::new();
+    finder.kinds(&[LinkKind::Url]);
+
+    // Process all links, stripping anchors and
+    // only taking up to `max_embeds` of links.
+    let links: Vec<String> = finder
+        .links(&content)
+        .map(|x| {
+            x.as_str()
+                .chars()
+                .take_while(|&ch| ch != '#')
+                .collect::<String>()
+        })
+        .collect::<HashSet<String>>()
+        .into_iter()
+        .take(max_embeds)
+        .collect();
+
+    // If no links, fail out.
+    if links.is_empty() {
+        return Err(create_error!(LabelMe));
+    }
+
+    // ! FIXME: batch request to january
+    let mut tasks = Vec::new();
+
+    for link in links {
+        let semaphore = semaphore.clone();
+        let host = host.to_string();
+        tasks.push(spawn(async move {
+            let guard = semaphore.acquire().await;
+
+            if let Ok(mut response) = isahc::get_async(format!(
+                "{host}/embed?url={}",
+                url_escape::encode_component(&link)
+            ))
+            .await
+            {
+                drop(guard);
+                response.json::<Embed>().await.ok()
+            } else {
+                None
+            }
+        }));
+    }
+
+    let embeds = join_all(tasks)
+        .await
+        .into_iter()
+        .flatten()
+        .collect::<Vec<Embed>>();
+
+    // Prevent database update when no embeds are found.
+    if !embeds.is_empty() {
+        Ok(embeds)
+    } else {
+        Err(create_error!(LabelMe))
+    }
+}

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

@@ -0,0 +1,122 @@
+use std::collections::HashSet;
+
+use authifier::Database;
+use base64::{
+    alphabet,
+    engine::{self},
+    Engine as _,
+};
+use deadqueue::limited::Queue;
+use once_cell::sync::Lazy;
+use revolt_presence::filter_online;
+use web_push::{
+    ContentEncoding, IsahcWebPushClient, 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,
+}
+
+static Q: Lazy<Queue<PushTask>> = Lazy::new(|| 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;
+    }
+
+    let online_ids = filter_online(&recipients).await;
+    let recipients = (&recipients.into_iter().collect::<HashSet<String>>() - &online_ids)
+        .into_iter()
+        .collect::<Vec<String>>();
+
+    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 = IsahcWebPushClient::new().unwrap();
+    let key = engine::GeneralPurpose::new(&alphabet::URL_SAFE, engine::general_purpose::NO_PAD)
+        .decode(std::env::var("REVOLT_VAPID_PRIVATE_KEY").unwrap())
+        .expect("valid `VAPID_PRIVATE_KEY`");
+
+    loop {
+        let task = Q.pop().await;
+
+        if let Ok(sessions) = db.find_sessions_with_subscription(&task.recipients).await {
+            for session in sessions {
+                if let Some(sub) = session.subscription {
+                    if sub.endpoint == "fcm" {
+                        // Use Firebase Cloud Messaging
+                        info!("Skipping FCM client, unimplemented...");
+                    } else {
+                        // Use Web Push Standard
+                        let subscription = SubscriptionInfo {
+                            endpoint: sub.endpoint,
+                            keys: SubscriptionKeys {
+                                auth: sub.auth,
+                                p256dh: sub.p256dh,
+                            },
+                        };
+
+                        match VapidSignatureBuilder::from_pem(
+                            std::io::Cursor::new(&key),
+                            &subscription,
+                        ) {
+                            Ok(sig_builder) => match sig_builder.build() {
+                                Ok(signature) => {
+                                    let mut builder = WebPushMessageBuilder::new(&subscription);
+                                    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
+                            ),
+                        }
+                    }
+                }
+            }
+        }
+    }
+}