From c97cc92df22b6447e4539ab1c67d929377a1c0e4 Mon Sep 17 00:00:00 2001
From: cat_or_not <41955154+catornot@users.noreply.github.com>
Date: Thu, 2 Jan 2025 16:11:24 -0500
Subject: [PATCH] drop the runtime on reload

---
 src/discord.rs |  7 +++++--
 src/lib.rs     | 26 +++++++++++++++++++++++---
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/discord.rs b/src/discord.rs
index 69d227a..b1cf5f9 100644
--- a/src/discord.rs
+++ b/src/discord.rs
@@ -1,6 +1,6 @@
 #![deny(non_snake_case)]
 
-use std::num::NonZeroU32;
+use std::{num::NonZeroU32, sync::atomic::Ordering};
 
 use discord_sdk::{
     activity::{events::ActivityEvent, ActivityBuilder, Assets, JoinRequestReply, PartyPrivacy},
@@ -46,7 +46,10 @@ pub async fn async_main() {
 
     let mut events = client.wheel.activity().0;
 
-    loop {
+    while PLUGIN
+        .get()
+        .is_some_and(|plugin| plugin.is_active.load(Ordering::Acquire))
+    {
         let data = activity.lock().clone();
 
         if let Some(img) = &data.large_image {
diff --git a/src/lib.rs b/src/lib.rs
index 9365689..d14affc 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,10 +1,13 @@
 #![allow(non_snake_case)]
 
+use std::sync::atomic::{AtomicBool, Ordering};
+
 use discord_sdk::activity::Secrets;
 use parking_lot::Mutex;
 use rrplug::prelude::*;
 use rrplug::{bindings::plugin_abi::PluginColor, interfaces::manager::register_interface};
 use tokio::runtime::Runtime;
+use tokio::task::JoinHandle;
 
 use crate::{
     discord::async_main,
@@ -39,6 +42,8 @@ pub struct ActivityData {
 pub struct DiscordRpcPlugin {
     pub activity: Mutex<ActivityData>,
     pub presence_data: Mutex<(GameStateStruct, UIPresenceStruct)>,
+    pub runtime: Mutex<Option<(Runtime, JoinHandle<()>)>>,
+    pub is_active: AtomicBool,
 }
 
 #[deny(non_snake_case)]
@@ -66,16 +71,22 @@ impl Plugin for DiscordRpcPlugin {
             ..Default::default()
         });
 
-        std::thread::spawn(|| match Runtime::new() {
-            Ok(rt) => rt.block_on(async_main()),
+        let rt = match Runtime::new() {
+            Ok(rt) => {
+                let handle = rt.spawn(async_main());
+                Some((rt, handle))
+            }
             Err(err) => {
                 log::error!("failed to create a runtime; {:?}", err);
+                None
             }
-        });
+        };
 
         Self {
             activity,
             presence_data: Mutex::new((GameStateStruct::default(), UIPresenceStruct::default())),
+            runtime: Mutex::new(rt),
+            is_active: AtomicBool::new(true),
         }
     }
 
@@ -89,6 +100,15 @@ impl Plugin for DiscordRpcPlugin {
     }
 
     fn on_reload_request(&self) -> reloading::ReloadResponse {
+        let (runtime, handle) = self
+            .runtime
+            .lock()
+            .take()
+            .expect("runtime should exist when unloading the plugin");
+        self.is_active.store(false, Ordering::Release);
+        runtime.block_on(async { _ = handle.await });
+        runtime.shutdown_background();
+
         log::info!("reloading request for DiscordRpc");
         // SAFETY: this plugin doesn't have anything that gets called or referenced from the game (usally) so it's safe
         unsafe { reloading::ReloadResponse::allow_reload() }