Allow user code to register for a callback on shutdown.

Benoit Hudson · filmor · commit 45b9130cd38d · 2018-11-06T13:17:29.000+01:00
This is needed to cleanly close sockets, UI, etc.
diff --git a/src/embed_tests/pyinitialize.cs b/src/embed_tests/pyinitialize.cs
@@ -74,5 +74,65 @@ public void ReInitialize()
             }
             PythonEngine.Shutdown();
         }
+
+        [Test]
+        public void TestScopeIsShutdown()
+        {
+            PythonEngine.Initialize();
+            var scope = PyScopeManager.Global.Create("test");
+            PythonEngine.Shutdown();
+            Assert.That(PyScopeManager.Global.Contains("test"), Is.False);
+        }
+
+        /// <summary>
+        /// Helper for testing the shutdown handlers.
+        /// </summary>
+        int shutdown_count = 0;
+        void OnShutdownIncrement()
+        {
+            shutdown_count++;
+        }
+        void OnShutdownDouble()
+        {
+            shutdown_count *= 2;
+        }
+
+        /// <summary>
+        /// Test the shutdown handlers.
+        /// </summary>
+        [Test]
+        public void ShutdownHandlers()
+        {
+            // Test we can run one shutdown handler.
+            shutdown_count = 0;
+            PythonEngine.Initialize();
+            PythonEngine.AddShutdownHandler(OnShutdownIncrement);
+            PythonEngine.Shutdown();
+            Assert.That(shutdown_count, Is.EqualTo(1));
+
+            // Test we can run multiple shutdown handlers in the right order.
+            shutdown_count = 4;
+            PythonEngine.Initialize();
+            PythonEngine.AddShutdownHandler(OnShutdownIncrement);
+            PythonEngine.AddShutdownHandler(OnShutdownDouble);
+            PythonEngine.Shutdown();
+            // Correct: 4 * 2 + 1 = 9
+            // Wrong:  (4 + 1) * 2 = 10
+            Assert.That(shutdown_count, Is.EqualTo(9));
+
+            // Test we can remove shutdown handlers, handling duplicates.
+            shutdown_count = 4;
+            PythonEngine.Initialize();
+            PythonEngine.AddShutdownHandler(OnShutdownIncrement);
+            PythonEngine.AddShutdownHandler(OnShutdownIncrement);
+            PythonEngine.AddShutdownHandler(OnShutdownDouble);
+            PythonEngine.AddShutdownHandler(OnShutdownIncrement);
+            PythonEngine.AddShutdownHandler(OnShutdownDouble);
+            PythonEngine.RemoveShutdownHandler(OnShutdownDouble);
+            PythonEngine.Shutdown();
+            // Correct: (4 + 1) * 2 + 1 + 1 = 12
+            // Wrong:   (4 * 2) + 1 + 1 + 1 = 11
+            Assert.That(shutdown_count, Is.EqualTo(12));
+        }
     }
 }
diff --git a/src/runtime/pythonengine.cs b/src/runtime/pythonengine.cs
@@ -168,6 +168,16 @@ public static void Initialize(IEnumerable<string> args, bool setSysArgv = true)
                 initialized = true;
                 Exceptions.Clear();
 
+                // Make sure we clean up properly on app domain unload.
+                AppDomain.CurrentDomain.DomainUnload += OnDomainUnload;
+
+                // Remember to shut down the runtime.
+                AddShutdownHandler(Runtime.Shutdown);
+
+                // The global scope gets used implicitly quite early on, remember
+                // to clear it out when we shut down.
+                AddShutdownHandler(PyScopeManager.Global.Clear);
+
                 if (setSysArgv)
                 {
                     Py.SetArgv(args);
@@ -220,9 +230,6 @@ public static void Initialize(IEnumerable<string> args, bool setSysArgv = true)
                 {
                     locals.Dispose();
                 }
-
-                // Make sure we clean up properly on app domain unload.
-                AppDomain.CurrentDomain.DomainUnload += OnDomainUnload;
             }
         }
 
@@ -302,21 +309,86 @@ public static void Shutdown()
         {
             if (initialized)
             {
-                PyScopeManager.Global.Clear();
+                // If the shutdown handlers trigger a domain unload,
+                // don't call shutdown again.
+                AppDomain.CurrentDomain.DomainUnload -= OnDomainUnload;
+
+                ExecuteShutdownHandlers();
+
                 Marshal.FreeHGlobal(_pythonHome);
                 _pythonHome = IntPtr.Zero;
                 Marshal.FreeHGlobal(_programName);
                 _programName = IntPtr.Zero;
                 Marshal.FreeHGlobal(_pythonPath);
                 _pythonPath = IntPtr.Zero;
 
-                Runtime.Shutdown();
-
-                AppDomain.CurrentDomain.DomainUnload -= OnDomainUnload;
                 initialized = false;
             }
         }
 
+        /// <summary>
+        /// Called when the engine is shut down.
+        ///
+        /// Shutdown handlers are run in reverse order they were added, so that
+        /// resources available when running a shutdown handler are the same as
+        /// what was available when it was added.
+        /// </summary>
+        public delegate void ShutdownHandler();
+
+        static List<ShutdownHandler> ShutdownHandlers = new List<ShutdownHandler>();
+
+        /// <summary>
+        /// Add a function to be called when the engine is shut down.
+        ///
+        /// Shutdown handlers are executed in the opposite order they were
+        /// added, so that you can be sure that everything that was initialized
+        /// when you added the handler is still initialized when you need to shut
+        /// down.
+        ///
+        /// If the same shutdown handler is added several times, it will be run
+        /// several times.
+        ///
+        /// Don't add shutdown handlers while running a shutdown handler.
+        /// </summary>
+        public static void AddShutdownHandler(ShutdownHandler handler)
+        {
+            ShutdownHandlers.Add(handler);
+        }
+
+        /// <summary>
+        /// Remove a shutdown handler.
+        ///
+        /// If the same shutdown handler is added several times, only the last
+        /// one is removed.
+        ///
+        /// Don't remove shutdown handlers while running a shutdown handler.
+        /// </summary>
+        public static void RemoveShutdownHandler(ShutdownHandler handler)
+        {
+            for (int index = ShutdownHandlers.Count - 1; index >= 0; --index)
+            {
+                if (ShutdownHandlers[index] == handler)
+                {
+                    ShutdownHandlers.RemoveAt(index);
+                    break;
+                }
+            }
+        }
+
+        /// <summary>
+        /// Run all the shutdown handlers.
+        ///
+        /// They're run in opposite order they were added.
+        /// </summary>
+        static void ExecuteShutdownHandlers()
+        {
+            while(ShutdownHandlers.Count > 0)
+            {
+                var handler = ShutdownHandlers[ShutdownHandlers.Count - 1];
+                ShutdownHandlers.RemoveAt(ShutdownHandlers.Count - 1);
+                handler();
+            }
+        }
 
         /// <summary>
         /// AcquireLock Method
