feat(task_queue): ack task

closes #111 and closes #86
This commit is contained in:
Paul
2021-11-02 13:30:36 +00:00
parent 0025bea23d
commit 012a30ce61
7 changed files with 178 additions and 76 deletions

View File

@@ -1,20 +1,16 @@
use crate::util::variables::{USE_JANUARY, VAPID_PRIVATE_KEY, PUBLIC_URL};
use crate::util::variables::{USE_JANUARY, PUBLIC_URL};
use crate::{
database::*,
notifications::{events::ClientboundNotification, websocket::is_online},
util::result::{Error, Result},
};
use futures::StreamExt;
use mongodb::options::UpdateOptions;
use mongodb::{
bson::{doc, to_bson, DateTime, Document},
};
use rauth::entities::{Model, Session};
use rocket::serde::json::Value;
use serde::{Deserialize, Serialize};
use ulid::Ulid;
use web_push::{ContentEncoding, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder, WebPushClient, WebPushMessageBuilder};
use std::time::SystemTime;
#[derive(Serialize, Deserialize, Debug, Clone)]
@@ -211,9 +207,17 @@ impl Message {
// if mentions {
// spawn task_queue ( update channel_unreads )
// }
/*if let Some(mentions) = &self.mentions {
}*/
if let Some(mentions) = &self.mentions {
for user in mentions {
crate::task_queue::task_ack::queue(
channel.id().into(),
user.clone(),
crate::task_queue::task_ack::AckEvent::AddMention {
ids: vec![ self.id.clone() ]
}
).await;
}
}
// if (channel => DM | Group) | mentions {
// spawn task_queue ( web push )
@@ -241,39 +245,6 @@ impl Message {
}
}
/*
// ! FIXME: also temp code
// ! THIS ADDS ANY MENTIONS
if let Some(mentions) = &self.mentions {
let message = self.id.clone();
let channel = self.channel.clone();
let mentions = mentions.clone();
async_std::task::spawn(async move {
get_collection("channel_unreads")
.update_many(
doc! {
"_id.channel": channel,
"_id.user": {
"$in": mentions
}
},
doc! {
"$push": {
"mentions": message
}
},
UpdateOptions::builder().upsert(true).build(),
)
.await
/*.map_err(|_| Error::DatabaseError {
operation: "update_many",
with: "channel_unreads",
})?;*/
.unwrap();
});
}*/
Ok(())
}

View File

@@ -2,9 +2,6 @@ use crate::database::*;
use crate::notifications::events::ClientboundNotification;
use crate::util::result::{Error, Result, EmptyResponse};
use mongodb::bson::doc;
use mongodb::options::UpdateOptions;
#[put("/<target>/ack/<message>")]
pub async fn req(user: User, target: Ref, message: Ref) -> Result<EmptyResponse> {
if user.bot.is_some() {
@@ -22,31 +19,16 @@ pub async fn req(user: User, target: Ref, message: Ref) -> Result<EmptyResponse>
Err(Error::MissingPermission)?
}
let id = target.id();
get_collection("channel_unreads")
.update_one(
doc! {
"_id.channel": id,
"_id.user": &user.id
},
doc! {
"$unset": {
"mentions": 1
},
"$set": {
"last_id": &message.id
}
},
UpdateOptions::builder().upsert(true).build(),
)
.await
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "channel_unreads",
})?;
crate::task_queue::task_ack::queue(
target.id().into(),
user.id.clone(),
crate::task_queue::task_ack::AckEvent::AckMessage {
id: message.id.clone()
}
).await;
ClientboundNotification::ChannelAck {
id: id.to_string(),
id: target.id().into(),
user: user.id.clone(),
message_id: message.id,
}

View File

@@ -1,9 +1,11 @@
pub mod task_last_message_id;
pub mod task_process_embeds;
pub mod task_web_push;
pub mod task_ack;
pub async fn start_queues() {
async_std::task::spawn(task_last_message_id::run());
async_std::task::spawn(task_process_embeds::run());
async_std::task::spawn(task_web_push::run());
async_std::task::spawn(task_ack::run());
}

136
src/task_queue/task_ack.rs Normal file
View File

@@ -0,0 +1,136 @@
// Queue Type: Debounced
use std::{collections::HashMap, time::{Duration, Instant}};
use async_channel::{ Sender, Receiver, bounded };
use mongodb::{bson::doc, options::UpdateOptions};
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;
struct Message {
channel: String,
user: String,
event: AckEvent,
}
#[derive(Debug, Eq, PartialEq)]
pub enum AckEvent {
AddMention { ids: Vec<String> },
AckMessage { id: String }
}
#[derive(Debug)]
struct Task {
event: AckEvent,
last_updated: Instant,
first_seen: Instant,
}
lazy_static! {
static ref CHANNEL: (Sender<Message>, Receiver<Message>) = bounded(100);
}
pub async fn queue(channel: String, user: String, event: AckEvent) {
CHANNEL.0.send(Message { channel, user, event }).await.ok();
}
pub async fn run() {
let unreads = get_collection("channel_unreads");
let mut tasks = HashMap::<(String, String), Task>::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 { event, .. }) = tasks.remove(key) {
let (user, channel) = key;
match event {
AckEvent::AddMention { ids } => {
unreads.update_one(
doc! {
"_id.channel": channel,
"_id.user": user,
},
doc! {
"$push": {
"mentions": {
"$each": ids
}
}
},
UpdateOptions::builder().upsert(true).build(),
)
.await
.ok();
},
AckEvent::AckMessage { id } => {
unreads.update_one(
doc! {
"_id.channel": channel,
"_id.user": user,
},
doc! {
"$unset": {
"mentions": 1
},
"$set": {
"last_id": id
}
},
UpdateOptions::builder().upsert(true).build(),
)
.await
.ok();
}
}
}
}
// Clear keys
keys.clear();
// Queue incoming tasks.
while let Ok(Message { channel, user, mut event }) = CHANNEL.1.try_recv() {
let key = (user, channel);
if let Some(mut existing_task) = tasks.get_mut(&key) {
existing_task.last_updated = Instant::now();
match &mut event {
AckEvent::AddMention { ids } => {
if let AckEvent::AddMention { ids: existing } = &mut existing_task.event {
existing.append(ids);
} else {
existing_task.event = event;
}
}
AckEvent::AckMessage { .. } => {
existing_task.event = event;
}
}
} else {
tasks.insert(key, Task {
event,
last_updated: Instant::now(),
first_seen: Instant::now()
});
}
}
// Sleep for an arbitrary amount of time.
async_std::task::sleep(Duration::from_secs(1)).await;
}
}

View File

@@ -12,7 +12,11 @@ static EXPIRE_CONSTANT: u64 = 30;
// Otherwise, commit to database after 5 seconds.
static SAVE_CONSTANT: u64 = 5;
type Message = (String, String, bool);
struct Message {
channel: String,
id: String,
is_dm: bool
}
#[derive(Debug)]
struct Task {
@@ -27,7 +31,7 @@ lazy_static! {
}
pub async fn queue(channel: String, id: String, is_dm: bool) {
CHANNEL.0.send((channel, id, is_dm)).await.ok();
CHANNEL.0.send(Message { channel, id, is_dm }).await.ok();
}
pub async fn run() {
@@ -68,7 +72,7 @@ pub async fn run() {
keys.clear();
// Queue incoming tasks.
while let Ok((channel, id, is_dm)) = CHANNEL.1.try_recv() {
while let Ok(Message { 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();

View File

@@ -4,20 +4,24 @@ use mongodb::bson::{doc, to_bson};
use crate::{database::*, notifications::events::ClientboundNotification};
type Message = (String, String, String);
struct Message {
channel: String,
id: String,
content: String
}
lazy_static! {
static ref CHANNEL: (Sender<Message>, Receiver<Message>) = bounded(100);
}
pub async fn queue(channel: String, id: String, content: String) {
CHANNEL.0.send((channel, id, content)).await.ok();
CHANNEL.0.send(Message { channel, id, content }).await.ok();
}
pub async fn run() {
let messages = get_collection("messages");
while let Ok((channel, id, content)) = CHANNEL.1.recv().await {
while let Ok(Message { channel, id, content }) = CHANNEL.1.recv().await {
if let Ok(embeds) = Embed::generate(content).await {
if let Ok(bson) = to_bson(&embeds) {
if let Ok(_) = messages

View File

@@ -8,14 +8,17 @@ use futures::StreamExt;
use crate::util::variables::VAPID_PRIVATE_KEY;
use crate::database::*;
type Message = (Vec<String>, String);
struct Message {
recipients: Vec<String>,
payload: String
}
lazy_static! {
static ref CHANNEL: (Sender<Message>, Receiver<Message>) = bounded(100);
}
pub async fn queue(recipients: Vec<String>, payload: String) {
CHANNEL.0.send((recipients, payload)).await.ok();
CHANNEL.0.send(Message { recipients, payload }).await.ok();
}
pub async fn run() {
@@ -23,7 +26,7 @@ pub async fn run() {
let key = base64::decode_config(VAPID_PRIVATE_KEY.clone(), base64::URL_SAFE)
.expect("valid `VAPID_PRIVATE_KEY`");
while let Ok((recipients, payload)) = CHANNEL.1.recv().await {
while let Ok(Message { recipients, payload }) = CHANNEL.1.recv().await {
if let Ok(mut cursor) = Session::find(
&get_db(),
doc! {