feat: add link embed generation

2026-07-14 05:26:59 +00:00 · 2022-04-14 20:46:46 +01:00
parent d2f0fb3414
commit 87391a49a5
7 changed files with 170 additions and 25 deletions
--- 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",
 ]
--- 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" }
--- 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()
--- 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))
 }
--- 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))
 }
--- a/src/tasks/mod.rs
+++ 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()));
+    }
+}
--- a/src/tasks/process_embeds.rs
+++ 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<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);
+            }
+        }
+    }
+}