diff --git a/Cargo.lock b/Cargo.lock index 9a4ddb28..795609e8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -654,6 +654,16 @@ dependencies = [ "cfg-if 1.0.0", ] +[[package]] +name = "crossbeam-queue" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f25d8400f4a7a5778f0e4e52384a48cbd9b5c495d110786187fc750075277a2" +dependencies = [ + "cfg-if 1.0.0", + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.5" @@ -781,6 +791,16 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3ee2393c4a91429dffb4bedf19f4d6abf27d8a732c8ce4980305d782e5426d57" +[[package]] +name = "deadqueue" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f13842c4b01ed03514527b625ae001f43b74158c7b0402a5369010cc38c3c8cf" +dependencies = [ + "crossbeam-queue", + "tokio 1.11.0", +] + [[package]] name = "derivative" version = "2.2.0" @@ -1247,9 +1267,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.3.4" +version = "0.3.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7f3675cfef6a30c8031cf9e6493ebdc3bb3272a3fea3923c4210d1830e6a472" +checksum = "37a82c6d637fc9515a4694bbf1cb2457b79d81ce52b3108bdeea58b07dd34a57" dependencies = [ "bytes 1.1.0", "fnv", @@ -1260,7 +1280,7 @@ dependencies = [ "indexmap", "slab", "tokio 1.11.0", - "tokio-util 0.6.8", + "tokio-util 0.7.1", "tracing", ] @@ -1440,7 +1460,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "h2 0.3.4", + "h2 0.3.13", "http", "http-body 0.4.3", "httparse", @@ -1664,6 +1684,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "linkify" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28d9967eb7d0bc31c39c6f52e8fce42991c0cd1f7a2078326f0b7a399a584c8d" +dependencies = [ + "memchr", +] + [[package]] name = "lock_api" version = "0.4.7" @@ -2929,9 +2958,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.5.4" +version = "1.5.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d07a8629359eb56f1e2fb1652bb04212c072a87ba68546a04065d525673ac461" +checksum = "1a11647b6b25ff05a515cb92c365cec08801e83423a235b51e231e1808747286" dependencies = [ "aho-corasick", "memchr", @@ -2955,15 +2984,16 @@ dependencies = [ [[package]] name = "reqwest" -version = "0.11.4" +version = "0.11.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "246e9f61b9bb77df069a947682be06e31ac43ea37862e244a69f177694ea6d22" +checksum = "46a1f7aa4f35e5e8b4160449f51afc758f0ce6454315a9fa7d0d113e958c41eb" dependencies = [ "base64 0.13.0", "bytes 1.1.0", "encoding_rs", "futures-core", "futures-util", + "h2 0.3.13", "http", "http-body 0.4.3", "hyper 0.14.12", @@ -2985,7 +3015,7 @@ dependencies = [ "wasm-bindgen", "wasm-bindgen-futures", "web-sys", - "winreg 0.7.0", + "winreg 0.10.1", ] [[package]] @@ -3010,6 +3040,7 @@ dependencies = [ "chrono", "ctrlc", "dashmap", + "deadqueue", "dotenv", "env_logger", "futures", @@ -3017,7 +3048,7 @@ dependencies = [ "impl_ops", "lazy_static", "lettre", - "linkify", + "linkify 0.6.0", "log", "lru", "many-to-many", @@ -3049,7 +3080,7 @@ dependencies = [ [[package]] name = "revolt-quark" version = "0.1.0" -source = "git+https://github.com/revoltchat/quark?rev=d65f49168c186b15a8465f16cbeccebac9b2bb7b#d65f49168c186b15a8465f16cbeccebac9b2bb7b" +source = "git+https://github.com/revoltchat/quark?rev=6e9713186e8cfef0b067aa87af1e71a66bb44fd0#6e9713186e8cfef0b067aa87af1e71a66bb44fd0" dependencies = [ "async-recursion", "async-std", @@ -3061,6 +3092,7 @@ dependencies = [ "futures", "impl_ops", "lazy_static", + "linkify 0.8.1", "log", "mongodb 2.1.0", "nanoid", @@ -3071,6 +3103,8 @@ dependencies = [ "pretty_env_logger", "rauth", "redis-kiss", + "regex", + "reqwest", "rocket", "rocket_okapi", "schemars", @@ -3498,12 +3532,12 @@ dependencies = [ [[package]] name = "serde_urlencoded" -version = "0.7.0" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "edfa57a7f8d9c1d260a549e7224100f6c43d43f9103e06dd8b4095a9b2b43ce9" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" dependencies = [ "form_urlencoded", - "itoa 0.4.8", + "itoa 1.0.1", "ryu", "serde", ] @@ -4053,6 +4087,20 @@ dependencies = [ "tokio 1.11.0", ] +[[package]] +name = "tokio-util" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0edfdeb067411dba2044da6d1cb2df793dd35add7888d73c16e3381ded401764" +dependencies = [ + "bytes 1.1.0", + "futures-core", + "futures-sink", + "pin-project-lite 0.2.7", + "tokio 1.11.0", + "tracing", +] + [[package]] name = "toml" version = "0.5.8" @@ -4089,9 +4137,21 @@ dependencies = [ "cfg-if 1.0.0", "log", "pin-project-lite 0.2.7", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e65ce065b4b5c53e73bb28912318cb8c9e9ad3921f1d669eb0e68b4c8143a2b" +dependencies = [ + "proc-macro2", + "quote 1.0.9", + "syn 1.0.76", +] + [[package]] name = "tracing-core" version = "0.1.20" @@ -4503,8 +4563,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "632f73e236b219150ea279196e54e610f5dbafa5d61786303d4da54f84e47fce" dependencies = [ "cfg-if 1.0.0", - "serde", - "serde_json", "wasm-bindgen-macro", ] @@ -4730,9 +4788,9 @@ dependencies = [ [[package]] name = "winreg" -version = "0.7.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0120db82e8a1e0b9fb3345a539c478767c0048d842860994d96113d5b667bd69" +checksum = "80d0f4e272c85def139476380b12f9ac60926689dd2e01d4923222f40580869d" dependencies = [ "winapi 0.3.9", ] diff --git a/Cargo.toml b/Cargo.toml index 860ce85c..de3c6f4e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ rmp-serde = { git = "https://github.com/insertish/msgpack-rust", rev = "5bf2c242 # async futures = "0.3.8" chrono = "0.4.15" +deadqueue = "0.2.1" async-channel = "1.6.1" reqwest = { version = "0.11.4", features = ["json"] } async-std = { version = "1.8.0", features = ["tokio1", "tokio02", "attributes"] } @@ -71,5 +72,5 @@ schemars = "0.8.8" rocket_okapi = { git = "https://github.com/insertish/okapi", rev = "dcf0df115596ee07a587a7a543cddf3d7944645b", features = [ "swagger" ] } # quark -revolt-quark = { git = "https://github.com/revoltchat/quark", rev = "d65f49168c186b15a8465f16cbeccebac9b2bb7b" } +revolt-quark = { git = "https://github.com/revoltchat/quark", rev = "6e9713186e8cfef0b067aa87af1e71a66bb44fd0" } # revolt-quark = { path = "../quark" } diff --git a/src/main.rs b/src/main.rs index fb36d705..1d5f8e86 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,9 +9,9 @@ extern crate lazy_static; extern crate ctrlc; pub mod routes; +pub mod tasks; pub mod util; pub mod version; -//pub mod task_queue; use log::info; use rauth::{ @@ -112,6 +112,9 @@ 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())); + let auth = Auth::new(mongo_db.database("revolt"), config); let rocket = rocket::build(); routes::mount(rocket) @@ -127,7 +130,6 @@ async fn main() { .manage(auth) .manage(db) .manage(cors.clone()) - //.manage(RatelimitState::new()) .attach(util::ratelimiter::RatelimitFairing) .attach(cors) .launch() diff --git a/src/routes/channels/message_edit.rs b/src/routes/channels/message_edit.rs index 0e353f0c..6cfd8d14 100644 --- a/src/routes/channels/message_edit.rs +++ b/src/routes/channels/message_edit.rs @@ -54,8 +54,8 @@ pub async fn req( }; // 1. Handle content update - if let Some(content) = edit.content { - partial.content = Some(Content::Text(content)); + if let Some(content) = &edit.content { + partial.content = Some(Content::Text(content.clone())); } // 2. Clear any auto generated embeds @@ -77,8 +77,19 @@ pub async fn req( } } - partial.embeds = message.embeds.clone(); + partial.embeds = Some(new_embeds); message.update(db, partial).await?; + + // Queue up a task for processing embeds + if let Some(content) = edit.content { + crate::tasks::process_embeds::queue( + message.channel.to_string(), + message.id.to_string(), + content, + ) + .await; + } + Ok(Json(message)) } diff --git a/src/routes/channels/message_send.rs b/src/routes/channels/message_send.rs index 6905a821..cf5f04f5 100644 --- a/src/routes/channels/message_send.rs +++ b/src/routes/channels/message_send.rs @@ -166,11 +166,20 @@ pub async fn message_send( } // 6. Set content - message.content = Content::Text(data.content); + message.content = Content::Text(data.content.clone()); // 7. Pass-through nonce value for clients message.nonce = Some(idempotency.into_key()); message.create(db).await?; + + // Queue up a task for processing embeds + crate::tasks::process_embeds::queue( + channel.id().to_string(), + message.id.to_string(), + data.content, + ) + .await; + Ok(Json(message)) } diff --git a/src/tasks/mod.rs b/src/tasks/mod.rs new file mode 100644 index 00000000..a5106325 --- /dev/null +++ b/src/tasks/mod.rs @@ -0,0 +1,15 @@ +//! Semi-important background task management + +use async_std::task; +use revolt_quark::Database; + +const WORKER_COUNT: usize = 10; + +pub mod process_embeds; + +/// Spawn background workers +pub async fn start_workers(db: Database) { + for _ in 0..WORKER_COUNT { + task::spawn(process_embeds::worker(db.clone())); + } +} diff --git a/src/tasks/process_embeds.rs b/src/tasks/process_embeds.rs new file mode 100644 index 00000000..c10bba7d --- /dev/null +++ b/src/tasks/process_embeds.rs @@ -0,0 +1,49 @@ +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 = 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); + } + } + } +}