From c88b4dc52e1686b4a35dc1fa031e1277c363a35f Mon Sep 17 00:00:00 2001
From: Andy Gayton <andy@thecablelounge.com>
Date: Fri, 21 Jun 2024 16:16:00 -0400
Subject: [PATCH] wip

---
 crates/nu-plugin-engine/src/gc.rs         | 13 ++++-
 crates/nu-plugin-engine/src/persistent.rs | 10 +---
 crates/nu-protocol/src/engine/ctrlc.rs    | 58 +++++++++++++++++++----
 3 files changed, 60 insertions(+), 21 deletions(-)

diff --git a/crates/nu-plugin-engine/src/gc.rs b/crates/nu-plugin-engine/src/gc.rs
index 572a61b2a9..00b66405ce 100644
--- a/crates/nu-plugin-engine/src/gc.rs
+++ b/crates/nu-plugin-engine/src/gc.rs
@@ -1,5 +1,5 @@
 use crate::PersistentPlugin;
-use nu_protocol::{PluginGcConfig, RegisteredPlugin};
+use nu_protocol::{engine::ctrlc, PluginGcConfig, RegisteredPlugin};
 use std::{
     sync::{mpsc, Arc, Weak},
     thread,
@@ -14,6 +14,7 @@ use std::{
 #[derive(Debug, Clone)]
 pub struct PluginGc {
     sender: mpsc::Sender<PluginGcMsg>,
+    _ctrlc_guard: Arc<Option<ctrlc::HandlerGuard>>,
 }
 
 impl PluginGc {
@@ -21,6 +22,7 @@ impl PluginGc {
     pub fn new(
         config: PluginGcConfig,
         plugin: &Arc<PersistentPlugin>,
+        ctrlc_handlers: Option<ctrlc::CtrlcHandlers>,
     ) -> std::io::Result<PluginGc> {
         let (sender, receiver) = mpsc::channel();
 
@@ -37,7 +39,14 @@ impl PluginGc {
             .name(format!("plugin gc ({})", plugin.identity().name()))
             .spawn(move || state.run(receiver))?;
 
-        Ok(PluginGc { sender })
+        let guard = ctrlc_handlers.map(|ctrlc_handlers| {
+            let sender = sender.clone();
+            ctrlc_handlers.add(Box::new(move || {
+                let _ = sender.send(PluginGcMsg::Ctrlc);
+            }))
+        });
+
+        Ok(PluginGc { sender, _ctrlc_guard: Arc::new(guard) })
     }
 
     /// Update the garbage collector config
diff --git a/crates/nu-plugin-engine/src/persistent.rs b/crates/nu-plugin-engine/src/persistent.rs
index bec55b94e7..3b66fd3dd4 100644
--- a/crates/nu-plugin-engine/src/persistent.rs
+++ b/crates/nu-plugin-engine/src/persistent.rs
@@ -1,5 +1,4 @@
 use crate::{
-    gc::PluginGcMsg,
     init::{create_command, make_plugin_interface},
     PluginGc,
 };
@@ -180,14 +179,7 @@ impl PersistentPlugin {
         })?;
 
         // Start the plugin garbage collector
-        let gc = PluginGc::new(mutable.gc_config.clone(), &self)?;
-
-        if let Some(ref ctrlc_handlers) = ctrlc_handlers {
-            let tx = gc.clone_sender();
-            ctrlc_handlers.add(Box::new(move || {
-                let _ = tx.send(PluginGcMsg::Ctrlc);
-            }));
-        }
+        let gc = PluginGc::new(mutable.gc_config.clone(), &self, ctrlc_handlers.clone())?;
 
         let pid = child.id();
         let interface = make_plugin_interface(
diff --git a/crates/nu-protocol/src/engine/ctrlc.rs b/crates/nu-protocol/src/engine/ctrlc.rs
index 9fb9d32e84..f5ff966fac 100644
--- a/crates/nu-protocol/src/engine/ctrlc.rs
+++ b/crates/nu-protocol/src/engine/ctrlc.rs
@@ -1,34 +1,72 @@
-use std::sync::{Arc, Mutex};
+use std::fmt::Debug;
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc, Mutex,
+};
 
 type CtrlcHandler = Box<dyn Fn() + Send + Sync>;
 
 #[derive(Clone)]
 pub struct CtrlcHandlers {
-    handlers: Arc<Mutex<Vec<CtrlcHandler>>>,
+    handlers: Arc<Mutex<Vec<(usize, CtrlcHandler)>>>,
+    next_id: Arc<AtomicUsize>,
+}
+
+#[derive(Clone)]
+pub struct HandlerGuard {
+    id: usize,
+    handlers: Arc<Mutex<Vec<(usize, CtrlcHandler)>>>,
+}
+
+impl Drop for HandlerGuard {
+    fn drop(&mut self) {
+        if let Ok(mut handlers) = self.handlers.lock() {
+            handlers.retain(|(id, _)| *id != self.id);
+        }
+    }
+}
+
+impl Debug for HandlerGuard {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("HandlerGuard")
+            .field("id", &self.id)
+            .finish()
+    }
 }
 
 impl CtrlcHandlers {
     pub fn new() -> CtrlcHandlers {
-        CtrlcHandlers {
-            handlers: Arc::new(Mutex::new(vec![])),
-        }
+        let handlers = Arc::new(Mutex::new(vec![]));
+        let next_id = Arc::new(AtomicUsize::new(0));
+        CtrlcHandlers { handlers, next_id }
     }
 
-    pub fn add(&self, handler: CtrlcHandler) {
-        if let Some(mut handlers) = self.handlers.lock().ok() {
-            handlers.push(handler);
+    pub fn add(&self, handler: CtrlcHandler) -> HandlerGuard {
+        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
+        if let Ok(mut handlers) = self.handlers.lock() {
+            handlers.push((id, handler));
+        }
+        HandlerGuard {
+            id,
+            handlers: Arc::clone(&self.handlers),
         }
     }
 
     pub fn run(&self) {
-        if let Some(handlers) = self.handlers.lock().ok() {
-            for handler in handlers.iter() {
+        if let Ok(handlers) = self.handlers.lock() {
+            for (_, handler) in handlers.iter() {
                 handler();
             }
         }
     }
 }
 
+impl Default for CtrlcHandlers {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl std::fmt::Debug for CtrlcHandlers {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("CtrlcHandlers")