jmug/stoatchat
1
0
Fork 1
mirror of https://github.com/stoatchat/stoatchat.git synced 2026-07-14 05:26:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: refactor tasks into quark

Browse Source
This commit is contained in:
Paul Makles
2022-04-21 13:52:35 +01:00
parent 3adbf29060
commit 5ce0c133e8
15 changed files with 18 additions and 551 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/main.rs
View File

@@ -9,7 +9,6 @@ extern crate lazy_static;
extern crate ctrlc;
pub mod routes;
pub mod tasks;
pub mod util;
pub mod version;
@@ -18,13 +17,13 @@ use rauth::{
config::{Captcha, Config, EmailVerification, SMTPSettings, Template, Templates},
logic::Auth,
};
use revolt_quark::DatabaseInfo;
use rocket_cors::AllowedOrigins;
use std::str::FromStr;
use util::variables::{
use revolt_quark::variables::delta::{
APP_URL, HCAPTCHA_KEY, INVITE_ONLY, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD, SMTP_USERNAME,
USE_EMAIL, USE_HCAPTCHA,
};
use revolt_quark::DatabaseInfo;
use rocket_cors::AllowedOrigins;
use std::str::FromStr;
#[async_std::main]
async fn main() {
@@ -36,7 +35,7 @@ async fn main() {
crate::version::VERSION
);
util::variables::preflight_checks();
revolt_quark::variables::delta::preflight_checks();
#[cfg(debug_assertions)]
ctrlc::set_handler(move || {
@@ -113,7 +112,7 @@ async fn main() {
rauth::entities::sync_models(&mongo_db.database("revolt")).await;
// Launch background task workers.
async_std::task::spawn(tasks::start_workers(db.clone()));
async_std::task::spawn(revolt_quark::tasks::start_workers(db.clone()));
let auth = Auth::new(mongo_db.database("revolt"), config);
let rocket = rocket::build();

3
src/routes/bots/create.rs
View File

@@ -1,8 +1,9 @@
use crate::util::{regex::RE_USERNAME, variables::MAX_BOT_COUNT};
use crate::util::regex::RE_USERNAME;
use nanoid::nanoid;
use revolt_quark::{
models::{user::BotInformation, Bot, User},
variables::delta::MAX_BOT_COUNT,
Db, Error, Result,
};

3
src/routes/channels/group_create.rs
View File

@@ -3,6 +3,7 @@ use std::{collections::HashSet, iter::FromIterator};
use revolt_quark::{
get_relationship,
models::{user::RelationshipStatus, Channel, User},
variables::delta::MAX_GROUP_SIZE,
Db, Error, Result,
};
@@ -11,8 +12,6 @@ use serde::{Deserialize, Serialize};
use ulid::Ulid;
use validator::Validate;
use crate::util::variables::MAX_GROUP_SIZE;
/// # Group Data
#[derive(Validate, Serialize, Deserialize, JsonSchema)]
pub struct DataCreateGroup {

2
src/routes/channels/message_edit.rs
View File

@@ -81,7 +81,7 @@ pub async fn req(
// Queue up a task for processing embeds
if let Some(content) = edit.content {
crate::tasks::process_embeds::queue(
revolt_quark::tasks::process_embeds::queue(
message.channel.to_string(),
message.id.to_string(),
content,

55
src/routes/channels/message_send.rs
View File

@@ -3,12 +3,9 @@ use std::collections::HashSet;
use revolt_quark::{
models::{
message::{Content, Masquerade, Reply, SendableEmbed},
Channel, Message, User,
Message, User,
},
perms,
presence::presence_filter_online,
types::push::PushNotification,
Db, Error, Permission, Ref, Result,
perms, Db, Error, Permission, Ref, Result,
};
use regex::Regex;
@@ -18,7 +15,6 @@ use ulid::Ulid;
use validator::Validate;
use crate::util::idempotency::IdempotencyKey;
use crate::util::variables::{APP_URL, AUTUMN_URL, PUBLIC_URL};
#[derive(Validate, Serialize, Deserialize, JsonSchema)]
pub struct DataMessageSend {
@@ -175,58 +171,15 @@ pub async fn message_send(
// 7. Pass-through nonce value for clients
message.nonce = Some(idempotency.into_key());
message.create(db).await?;
message.create(db, &channel, Some(&user)).await?;
// Queue up a task for processing embeds
crate::tasks::process_embeds::queue(
revolt_quark::tasks::process_embeds::queue(
channel.id().to_string(),
message.id.to_string(),
data.content,
)
.await;
// ! FIXME: this should be part of quark.
// ! Actually, so should be the thing above
// ! probably the entire task queue system
// ! should be moved
crate::tasks::web_push::queue(
{
let mut target_ids = vec![];
match &channel {
Channel::DirectMessage { recipients, .. } | Channel::Group { recipients, .. } => {
target_ids = (&recipients.iter().cloned().collect::<HashSet<String>>()
- &presence_filter_online(recipients).await)
.into_iter()
.collect::<Vec<String>>();
}
Channel::TextChannel { .. } => {
if let Some(mentions) = &message.mentions {
target_ids.append(&mut mentions.clone());
}
}
_ => {}
};
target_ids
},
json!(PushNotification::new(
message.clone(),
user,
channel.id(),
&*AUTUMN_URL,
&*PUBLIC_URL,
&*APP_URL,
))
.to_string(),
)
.await;
// ! ANOTHER ONE
crate::tasks::last_message_id::queue(
channel.id().to_string(),
message.id.to_string(),
channel.is_direct_dm(),
)
.await;
Ok(Json(message))
}

3
src/routes/invites/invite_join.rs
View File

@@ -1,13 +1,12 @@
use revolt_quark::{
models::{Channel, Invite, Server, User},
variables::delta::MAX_SERVER_COUNT,
Db, Error, Ref, Result,
};
use rocket::serde::json::Json;
use serde::Serialize;
use crate::util::variables::MAX_SERVER_COUNT;
/// # Join Response
#[derive(Serialize, JsonSchema)]
#[serde(tag = "type")]

5
src/routes/root.rs
View File

@@ -1,14 +1,13 @@
use crate::util::variables::{
use revolt_quark::variables::delta::{
APP_URL, AUTUMN_URL, EXTERNAL_WS_URL, HCAPTCHA_SITEKEY, INVITE_ONLY, JANUARY_URL, USE_AUTUMN,
USE_EMAIL, USE_HCAPTCHA, USE_JANUARY, USE_VOSO, VAPID_PUBLIC_KEY, VOSO_URL, VOSO_WS_HOST,
};
use revolt_quark::Result;
use rocket::http::Status;
use rocket::serde::json::Json;
use serde::Serialize;
use revolt_quark::Result;
/// # hCaptcha Configuration
#[derive(Serialize, JsonSchema, Debug)]
pub struct CaptchaFeature {

3
src/routes/servers/server_create.rs
View File

@@ -2,6 +2,7 @@ use std::collections::HashMap;
use revolt_quark::{
models::{Channel, Server, User},
variables::delta::MAX_SERVER_COUNT,
Db, Error, Result, DEFAULT_PERMISSION_SERVER,
};
@@ -10,8 +11,6 @@ use serde::{Deserialize, Serialize};
use ulid::Ulid;
use validator::Validate;
use crate::util::variables::MAX_SERVER_COUNT;
/// # Server Data
#[derive(Validate, Deserialize, JsonSchema)]
pub struct DataCreateServer {

102
src/tasks/ack.rs
View File

@@ -1,102 +0,0 @@
// 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;
}
}

79
src/tasks/last_message_id.rs
View File

@@ -1,79 +0,0 @@
// Queue Type: Debounced
use deadqueue::limited::Queue;
use log::info;
use mongodb::bson::doc;
use revolt_quark::{models::channel::PartialChannel, Database};
use std::{collections::HashMap, time::Duration};
use super::DelayedTask;
struct Data {
channel: String,
id: String,
is_dm: bool,
}
#[derive(Debug)]
struct Task {
id: String,
is_dm: bool,
}
lazy_static! {
static ref Q: Queue<Data> = Queue::new(10_000);
}
pub async fn queue(channel: String, id: String, is_dm: bool) {
Q.push(Data { channel, id, is_dm }).await;
}
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;
}
}

58
src/tasks/mod.rs
View File

@@ -1,58 +0,0 @@
//! Semi-important background task management
use std::time::Instant;
use async_std::task;
use revolt_quark::Database;
const WORKER_COUNT: usize = 5;
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
}
}

49
src/tasks/process_embeds.rs
View File

@@ -1,49 +0,0 @@
use crate::util::variables::{JANUARY_URL, MAX_EMBED_COUNT};
use deadqueue::limited::Queue;
use log::error;
use revolt_quark::{
models::{message::AppendMessage, Message},
types::january::Embed,
Database,
};
#[derive(Debug)]
struct EmbedTask {
channel: String,
id: String,
content: String,
}
lazy_static! {
static ref Q: Queue<EmbedTask> = Queue::new(10_000);
}
pub async fn queue(channel: String, id: String, content: String) {
Q.push(EmbedTask {
channel,
id,
content,
})
.await;
}
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);
}
}
}
}

91
src/tasks/web_push.rs
View File

@@ -1,91 +0,0 @@
use crate::util::variables::VAPID_PRIVATE_KEY;
use deadqueue::limited::Queue;
use futures::StreamExt;
use rauth::entities::Session;
use revolt_quark::{bson::doc, r#impl::mongo::MongoDb, Database};
use web_push::{
ContentEncoding, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder, WebPushClient,
WebPushMessageBuilder,
};
#[derive(Debug)]
struct PushTask {
recipients: Vec<String>,
payload: String,
}
lazy_static! {
static ref Q: Queue<PushTask> = Queue::new(10_000);
}
pub async fn queue(recipients: Vec<String>, payload: String) {
if recipients.is_empty() {
return;
}
Q.push(PushTask {
recipients,
payload,
})
.await;
}
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,
},
};
let mut builder = WebPushMessageBuilder::new(&subscription).unwrap();
let sig_builder = VapidSignatureBuilder::from_pem(
std::io::Cursor::new(&key),
&subscription,
)
.unwrap();
let signature = sig_builder.build().unwrap();
builder.set_vapid_signature(signature);
builder.set_payload(ContentEncoding::AesGcm, task.payload.as_bytes());
let msg = builder.build().unwrap();
match client.send(msg).await {
Ok(_) => info!("Sent Web Push notification to {:?}.", session.id),
Err(err) => error!("Hit error sending Web Push! {:?}", err),
}
}
}
}
}
}
}

1
src/util/mod.rs
View File

@@ -1,4 +1,3 @@
pub mod idempotency;
pub mod ratelimiter;
pub mod regex;
pub mod variables;

102
src/util/variables.rs
View File

@@ -1,102 +0,0 @@
use std::env;
#[cfg(debug_assertions)]
use log::warn;
lazy_static! {
// Application Settings
pub static ref REDIS_URI: String =
env::var("REVOLT_REDIS_URI").expect("Missing REVOLT_REDIS_URI environment variable.");
pub static ref WS_HOST: String =
env::var("REVOLT_WS_HOST").unwrap_or_else(|_| "0.0.0.0:9000".to_string());
pub static ref PUBLIC_URL: String =
env::var("REVOLT_PUBLIC_URL").expect("Missing REVOLT_PUBLIC_URL environment variable.");
pub static ref APP_URL: String =
env::var("REVOLT_APP_URL").expect("Missing REVOLT_APP_URL environment variable.");
pub static ref EXTERNAL_WS_URL: String =
env::var("REVOLT_EXTERNAL_WS_URL").expect("Missing REVOLT_EXTERNAL_WS_URL environment variable.");
pub static ref AUTUMN_URL: String =
env::var("AUTUMN_PUBLIC_URL").unwrap_or_else(|_| "https://example.com".to_string());
pub static ref JANUARY_URL: String =
env::var("JANUARY_PUBLIC_URL").unwrap_or_else(|_| "https://example.com".to_string());
pub static ref VOSO_URL: String =
env::var("VOSO_PUBLIC_URL").unwrap_or_else(|_| "https://example.com".to_string());
pub static ref VOSO_WS_HOST: String =
env::var("VOSO_WS_HOST").unwrap_or_else(|_| "wss://example.com".to_string());
pub static ref VOSO_MANAGE_TOKEN: String =
env::var("VOSO_MANAGE_TOKEN").unwrap_or_else(|_| "0".to_string());
pub static ref HCAPTCHA_KEY: String =
env::var("REVOLT_HCAPTCHA_KEY").unwrap_or_else(|_| "0x0000000000000000000000000000000000000000".to_string());
pub static ref HCAPTCHA_SITEKEY: String =
env::var("REVOLT_HCAPTCHA_SITEKEY").unwrap_or_else(|_| "10000000-ffff-ffff-ffff-000000000001".to_string());
pub static ref VAPID_PRIVATE_KEY: String =
env::var("REVOLT_VAPID_PRIVATE_KEY").expect("Missing REVOLT_VAPID_PRIVATE_KEY environment variable.");
pub static ref VAPID_PUBLIC_KEY: String =
env::var("REVOLT_VAPID_PUBLIC_KEY").expect("Missing REVOLT_VAPID_PUBLIC_KEY environment variable.");
// Application Flags
pub static ref INVITE_ONLY: bool = env::var("REVOLT_INVITE_ONLY").map_or(false, |v| v == "1");
pub static ref USE_EMAIL: bool = env::var("REVOLT_USE_EMAIL_VERIFICATION").map_or(
env::var("REVOLT_SMTP_HOST").is_ok()
&& env::var("REVOLT_SMTP_USERNAME").is_ok()
&& env::var("REVOLT_SMTP_PASSWORD").is_ok()
&& env::var("REVOLT_SMTP_FROM").is_ok(),
|v| v == *"1"
);
pub static ref USE_HCAPTCHA: bool = env::var("REVOLT_HCAPTCHA_KEY").is_ok();
pub static ref USE_AUTUMN: bool = env::var("AUTUMN_PUBLIC_URL").is_ok();
pub static ref USE_JANUARY: bool = env::var("JANUARY_PUBLIC_URL").is_ok();
pub static ref USE_VOSO: bool = env::var("VOSO_PUBLIC_URL").is_ok() && env::var("VOSO_MANAGE_TOKEN").is_ok();
// SMTP Settings
pub static ref SMTP_HOST: String =
env::var("REVOLT_SMTP_HOST").unwrap_or_else(|_| "".to_string());
pub static ref SMTP_USERNAME: String =
env::var("REVOLT_SMTP_USERNAME").unwrap_or_else(|_| "".to_string());
pub static ref SMTP_PASSWORD: String =
env::var("REVOLT_SMTP_PASSWORD").unwrap_or_else(|_| "".to_string());
pub static ref SMTP_FROM: String = env::var("REVOLT_SMTP_FROM").unwrap_or_else(|_| "".to_string());
// Application Logic Settings
pub static ref MAX_GROUP_SIZE: usize =
env::var("REVOLT_MAX_GROUP_SIZE").unwrap_or_else(|_| "50".to_string()).parse().unwrap();
pub static ref MAX_BOT_COUNT: usize =
env::var("REVOLT_MAX_BOT_COUNT").unwrap_or_else(|_| "5".to_string()).parse().unwrap();
pub static ref MAX_EMBED_COUNT: usize =
env::var("REVOLT_MAX_EMBED_COUNT").unwrap_or_else(|_| "5".to_string()).parse().unwrap();
pub static ref MAX_SERVER_COUNT: usize =
env::var("REVOLT_MAX_SERVER_COUNT").unwrap_or_else(|_| "100".to_string()).parse().unwrap();
pub static ref EARLY_ADOPTER_BADGE: i64 =
env::var("REVOLT_EARLY_ADOPTER_BADGE").unwrap_or_else(|_| "0".to_string()).parse().unwrap();
}
pub fn preflight_checks() {
format!("url = {}", *APP_URL);
format!("public = {}", *PUBLIC_URL);
format!("external = {}", *EXTERNAL_WS_URL);
format!("privkey = {}", *VAPID_PRIVATE_KEY);
format!("pubkey = {}", *VAPID_PUBLIC_KEY);
if !(*USE_EMAIL) {
#[cfg(not(debug_assertions))]
if !env::var("REVOLT_UNSAFE_NO_EMAIL").map_or(false, |v| v == *"1") {
panic!("Running in production without email is not recommended, set REVOLT_UNSAFE_NO_EMAIL=1 to override.");
}
#[cfg(debug_assertions)]
warn!("No SMTP settings specified! Remember to configure email.");
}
if !(*USE_HCAPTCHA) {
#[cfg(not(debug_assertions))]
if !env::var("REVOLT_UNSAFE_NO_CAPTCHA").map_or(false, |v| v == *"1") {
panic!("Running in production without CAPTCHA is not recommended, set REVOLT_UNSAFE_NO_CAPTCHA=1 to override.");
}
#[cfg(debug_assertions)]
warn!("No Captcha key specified! Remember to add hCaptcha key.");
}
}