multithreading and message filtering from multiple clients

5 files changed, 203 insertions, 43 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 6daeffd..f642e3f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -18,6 +18,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
 [[package]]
+name = "ahash"
+version = "0.8.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011"
+dependencies = [
+ "cfg-if",
+ "once_cell",
+ "version_check",
+ "zerocopy",
+]
+
+[[package]]
 name = "aho-corasick"
 version = "1.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -42,6 +54,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "allocator-api2"
+version = "0.2.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c6cb57a04249c6480766f7f7cef5467412af1490f8d1e243141daddada3264f"
+
+[[package]]
 name = "autocfg"
 version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -327,6 +345,10 @@ name = "hashbrown"
 version = "0.14.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "290f1a1d9242c78d09ce40a5e87e7554ee637af1351968159f4952f028f75604"
+dependencies = [
+ "ahash",
+ "allocator-api2",
+]
 
 [[package]]
 name = "http"
@@ -459,6 +481,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f"
 
 [[package]]
+name = "lru"
+version = "0.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d3262e75e648fce39813cb56ac41f3c3e3f65217ebf3844d818d1f9398cfb0dc"
+dependencies = [
+ "hashbrown",
+]
+
+[[package]]
 name = "memchr"
 version = "2.7.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -733,6 +764,7 @@ dependencies = [
 name = "seance-rs"
 version = "0.1.0"
 dependencies = [
+ "lru",
  "regex",
  "serde",
  "serde_regex",
@@ -740,6 +772,7 @@ dependencies = [
  "toml",
  "twilight-gateway",
  "twilight-http",
+ "twilight-model",
 ]
 
 [[package]]
@@ -1523,3 +1556,23 @@ checksum = "1931d78a9c73861da0134f453bb1f790ce49b2e30eba8410b4b79bac72b46a2d"
 dependencies = [
  "memchr",
 ]
+
+[[package]]
+name = "zerocopy"
+version = "0.7.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0"
+dependencies = [
+ "zerocopy-derive",
+]
+
+[[package]]
+name = "zerocopy-derive"
+version = "0.7.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
diff --git a/Cargo.toml b/Cargo.toml
index 3ff6966..43f49c6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,6 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+lru = "0.12.3"
 regex = "1.10.2"
 serde = { version = "1.0.196", features = [ "derive" ] }
 serde_regex = "1.1.0"
@@ -13,3 +14,4 @@ tokio = { version = "1.38.0", features = [ "rt", "macros" ] }
 toml = "0.8.8"
 twilight-gateway = "0.15.4"
 twilight-http = "0.15.4"
+twilight-model = "0.15.4"
diff --git a/src/config.rs b/src/config.rs
index 37a967f..26d3666 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -10,7 +10,7 @@ pub enum AutoProxyScope {
     Channel
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Clone)]
 pub enum PresenceMode {
     Online,
     Busy,
@@ -18,7 +18,7 @@ pub enum PresenceMode {
     Invisible,
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Clone)]
 pub enum AutoproxyConfig {
     Member(String),
     Latch {
@@ -28,20 +28,20 @@ pub enum AutoproxyConfig {
     }
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Clone)]
 pub enum AutoproxyLatchScope {
     Global,
     Server
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Clone)]
 pub struct PluralkitConfig {
     #[serde(with = "serde_regex")]
     pub message_pattern: Regex,
     pub api_token: String,
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Clone)]
 pub struct System {
     pub reference_user_id: String,
     pub members: Vec<Member>,
@@ -55,7 +55,7 @@ fn default_forward_pings() -> bool {
     false
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Clone)]
 pub struct Member {
     pub name: String,
     #[serde(with = "serde_regex")]
diff --git a/src/main.rs b/src/main.rs
index 02fba86..a63c4bd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,44 +1,28 @@
 mod config;
-use config::Config;
-use twilight_gateway::{Intents, Shard, ShardId};
+mod system;
 
-#[tokio::main(flavor = "current_thread")]
-async fn main() {
-    println!("Hello, world!");
+use std::thread::{self, JoinHandle};
+use tokio::runtime;
 
-    let config_str = include_str!("../config.toml");
-    let config = Config::load(config_str.to_string());
-
-    let token = config.systems.get("ashe-test").unwrap().members
-        .iter().find(|member| member.name == "test").unwrap()
-        .discord_token.clone();
-
-    println!("Token: {}", token);
-
-    let intents = Intents::GUILD_MEMBERS | Intents::GUILD_PRESENCES | Intents::GUILD_MESSAGES | Intents::MESSAGE_CONTENT;
-
-    let mut shard = Shard::new(ShardId::ONE, token, intents);
+use system::System;
 
-    loop {
-        let event = match shard.next_event().await {
-            Ok(event) => event,
-            Err(source) => {
-                println!("error receiving event");
-
-                if source.is_fatal() {
-                    break;
-                }
-
-                continue;
-            }
-        };
-
-        match event {
-            twilight_gateway::Event::MessageCreate(message) => println!("Message: {:?}", message),
-            twilight_gateway::Event::MessageUpdate(_) => println!("Message updated"),
-            twilight_gateway::Event::Ready(_) => println!("Bot ready!"),
-            _ => (),
-        }
+fn main() {
+    let config_str = include_str!("../config.toml");
+    let config = config::Config::load(config_str.to_string());
+
+    let handles : Vec<_> = config.systems.into_iter().map(|(system_name, system_config)| -> JoinHandle<()>  {
+        println!("Starting thread for system {}", system_name);
+        thread::spawn(move || {
+            let runtime = runtime::Builder::new_current_thread().enable_all().build().expect("Could not construct Tokio runtime");
+            runtime.block_on(async {
+                let mut system = System::new(system_config);
+                system.start_clients().await;
+            })
+        })
+    }).collect();
+
+    for thread_handle in handles.into_iter() {
+        thread_handle.join().expect("Child thread has panicked");
     }
 }
 
diff --git a/src/system.rs b/src/system.rs
new file mode 100644
index 0000000..e707d7f
--- /dev/null
+++ b/src/system.rs
@@ -0,0 +1,121 @@
+use std::{num::NonZeroUsize, str::FromStr};
+
+use std::sync::Arc;
+use tokio::task::JoinSet;
+use tokio::sync::Mutex;
+use twilight_http::Client;
+use twilight_model::id::{Id, marker::{MessageMarker, UserMarker}};
+use twilight_gateway::{Intents, Shard, ShardId};
+
+#[derive(Clone)]
+pub struct System {
+    pub config: crate::config::System,
+    pub message_dedup_cache: Arc<Mutex<lru::LruCache<Id<MessageMarker>, ()>>>,
+}
+
+impl System {
+    pub fn new(system_config: crate::config::System) -> Self {
+        System {
+            config: system_config,
+            message_dedup_cache: Arc::new(Mutex::new(lru::LruCache::new(NonZeroUsize::new(100).unwrap())))
+        }
+    }
+
+    pub async fn start_clients(&mut self) {
+        println!("Starting clients for system");
+
+        let reference_user_id : Id<UserMarker> = Id::from_str(self.config.reference_user_id.as_str())
+            .expect(format!("Invalid user ID: {}", self.config.reference_user_id).as_str());
+
+        let mut set = JoinSet::new();
+
+        for member in self.config.members.iter() {
+            let token = member.discord_token.clone();
+            let self_clone = self.clone();
+            set.spawn(async move {
+                println!("Starting client with token: {}", token);
+
+                let intents = Intents::GUILD_MEMBERS | Intents::GUILD_PRESENCES | Intents::GUILD_MESSAGES | Intents::MESSAGE_CONTENT;
+
+                let mut shard = Shard::new(ShardId::ONE, token.clone(), intents);
+                let mut client = Client::new(token.clone());
+
+                loop {
+                    let event = match shard.next_event().await {
+                        Ok(event) => event,
+                        Err(source) => {
+                            println!("error receiving event");
+
+                            if source.is_fatal() {
+                                break;
+                            }
+
+                            continue;
+                        }
+                    };
+
+                    match event {
+                        twilight_gateway::Event::Ready(client) => {
+                            println!("Bot started for {}#{}", client.user.name, client.user.discriminator);
+                        },
+
+                        twilight_gateway::Event::MessageCreate(message) => {
+                            if message.author.id != reference_user_id {
+                                continue
+                            }
+
+                            if self_clone.is_new_message(message.id).await {
+                                self_clone.handle_message_create(message, &mut client).await;
+                            }
+                        },
+
+                        twilight_gateway::Event::MessageUpdate(message) => {
+                            if message.author.is_none() || message.author.as_ref().unwrap().id != reference_user_id {
+                                continue
+                            }
+
+                            if self_clone.is_new_message(message.id).await {
+                                self_clone.handle_message_update(message, &mut client).await;
+                            }
+                        },
+
+                        _ => (),
+                    }
+                }
+            });
+        }
+
+        while let Some(join_result) = set.join_next().await {
+            if let Err(join_error) = join_result {
+                println!("Task encountered error: {}", join_error);
+            } else {
+                println!("Task joined cleanly");
+            }
+        }
+    }
+
+    async fn is_new_message(&self, message_id: Id<MessageMarker>) -> bool {
+        let mut message_cache = self.message_dedup_cache.lock().await;
+        if let None = message_cache.get(&message_id) {
+            message_cache.put(message_id, ());
+            true
+        } else {
+            false
+        }
+    }
+
+    async fn handle_message_create(&self, message: Box<twilight_model::gateway::payload::incoming::MessageCreate>, client: &mut Client) {
+        println!("Message created: {}", message.content);
+
+        if let Err(err) = client.create_message(message.channel_id)
+            .reply(message.id)
+            .content(&format!("Recognized message from authorized user {}", message.author.name)).expect("Error: reply too long")
+            .await {
+            println!("Error: {}", err);
+        }
+    }
+
+    async fn handle_message_update(&self, _message: Box<twilight_model::gateway::payload::incoming::MessageUpdate>, _client: &mut Client) {
+        // TODO: handle message edits and stuff
+    }
+}