From 4287d5835c96512ce5c73a95f779457c35da278a Mon Sep 17 00:00:00 2001
From: Pavel Koneski <pavel.koneski@gmail.com>
Date: Tue, 22 Apr 2025 15:41:36 -0700
Subject: [PATCH] Unregister signal handlers during engine shutdown

---
 .../signal.NtSignalState.cs                   | 11 ++++++++++
 .../signal.SimpleSignalState.cs               | 14 ++++++++++++-
 src/core/IronPython.Modules/signal.cs         | 21 ++++++++++++++++++-
 src/core/IronPython/Runtime/PythonContext.cs  | 11 ++++++++++
 4 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/src/core/IronPython.Modules/signal.NtSignalState.cs b/src/core/IronPython.Modules/signal.NtSignalState.cs
index 2066929de..0d75dedae 100644
--- a/src/core/IronPython.Modules/signal.NtSignalState.cs
+++ b/src/core/IronPython.Modules/signal.NtSignalState.cs
@@ -27,6 +27,17 @@ public NtSignalState(PythonContext pc) : base(pc) {
             }
 
 
+            protected override void Dispose(bool disposing) {
+                if (disposing) {
+                    NativeWindowsSignal.SetConsoleCtrlHandler(this.WinAllSignalsHandlerDelegate, false);
+                } else {
+                    var winAllSignalsHandlerDelegate = new NativeWindowsSignal.WinSignalsHandler(WindowsEventHandler);
+                    NativeWindowsSignal.SetConsoleCtrlHandler(winAllSignalsHandlerDelegate, false);
+                }
+                base.Dispose(disposing);
+            }
+
+
             // Our implementation of WinSignalsHandler
             private bool WindowsEventHandler(uint winSignal) {
                 bool retVal;
diff --git a/src/core/IronPython.Modules/signal.SimpleSignalState.cs b/src/core/IronPython.Modules/signal.SimpleSignalState.cs
index b9084d6f6..cc70f66f7 100644
--- a/src/core/IronPython.Modules/signal.SimpleSignalState.cs
+++ b/src/core/IronPython.Modules/signal.SimpleSignalState.cs
@@ -9,6 +9,8 @@
 using System;
 using System.Runtime.InteropServices;
 
+using Microsoft.Scripting.Hosting.Shell;
+
 using IronPython.Runtime;
 
 namespace IronPython.Modules {
@@ -18,13 +20,23 @@ private class SimpleSignalState : PythonSignalState {
 
             public SimpleSignalState(PythonContext pc) : base(pc) {
                 Console.CancelKeyPress += new ConsoleCancelEventHandler(Console_CancelKeyPress);
-                if (pc.Console is Microsoft.Scripting.Hosting.Shell.BasicConsole console) {
+                if (pc.Console is BasicConsole console) {
                     // in console hosting scenarios, we need to override the console handler of Ctrl+C
                     _consoleHandler = console.ConsoleCancelEventHandler;
                     console.ConsoleCancelEventHandler = null;
                 }
             }
 
+            protected override void Dispose(bool disposing) {
+                if (disposing) {
+                    Console.CancelKeyPress -= new ConsoleCancelEventHandler(Console_CancelKeyPress);
+                    if (_consoleHandler != null && DefaultContext.DefaultPythonContext.Console is BasicConsole console) {
+                        // restore the original console handler
+                        console.ConsoleCancelEventHandler = _consoleHandler;
+                    }
+                }
+                base.Dispose(disposing);
+            }
 
             private void Console_CancelKeyPress(object? sender, ConsoleCancelEventArgs e) {
                 int pySignal = !RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? SIGINT
diff --git a/src/core/IronPython.Modules/signal.cs b/src/core/IronPython.Modules/signal.cs
index e96b2e981..3d80d8a54 100644
--- a/src/core/IronPython.Modules/signal.cs
+++ b/src/core/IronPython.Modules/signal.cs
@@ -421,7 +421,7 @@ private static void SetPythonSignalState(CodeContext/*!*/ context, PythonSignalS
         /// <summary>
         /// This class is used to store the installed signal handlers.
         /// </summary>
-        private class PythonSignalState {
+        private class PythonSignalState : IDisposable {
             // this provides us with access to the Main thread's stack
             public readonly PythonContext SignalPythonContext;
 
@@ -511,6 +511,25 @@ protected void CallPythonHandler(int signum, object? handler) {
                     }
                 }
             }
+
+            public void Dispose() {
+                Dispose(true);
+                GC.SuppressFinalize(this);
+            }
+
+            ~PythonSignalState() {
+                Dispose(false);
+            }
+
+            protected virtual void Dispose(bool disposing) {
+                if (!_disposed) {
+                    if (disposing) {
+                        Array.Clear(PySignalToPyHandler, 0, PySignalToPyHandler.Length);
+                    }
+                    _disposed = true;
+                }
+            }
+            private bool _disposed = false;
         }
 
 
diff --git a/src/core/IronPython/Runtime/PythonContext.cs b/src/core/IronPython/Runtime/PythonContext.cs
index 27f563dfb..505942a41 100644
--- a/src/core/IronPython/Runtime/PythonContext.cs
+++ b/src/core/IronPython/Runtime/PythonContext.cs
@@ -545,11 +545,16 @@ public object GetModuleState(object key) {
         /// Sets per-runtime state used by a module.  The module should have a unique key for
         /// each piece of state it needs to store.
         /// </summary>
+        /// <remarks>
+        /// The previous piece of state (if any) is disposed if it implements IDisposable.
+        /// </remarks>
         public void SetModuleState(object key, object value) {
             EnsureModuleState();
 
             lock (_moduleState) {
+                _moduleState.TryGetValue(key, out object oldState);
                 _moduleState[key] = value;
+                (oldState as IDisposable)?.Dispose();
             }
         }
 
@@ -1319,6 +1324,12 @@ public override void Shutdown() {
             Flush(SharedContext, SystemStandardOut);
             Flush(SharedContext, SystemStandardError);
 
+            lock (_moduleState) {
+                foreach (var state in _moduleState.Values) {
+                    (state as IDisposable)?.Dispose();
+                }
+            }
+
             static void Flush(CodeContext context, object obj) {
                 if (obj is PythonIOModule._IOBase pf) {
                     if (!pf.closed)