Finalizer for PyObject

QuantConnect · jaredbroad · Mar 28, 2019 · Jun 23, 2018 · Jun 24, 2018 · Jun 24, 2018
commit 26e88d748c80c9bc585d3acc1077fe5e0530ea6f
diff --git a/src/embed_tests/TestFinalizer.cs b/src/embed_tests/TestFinalizer.cs
@@ -0,0 +1,143 @@
+using NUnit.Framework;
+using Python.Runtime;
+using System;
+using System.Threading;
+
+namespace Python.EmbeddingTest
+{
+    public class TestFinalizer
+    {
+        private string _PYTHONMALLOC = string.Empty;
+
+        [SetUp]
+        public void SetUp()
+        {
+            try
+            {
+                _PYTHONMALLOC = Environment.GetEnvironmentVariable("PYTHONMALLOC");
+            }
+            catch (ArgumentNullException)
+            {
+                _PYTHONMALLOC = string.Empty;
+            }
+            Environment.SetEnvironmentVariable("PYTHONMALLOC", "malloc");
+            PythonEngine.Initialize();
+        }
+
+        [TearDown]
+        public void TearDown()
+        {
+            PythonEngine.Shutdown();
+            if (string.IsNullOrEmpty(_PYTHONMALLOC))
+            {
+                Environment.SetEnvironmentVariable("PYTHONMALLOC", _PYTHONMALLOC);
+            }
+        }
+
+        private static void FullGCCollect()
+        {
+            GC.Collect(GC.MaxGeneration, GCCollectionMode.Forced);
+            GC.WaitForFullGCComplete();
+            GC.WaitForPendingFinalizers();
+        }
+
+        [Test]
+        public void CollectBasicObject()
+        {
+            int thId = Thread.CurrentThread.ManagedThreadId;
+            Finalizer.Instance.Threshold = 1;
+            bool called = false;
+            EventHandler<Finalizer.CollectArgs> handler = (s, e) =>
+            {
+                Assert.AreEqual(thId, Thread.CurrentThread.ManagedThreadId);
+                Assert.GreaterOrEqual(e.ObjectCount, 1);
+                called = true;
+            };
+            Finalizer.Instance.CollectOnce += handler;
+            FullGCCollect();
+            PyLong obj = new PyLong(1024);
+
+            WeakReference shortWeak = new WeakReference(obj);
+            WeakReference longWeak = new WeakReference(obj, true);
+            obj = null;
+            FullGCCollect();
+            // The object has been resurrected
+            Assert.IsFalse(shortWeak.IsAlive);
+            Assert.IsTrue(longWeak.IsAlive);
+
+            Assert.IsFalse(called);
+            var garbage = Finalizer.Instance.GetCollectedObjects();
+            // FIXME: If make some query for garbage,
+            // the above case will failed Assert.IsFalse(shortWeak.IsAlive)
+            //Assert.IsTrue(garbage.All(T => T.IsAlive));
+
+            Finalizer.Instance.CallPendingFinalizers();
+            Assert.IsTrue(called);
+
+            FullGCCollect();
+            //Assert.IsFalse(garbage.All(T => T.IsAlive));
+
+            Assert.IsNull(longWeak.Target);
+
+            Finalizer.Instance.CollectOnce -= handler;
+        }
+
+        private static long CompareWithFinalizerOn(PyObject pyCollect, bool enbale)
+        {
+            // Must larger than 512 bytes make sure Python use 
+            string str = new string('1', 1024);
+            Finalizer.Instance.Enable = true;
+            FullGCCollect();
+            FullGCCollect();
+            pyCollect.Invoke();
+            Finalizer.Instance.CallPendingFinalizers();
+            Finalizer.Instance.Enable = enbale;
+
+            // Estimate unmanaged memory size
+            long before = Environment.WorkingSet - GC.GetTotalMemory(true);
+            for (int i = 0; i < 10000; i++)
+            {
+                // Memory will leak when disable Finalizer
+                new PyString(str);
+            }
+            FullGCCollect();
+            FullGCCollect();
+            pyCollect.Invoke();
+            if (enbale)
+            {
+                Finalizer.Instance.CallPendingFinalizers();
+            }
+
+            FullGCCollect();
+            FullGCCollect();
+            long after = Environment.WorkingSet - GC.GetTotalMemory(true);
+            return after - before;
+
+        }
+
+        /// <summary>
+        /// Because of two vms both have their memory manager,
+        /// this test only prove the finalizer has take effect.
+        /// </summary>
+        [Test]
+        [Ignore("Too many uncertainties, only manual on when debugging")]
+        public void SimpleTestMemory()
+        {
+            bool oldState = Finalizer.Instance.Enable;
+            try
+            {
+                using (PyObject gcModule = PythonEngine.ImportModule("gc"))
+                using (PyObject pyCollect = gcModule.GetAttr("collect"))
+                {
+                    long span1 = CompareWithFinalizerOn(pyCollect, false);
+                    long span2 = CompareWithFinalizerOn(pyCollect, true);
+                    Assert.Less(span2, span1);
+                }
+            }
+            finally
+            {
+                Finalizer.Instance.Enable = oldState;
+            }
+        }
+    }
+}
diff --git a/src/runtime/delegatemanager.cs b/src/runtime/delegatemanager.cs
@@ -181,10 +181,12 @@ A possible alternate strategy would be to create custom subclasses
        too "special" for this to work. It would be more work, so for now
        the 80/20 rule applies :) */
 
-    public class Dispatcher
+    public class Dispatcher : IDisposable
     {
         public IntPtr target;
         public Type dtype;
+        private bool _disposed = false;
+        private bool _finalized = false;
 
         public Dispatcher(IntPtr target, Type dtype)
         {
@@ -195,18 +197,25 @@ public Dispatcher(IntPtr target, Type dtype)
 
         ~Dispatcher()
         {
-            // We needs to disable Finalizers until it's valid implementation.
-            // Current implementation can produce low probability floating bugs.
-            return;
+            if (_finalized || _disposed)
+            {
+                return;
+            }
+            _finalized = true;
+            Finalizer.Instance.AddFinalizedObject(this);
+        }
 
-            // Note: the managed GC thread can run and try to free one of
-            // these *after* the Python runtime has been finalized!
-            if (Runtime.Py_IsInitialized() > 0)
+        public void Dispose()
+        {
+            if (_disposed)
             {
-                IntPtr gs = PythonEngine.AcquireLock();
-                Runtime.XDecref(target);
-                PythonEngine.ReleaseLock(gs);
+                return;
             }
+            _disposed = true;
+            Runtime.XDecref(target);
+            target = IntPtr.Zero;
+            dtype = null;
+            GC.SuppressFinalize(this);
         }
 
         public object Dispatch(ArrayList args)

diff --git a/src/runtime/finalizer.cs b/src/runtime/finalizer.cs
@@ -0,0 +1,117 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
+using System.Runtime.InteropServices;
+using System.Threading;
+
+namespace Python.Runtime
+{
+    public class Finalizer
+    {
+        public class CollectArgs : EventArgs
+        {
+            public int ObjectCount { get; set; }
+        }
+
+        public static readonly Finalizer Instance = new Finalizer();
+
+        public event EventHandler<CollectArgs> CollectOnce;
+
+        private ConcurrentQueue<IDisposable> _objQueue = new ConcurrentQueue<IDisposable>();
+
+        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
+        private delegate int PedingCall(IntPtr arg);
+        private PedingCall _collectAction;
+        private bool _pending = false;
+        private readonly object _collectingLock = new object();
+        public int Threshold { get; set; }
+        public bool Enable { get; set; }
+
+        private Finalizer()
+        {
+            Enable = true;
+            Threshold = 200;
+            _collectAction = OnCollect;
+        }
+
+        public void CallPendingFinalizers()
+        {
+            if (Thread.CurrentThread.ManagedThreadId != Runtime.MainManagedThreadId)
+            {
+                throw new Exception("PendingCall should execute in main Python thread");
+            }
+            Runtime.Py_MakePendingCalls();
+        }
+
+        public List<WeakReference> GetCollectedObjects()
+        {
+            return _objQueue.Select(T => new WeakReference(T)).ToList();
+        }
+
+        internal void AddFinalizedObject(IDisposable obj)
+        {
+            if (!Enable)
+            {
+                return;
+            }
+            if (Runtime.Py_IsInitialized() == 0)
+            {
+                // XXX: Memory will leak if a PyObject finalized after Python shutdown,
+                // for avoiding that case, user should call GC.Collect manual before shutdown.
+                return;
+            }
+            _objQueue.Enqueue(obj);
+            GC.ReRegisterForFinalize(obj);
+            if (_objQueue.Count >= Threshold)
+            {
+                Collect();
+            }
+        }
+
+        internal static void Shutdown()
+        {
+            Instance.DisposeAll();
+            Instance.CallPendingFinalizers();
+            Runtime.PyErr_Clear();
+        }
+
+        private void Collect()
+        {
+            lock (_collectingLock)
+            {
+                if (_pending)
+                {
+                    return;
+                }
+                _pending = true;
+            }
+            IntPtr func = Marshal.GetFunctionPointerForDelegate(_collectAction);
+            if (Runtime.Py_AddPendingCall(func, IntPtr.Zero) != 0)
+            {
+                // Full queue, append next time
+                _pending = false;
+            }
+        }
+
+        private int OnCollect(IntPtr arg)
+        {
+            CollectOnce?.Invoke(this, new CollectArgs()
+            {
+                ObjectCount = _objQueue.Count
+            });
+            DisposeAll();
+            _pending = false;
+            return 0;
+        }
+
+        private void DisposeAll()
+        {
+            IDisposable obj;
+            while (_objQueue.TryDequeue(out obj))
+            {
+                obj.Dispose();
+            }
+        }
+    }
+}
diff --git a/src/runtime/pyobject.cs b/src/runtime/pyobject.cs
@@ -17,6 +17,7 @@ public class PyObject : DynamicObject, IEnumerable, IDisposable
     {
         protected internal IntPtr obj = IntPtr.Zero;
         private bool disposed = false;
+        private bool _finalized = false;
 
         /// <summary>
         /// PyObject Constructor
@@ -41,14 +42,15 @@ protected PyObject()
 
         // Ensure that encapsulated Python object is decref'ed appropriately
         // when the managed wrapper is garbage-collected.
-
         ~PyObject()
         {
-            // We needs to disable Finalizers until it's valid implementation.
-            // Current implementation can produce low probability floating bugs.
-            return;
-
-            Dispose();
+            if (_finalized || disposed)
+            {
+                return;
+            }
+            // Prevent a infinity loop by calling GC.WaitForPendingFinalizers
+            _finalized = true;
+            Finalizer.Instance.AddFinalizedObject(this);
         }
 
 

diff --git a/src/runtime/pyscope.cs b/src/runtime/pyscope.cs
@@ -37,6 +37,7 @@ public class PyScope : DynamicObject, IDisposable
         internal readonly IntPtr variables;
 
         private bool _isDisposed;
+        private bool _finalized = false;
 
         /// <summary>
         /// The Manager this scope associated with.
@@ -527,11 +528,12 @@ public void Dispose()
 
         ~PyScope()
         {
-            // We needs to disable Finalizers until it's valid implementation.
-            // Current implementation can produce low probability floating bugs.
-            return;
-
-            Dispose();
+            if (_finalized || _isDisposed)
+            {
+                return;
+            }
+            _finalized = true;
+            Finalizer.Instance.AddFinalizedObject(this);
         }
     }