chore: port ack task

This commit is contained in:
Paul Makles
2022-04-17 13:55:22 +01:00
parent 3d98f5dfbb
commit 3adbf29060
2 changed files with 105 additions and 1 deletions

102
src/tasks/ack.rs Normal file
View File

@@ -0,0 +1,102 @@
// Queue Type: Debounced
use deadqueue::limited::Queue;
use mongodb::bson::doc;
use revolt_quark::Database;
use std::{collections::HashMap, time::Duration};
use super::DelayedTask;
#[derive(Debug, Eq, PartialEq)]
pub enum AckEvent {
AddMention { ids: Vec<String> },
AckMessage { id: String },
}
struct Data {
channel: String,
user: String,
event: AckEvent,
}
#[derive(Debug)]
struct Task {
event: AckEvent,
}
lazy_static! {
static ref Q: Queue<Data> = Queue::new(10_000);
}
pub async fn queue(channel: String, user: String, event: AckEvent) {
Q.push(Data {
channel,
user,
event,
})
.await;
}
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})");
}
}
}
// 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;
}
}

View File

@@ -7,6 +7,7 @@ use revolt_quark::Database;
const WORKER_COUNT: usize = 5;
mod ack;
pub mod last_message_id;
pub mod process_embeds;
pub mod web_push;
@@ -14,9 +15,10 @@ 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::spawn(last_message_id::worker(db.clone()));
}
}