Run callbacks registered by atexit at Shutdown on soft-shutdown mode

amos402 · amos402 · commit 6d738bf53a1e · 2020-07-12T03:39:52.000+08:00
diff --git a/src/embed_tests/pyinitialize.cs b/src/embed_tests/pyinitialize.cs
@@ -136,5 +136,46 @@ public void ShutdownHandlers()
             // Wrong:   (4 * 2) + 1 + 1 + 1 = 11
             Assert.That(shutdown_count, Is.EqualTo(12));
         }
+
+        [Test]
+        public static void TestRunExitFuncs()
+        {
+            if (Runtime.Runtime.GetDefaultShutdownMode() == ShutdownMode.Normal)
+            {
+                // If the runtime using the normal mode,
+                // callback registered by atexit will be called after we release the clr information,
+                // thus there's no chance we can check it here.
+                Assert.Ignore("Skip on normal mode");
+            }
+            Runtime.Runtime.Initialize();
+            PyObject atexit;
+            try
+            {
+                atexit = Py.Import("atexit");
+            }
+            catch (PythonException e)
+            {
+                string msg = e.ToString();
+                Runtime.Runtime.Shutdown();
+
+                if (e.IsMatches(Exceptions.ImportError))
+                {
+                    Assert.Ignore("no atexit module");
+                }
+                else
+                {
+                    Assert.Fail(msg);
+                }
+                return;
+            }
+            bool called = false;
+            Action callback = () =>
+            {
+                called = true;
+            };
+            atexit.InvokeMethod("register", callback.ToPython());
+            Runtime.Runtime.Shutdown();
+            Assert.True(called);
+        }
     }
 }
diff --git a/src/runtime/pythonexception.cs b/src/runtime/pythonexception.cs
@@ -187,6 +187,11 @@ public string Format()
             return res;
         }
 
+        public bool IsMatches(IntPtr exc)
+        {
+            return Runtime.PyErr_GivenExceptionMatches(PyType, exc) != 0;
+        }
+
         /// <summary>
         /// Dispose Method
         /// </summary>
diff --git a/src/runtime/runtime.cs b/src/runtime/runtime.cs
@@ -317,6 +317,10 @@ internal static void Shutdown()
             PyGILState_Ensure();
 
             var mode = ShutdownMode;
+            if (mode != ShutdownMode.Normal)
+            {
+                RunExitFuncs();
+            }
 #if !NETSTANDARD
             if (mode == ShutdownMode.Reload)
             {
@@ -388,6 +392,34 @@ internal static int AtExit()
             return 0;
         }
 
+        private static void RunExitFuncs()
+        {
+            PyObject atexit;
+            try
+            {
+                atexit = Py.Import("atexit");
+            }
+            catch (PythonException e)
+            {
+                if (!e.IsMatches(Exceptions.ImportError))
+                {
+                    throw;
+                }
+                e.Dispose();
+                // The runtime may not provided `atexit` module.
+                return;
+            }
+            try
+            {
+                atexit.InvokeMethod("_run_exitfuncs").Dispose();
+            }
+            catch (PythonException e)
+            {
+                Console.Error.WriteLine(e);
+                e.Dispose();
+            }
+        }
+
         private static void SetPyMember(ref IntPtr obj, IntPtr value, Action onRelease)
         {
             // XXX: For current usages, value should not be null.
@@ -1849,7 +1881,7 @@ internal static IntPtr PyType_GenericAlloc(IntPtr type, long n)
         private static extern IntPtr PyType_GenericAlloc(IntPtr type, IntPtr n);
 
         /// <summary>
-        /// Finalize a type object. This should be called on all type objects to finish their initialization. This function is responsible for adding inherited slots from a type��s base class. Return 0 on success, or return -1 and sets an exception on error.
+        /// Finalize a type object. This should be called on all type objects to finish their initialization. This function is responsible for adding inherited slots from a type’s base class. Return 0 on success, or return -1 and sets an exception on error.
         /// </summary>
         [DllImport(_PythonDll, CallingConvention = CallingConvention.Cdecl)]
         internal static extern int PyType_Ready(IntPtr type);
