Add ref count check for helping discover the bugs of decref too much

QuantConnect · jaredbroad · Mar 28, 2019 · Jun 23, 2018 · Jun 24, 2018 · Jun 24, 2018
commit 361022034989257bd85fca22cbca694d1900bd52
diff --git a/src/embed_tests/TestFinalizer.cs b/src/embed_tests/TestFinalizer.cs
@@ -8,7 +8,6 @@ namespace Python.EmbeddingTest
 {
     public class TestFinalizer
     {
-        private string _PYTHONMALLOC = string.Empty;
         private int _oldThreshold;
 
         [SetUp]
@@ -195,5 +194,46 @@ public void ErrorHandling()
                 Finalizer.Instance.ErrorHandler -= handleFunc;
             }
         }
+
+        [Test]
+        public void ValidateRefCount()
+        {
+            if (!Finalizer.Instance.RefCountValidationEnabled)
+            {
+                Assert.Pass("Only run with FINALIZER_CHECK");
+            }
+            IntPtr ptr = IntPtr.Zero;
+            bool called = false;
+            Finalizer.IncorrectRefCntHandler handler = (s, e) =>
+            {
+                called = true;
+                Assert.AreEqual(ptr, e.Handle);
+                Assert.AreEqual(2, e.ImpactedObjects.Count);
+                // Fix for this test, don't do this on general environment
+                Runtime.Runtime.XIncref(e.Handle);
+                return false;
+            };
+            Finalizer.Instance.IncorrectRefCntResovler += handler;
+            try
+            {
+                ptr = CreateStringGarbage();
+                FullGCCollect();
+                Assert.Throws<Finalizer.IncorrectRefCountException>(() => Finalizer.Instance.Collect());
+                Assert.IsTrue(called);
+            }
+            finally
+            {
+                Finalizer.Instance.IncorrectRefCntResovler -= handler;
+            }
+        }
+
+        private static IntPtr CreateStringGarbage()
+        {
+            PyString s1 = new PyString("test_string");
+            // s2 steal a reference from s1
+            PyString s2 = new PyString(s1.Handle);
+            return s1.Handle;
+        }
+
     }
 }
diff --git a/src/runtime/Python.Runtime.15.csproj b/src/runtime/Python.Runtime.15.csproj
@@ -68,10 +68,10 @@
     <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON3;$(Python3Version);$(PythonMonoDefineConstants)</DefineConstants>
   </PropertyGroup>
   <PropertyGroup Condition=" '$(Configuration)' == 'DebugMono'">
-    <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON2;$(Python2Version);$(PythonMonoDefineConstants);TRACE;DEBUG</DefineConstants>
+    <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON2;$(Python2Version);$(PythonMonoDefineConstants);FINALIZER_CHECK;TRACE;DEBUG</DefineConstants>
   </PropertyGroup>
   <PropertyGroup Condition=" '$(Configuration)' == 'DebugMonoPY3'">
-    <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON3;$(Python3Version);$(PythonMonoDefineConstants);TRACE;DEBUG</DefineConstants>
+    <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON3;$(Python3Version);$(PythonMonoDefineConstants);FINALIZER_CHECK;TRACE;DEBUG</DefineConstants>
   </PropertyGroup>
   <PropertyGroup Condition=" '$(Configuration)' == 'ReleaseWin'">
     <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON2;$(Python2Version);$(PythonWinDefineConstants)</DefineConstants>
@@ -80,10 +80,10 @@
     <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON3;$(Python3Version);$(PythonWinDefineConstants)</DefineConstants>
   </PropertyGroup>
   <PropertyGroup Condition=" '$(Configuration)' == 'DebugWin'">
-    <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON2;$(Python2Version);$(PythonWinDefineConstants);TRACE;DEBUG</DefineConstants>
+    <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON2;$(Python2Version);$(PythonWinDefineConstants);FINALIZER_CHECK;TRACE;DEBUG</DefineConstants>
   </PropertyGroup>
   <PropertyGroup Condition=" '$(Configuration)' == 'DebugWinPY3'">
-    <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON3;$(Python3Version);$(PythonWinDefineConstants);TRACE;DEBUG</DefineConstants>
+    <DefineConstants Condition="'$(CustomDefineConstants)' == ''">$(DefineConstants);PYTHON3;$(Python3Version);$(PythonWinDefineConstants);FINALIZER_CHECK;TRACE;DEBUG</DefineConstants>
   </PropertyGroup>
 
   <ItemGroup Condition=" '$(PythonInteropFile)' != '' ">

diff --git a/src/runtime/delegatemanager.cs b/src/runtime/delegatemanager.cs
@@ -181,7 +181,7 @@ 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 : IDisposable
+    public class Dispatcher : IPyDisposable
     {
         public IntPtr target;
         public Type dtype;
@@ -276,6 +276,11 @@ public object TrueDispatch(ArrayList args)
             Runtime.XDecref(op);
             return result;
         }
+
+        public IntPtr[] GetTrackedHandles()
+        {
+            return new IntPtr[] { target };
+        }
     }
 
 

diff --git a/src/runtime/finalizer.cs b/src/runtime/finalizer.cs
@@ -38,11 +38,48 @@ struct PendingArgs
         private delegate int PendingCall(IntPtr arg);
         private readonly PendingCall _collectAction;
 
-        private ConcurrentQueue<IDisposable> _objQueue = new ConcurrentQueue<IDisposable>();
+        private ConcurrentQueue<IPyDisposable> _objQueue = new ConcurrentQueue<IPyDisposable>();
         private bool _pending = false;
         private readonly object _collectingLock = new object();
         private IntPtr _pendingArgs;
 
+        #region FINALIZER_CHECK
+
+#if FINALIZER_CHECK
+        private readonly object _queueLock = new object();
+        public bool RefCountValidationEnabled { get; set; } = true;
+#else
+        public readonly bool RefCountValidationEnabled = false;
+#endif
+        // Keep these declarations for compat even no FINALIZER_CHECK
+        public class IncorrectFinalizeArgs : EventArgs
+        {
+            public IntPtr Handle { get; internal set; }
+            public ICollection<IPyDisposable> ImpactedObjects { get; internal set; }
+        }
+
+        public class IncorrectRefCountException : Exception
+        {
+            public IntPtr PyPtr { get; internal set; }
+            private string _message;
+            public override string Message => _message;
+
+            public IncorrectRefCountException(IntPtr ptr)
+            {
+                PyPtr = ptr;
+                IntPtr pyname = Runtime.PyObject_Unicode(PyPtr);
+                string name = Runtime.GetManagedString(pyname);
+                Runtime.XDecref(pyname);
+                _message = $"{name} may has a incorrect ref count";
+            }
+        }
+
+        public delegate bool IncorrectRefCntHandler(object sender, IncorrectFinalizeArgs e);
+        public event IncorrectRefCntHandler IncorrectRefCntResovler;
+        public bool ThrowIfUnhandleIncorrectRefCount { get; set; } = true;
+
+        #endregion
+
         private Finalizer()
         {
             Enable = true;
@@ -72,7 +109,7 @@ public List<WeakReference> GetCollectedObjects()
             return _objQueue.Select(T => new WeakReference(T)).ToList();
         }
 
-        internal void AddFinalizedObject(IDisposable obj)
+        internal void AddFinalizedObject(IPyDisposable obj)
         {
             if (!Enable)
             {
@@ -84,7 +121,12 @@ internal void AddFinalizedObject(IDisposable obj)
                 // for avoiding that case, user should call GC.Collect manual before shutdown.
                 return;
             }
-            _objQueue.Enqueue(obj);
+#if FINALIZER_CHECK
+            lock (_queueLock)
+#endif
+            {
+                _objQueue.Enqueue(obj);
+            }
             GC.ReRegisterForFinalize(obj);
             if (_objQueue.Count >= Threshold)
             {
@@ -96,7 +138,7 @@ internal static void Shutdown()
         {
             if (Runtime.Py_IsInitialized() == 0)
             {
-                Instance._objQueue = new ConcurrentQueue<IDisposable>();
+                Instance._objQueue = new ConcurrentQueue<IPyDisposable>();
                 return;
             }
             Instance.DisposeAll();
@@ -175,21 +217,29 @@ private void DisposeAll()
             {
                 ObjectCount = _objQueue.Count
             });
-            IDisposable obj;
-            while (_objQueue.TryDequeue(out obj))
+#if FINALIZER_CHECK
+            lock (_queueLock)
+#endif
             {
-                try
-                {
-                    obj.Dispose();
-                    Runtime.CheckExceptionOccurred();
-                }
-                catch (Exception e)
+#if FINALIZER_CHECK
+                ValidateRefCount();
+#endif
+                IPyDisposable obj;
+                while (_objQueue.TryDequeue(out obj))
                 {
-                    // We should not bother the main thread
-                    ErrorHandler?.Invoke(this, new ErrorArgs()
+                    try
+                    {
+                        obj.Dispose();
+                        Runtime.CheckExceptionOccurred();
+                    }
+                    catch (Exception e)
                     {
-                        Error = e
-                    });
+                        // We should not bother the main thread
+                        ErrorHandler?.Invoke(this, new ErrorArgs()
+                        {
+                            Error = e
+                        });
+                    }
                 }
             }
         }
@@ -202,5 +252,80 @@ private void ResetPending()
                 _pendingArgs = IntPtr.Zero;
             }
         }
+
+#if FINALIZER_CHECK
+        private void ValidateRefCount()
+        {
+            if (!RefCountValidationEnabled)
+            {
+                return;
+            }
+            var counter = new Dictionary<IntPtr, long>();
+            var holdRefs = new Dictionary<IntPtr, long>();
+            var indexer = new Dictionary<IntPtr, List<IPyDisposable>>();
+            foreach (var obj in _objQueue)
+            {
+                IntPtr[] handles = obj.GetTrackedHandles();
+                foreach (var handle in handles)
+                {
+                    if (handle == IntPtr.Zero)
+                    {
+                        continue;
+                    }
+                    if (!counter.ContainsKey(handle))
+                    {
+                        counter[handle] = 0;
+                    }
+                    counter[handle]++;
+                    if (!holdRefs.ContainsKey(handle))
+                    {
+                        holdRefs[handle] = Runtime.Refcount(handle);
+                    }
+                    List<IPyDisposable> objs;
+                    if (!indexer.TryGetValue(handle, out objs))
+                    {
+                        objs = new List<IPyDisposable>();
+                        indexer.Add(handle, objs);
+                    }
+                    objs.Add(obj);
+                }
+            }
+            foreach (var pair in counter)
+            {
+                IntPtr handle = pair.Key;
+                long cnt = pair.Value;
+                // Tracked handle's ref count is larger than the object's holds
+                // it may take an unspecified behaviour if it decref in Dispose
+                if (cnt > holdRefs[handle])
+                {
+                    var args = new IncorrectFinalizeArgs()
+                    {
+                        Handle = handle,
+                        ImpactedObjects = indexer[handle]
+                    };
+                    bool handled = false;
+                    if (IncorrectRefCntResovler != null)
+                    {
+                        var funcList = IncorrectRefCntResovler.GetInvocationList();
+                        foreach (IncorrectRefCntHandler func in funcList)
+                        {
+                            if (func(this, args))
+                            {
+                                handled = true;
+                                break;
+                            }
+                        }
+                    }
+                    if (!handled && ThrowIfUnhandleIncorrectRefCount)
+                    {
+                        throw new IncorrectRefCountException(handle);
+                    }
+                }
+                // Make sure no other references for PyObjects after this method
+                indexer[handle].Clear();
+            }
+            indexer.Clear();
+        }
+#endif
     }
 }
diff --git a/src/runtime/pyobject.cs b/src/runtime/pyobject.cs
@@ -1,20 +1,33 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Dynamic;
 using System.Linq.Expressions;
 
 namespace Python.Runtime
 {
+    public interface IPyDisposable : IDisposable
+    {
+        IntPtr[] GetTrackedHandles();
+    }
+
     /// <summary>
     /// Represents a generic Python object. The methods of this class are
     /// generally equivalent to the Python "abstract object API". See
     /// PY2: https://docs.python.org/2/c-api/object.html
     /// PY3: https://docs.python.org/3/c-api/object.html
     /// for details.
     /// </summary>
-    public class PyObject : DynamicObject, IEnumerable, IDisposable
+    public class PyObject : DynamicObject, IEnumerable, IPyDisposable
     {
+#if TRACE_ALLOC
+        /// <summary>
+        /// Trace stack for PyObject's construction
+        /// </summary>
+        public StackTrace Traceback { get; private set; }
+#endif
+
         protected internal IntPtr obj = IntPtr.Zero;
         private bool disposed = false;
         private bool _finalized = false;
@@ -31,19 +44,29 @@ public class PyObject : DynamicObject, IEnumerable, IDisposable
         public PyObject(IntPtr ptr)
         {
             obj = ptr;
+#if TRACE_ALLOC
+            Traceback = new StackTrace(1);
+#endif
         }
 
         // Protected default constructor to allow subclasses to manage
         // initialization in different ways as appropriate.
 
         protected PyObject()
         {
+#if TRACE_ALLOC
+            Traceback = new StackTrace(1);
+#endif
         }
 
         // Ensure that encapsulated Python object is decref'ed appropriately
         // when the managed wrapper is garbage-collected.
         ~PyObject()
         {
+            if (obj == IntPtr.Zero)
+            {
+                return;
+            }
             if (_finalized || disposed)
             {
                 return;
@@ -176,6 +199,11 @@ public void UnsafeDispose()
             GC.SuppressFinalize(this);
         }
 
+        public IntPtr[] GetTrackedHandles()
+        {
+            return new IntPtr[] { obj };
+        }
+
         /// <summary>
         /// GetPythonType Method
         /// </summary>