fix ref bug with iter, views and istr (#1311)

Co-authored-by: J. Nick Koston <[email protected]>
Co-authored-by: J. Nick Koston <[email protected]>

Author: Vizonex

CHANGES/1311.bugfix.rst

@@ -0,0 +1,2 @@
+Fixed reference leak in iterators, views and ``istr``
+-- by :user:`Vizonex`.

multidict/_multilib/istr.h

@@ -22,8 +22,10 @@ PyDoc_STRVAR(istr__doc__, "istr class implementation");
 static inline void
 istr_dealloc(istrobject *self)
 {
+    PyTypeObject *tp = Py_TYPE(self);
     Py_XDECREF(self->canonical);
     PyUnicode_Type.tp_dealloc((PyObject *)self);
+    Py_DECREF(tp);
 }
 
 static inline PyObject *

multidict/_multilib/iter.h

@@ -134,14 +134,17 @@ multidict_keys_iter_iternext(MultidictIter *self)
 static inline void
 multidict_iter_dealloc(MultidictIter *self)
 {
+    PyTypeObject *tp = Py_TYPE(self);
     PyObject_GC_UnTrack(self);
     Py_XDECREF(self->md);
-    PyObject_GC_Del(self);
+    tp->tp_free(self);
+    Py_DECREF(tp);
 }
 
 static inline int
 multidict_iter_traverse(MultidictIter *self, visitproc visit, void *arg)
 {
+    Py_VISIT(Py_TYPE(self));
     Py_VISIT(self->md);
     return 0;
 }

multidict/_multilib/views.h

@@ -30,15 +30,18 @@ _init_view(_Multidict_ViewObject *self, MultiDictObject *md)
 static inline void
 multidict_view_dealloc(_Multidict_ViewObject *self)
 {
+    PyTypeObject *tp = Py_TYPE(self);
     PyObject_GC_UnTrack(self);
     Py_XDECREF(self->md);
-    PyObject_GC_Del(self);
+    tp->tp_free(self);
+    Py_DECREF(tp);
 }
 
 static inline int
 multidict_view_traverse(_Multidict_ViewObject *self, visitproc visit,
                         void *arg)
 {
+    Py_VISIT(Py_TYPE(self));
     Py_VISIT(self->md);
     return 0;
 }

tests/isolated/multidict_type_leak.py

@@ -0,0 +1,50 @@
+import gc
+import sys
+import sysconfig
+
+from multidict import MultiDict, istr
+
+# sys.getrefcount is not meaningful under the free-threaded build:
+# refcounts are biased per-thread and types may be immortalized, so
+# the simple baseline/after comparison below does not apply.
+FREETHREADED = bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
+
+
+if __name__ == "__main__":
+    if FREETHREADED:
+        raise SystemExit(0)
+
+    md = MultiDict([("a", "1"), ("b", "2")])
+
+    # Iterator type leak
+    iter_type = type(iter(md.keys()))
+    gc.collect()
+    baseline = sys.getrefcount(iter_type)
+    for _ in range(1000):
+        _it = iter(md.keys())
+        list(_it)
+        del _it
+    gc.collect()
+    after = sys.getrefcount(iter_type)
+    assert after == baseline, f"iterator type leaked: {after - baseline}"
+
+    # View type leak
+    view_type = type(md.keys())
+    gc.collect()
+    baseline = sys.getrefcount(view_type)
+    for _ in range(1000):
+        _v = md.keys()
+        del _v
+    gc.collect()
+    after = sys.getrefcount(view_type)
+    assert after == baseline, f"view type leaked: {after - baseline}"
+
+    # istr type leak
+    gc.collect()
+    baseline = sys.getrefcount(istr)
+    for _ in range(1000):
+        _s = istr("hello")
+        del _s
+    gc.collect()
+    after = sys.getrefcount(istr)
+    assert after == baseline, f"istr type leaked: {after - baseline}"

tests/isolated/multidict_type_leak_items_values.py

@@ -0,0 +1,49 @@
+import gc
+import sys
+import sysconfig
+
+from multidict import MultiDict
+
+# sys.getrefcount is not meaningful under the free-threaded build:
+# refcounts are biased per-thread and types may be immortalized, so
+# the simple baseline/after comparison below does not apply.
+FREETHREADED = bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
+
+
+if __name__ == "__main__":
+    if FREETHREADED:
+        raise SystemExit(0)
+
+    md = MultiDict([("a", "1"), ("b", "2")])
+
+    # items() / values() iterators each have their own PyType_Spec
+    # sharing multidict_iter_dealloc; test them independently so a
+    # regression in just one spec's slot table is still caught.
+    for view_name in ("items", "values"):
+        get_view = getattr(md, view_name)
+        iter_type = type(iter(get_view()))
+        gc.collect()
+        baseline = sys.getrefcount(iter_type)
+        for _ in range(1000):
+            _it = iter(get_view())
+            list(_it)
+            del _it
+        gc.collect()
+        after = sys.getrefcount(iter_type)
+        assert after == baseline, (
+            f"{view_name} iterator type leaked: {after - baseline}"
+        )
+
+    # items() / values() views each have their own PyType_Spec
+    # sharing multidict_view_dealloc; same rationale as above.
+    for view_name in ("items", "values"):
+        get_view = getattr(md, view_name)
+        view_type = type(get_view())
+        gc.collect()
+        baseline = sys.getrefcount(view_type)
+        for _ in range(1000):
+            _v = get_view()
+            del _v
+        gc.collect()
+        after = sys.getrefcount(view_type)
+        assert after == baseline, f"{view_name} view type leaked: {after - baseline}"

tests/test_leaks.py

@@ -15,6 +15,8 @@
         "multidict_extend_multidict.py",
         "multidict_extend_tuple.py",
         "multidict_update_multidict.py",
+        "multidict_type_leak.py",
+        "multidict_type_leak_items_values.py",
         "multidict_pop.py",
     ),
 )