Fix memory leak on objects passed to init, update, or extend (#1121)

fixes #1117

---------

Co-authored-by: Andrew Svetlov <andrew.svetlov@gmail.com>

Author: bdraco

CHANGES/1121.bugfix.rst

@@ -0,0 +1 @@
+Resolved a memory leak by ensuring proper reference count decrementation -- by :user:`asvetlov` and :user:`bdraco`.

docs/spelling_wordlist.txt

@@ -26,6 +26,7 @@ Ctrl
 cython
 Deprecations
 deallocation
+decrementation
 dev
 dict
 docstrings

multidict/_multidict.c

@@ -505,15 +505,19 @@ multidict_tp_init(MultiDictObject *self, PyObject *args, PyObject *kwds)
     PyObject *arg = NULL;
     Py_ssize_t size = _multidict_extend_parse_args(args, kwds, "MultiDict", &arg);
     if (size < 0) {
-        return -1;
+        goto fail;
     }
     if (pair_list_init(&self->pairs, size) < 0) {
-        return -1;
+        goto fail;
     }
     if (_multidict_extend(self, arg, kwds, "MultiDict", 1) < 0) {
-        return -1;
+        goto fail;
     }
+    Py_CLEAR(arg);
     return 0;
+fail:
+    Py_CLEAR(arg);
+    return -1;
 }
 
 static inline PyObject *
@@ -544,14 +548,17 @@ multidict_extend(MultiDictObject *self, PyObject *args, PyObject *kwds)
     PyObject *arg = NULL;
     Py_ssize_t size = _multidict_extend_parse_args(args, kwds, "extend", &arg);
     if (size < 0) {
-        return NULL;
+        goto fail;
     }
     pair_list_grow(&self->pairs, size);
     if (_multidict_extend(self, arg, kwds, "extend", 1) < 0) {
-        return NULL;
+        goto fail;
     }
-
+    Py_CLEAR(arg);
     Py_RETURN_NONE;
+fail:
+    Py_CLEAR(arg);
+    return NULL;
 }
 
 static inline PyObject *
@@ -692,12 +699,16 @@ multidict_update(MultiDictObject *self, PyObject *args, PyObject *kwds)
 {
     PyObject *arg = NULL;
     if (_multidict_extend_parse_args(args, kwds, "update", &arg) < 0) {
-        return NULL;
+        goto fail;
     }
     if (_multidict_extend(self, arg, kwds, "update", 0) < 0) {
-        return NULL;
+        goto fail;
     }
+    Py_CLEAR(arg);
     Py_RETURN_NONE;
+fail:
+    Py_CLEAR(arg);
+    return NULL;
 }
 
 PyDoc_STRVAR(multidict_add_doc,
@@ -917,17 +928,21 @@ cimultidict_tp_init(MultiDictObject *self, PyObject *args, PyObject *kwds)
     PyObject *arg = NULL;
     Py_ssize_t size = _multidict_extend_parse_args(args, kwds, "CIMultiDict", &arg);
     if (size < 0) {
-        return -1;
+        goto fail;
     }
 
     if (ci_pair_list_init(&self->pairs, size) < 0) {
-        return -1;
+        goto fail;
     }
 
     if (_multidict_extend(self, arg, kwds, "CIMultiDict", 1) < 0) {
-        return -1;
+        goto fail;
     }
+    Py_CLEAR(arg);
     return 0;
+fail:
+    Py_CLEAR(arg);
+    return -1;
 }
 
 

tests/test_leaks.py

@@ -19,7 +19,6 @@
 )
 @pytest.mark.leaks
 @pytest.mark.skipif(IS_PYPY, reason="leak testing is not supported on PyPy")
-@pytest.mark.xfail(reason="memory leak https://github.com/aio-libs/multidict/issues/1117")
 def test_leak(script: str) -> None:
     """Run isolated leak test script and check for leaks."""
     leak_test_script = pathlib.Path(__file__).parent.joinpath("isolated", script)

tests/test_multidict.py

@@ -2,6 +2,7 @@
 
 import gc
 import operator
+import platform
 import sys
 import weakref
 from collections import deque
@@ -22,6 +23,7 @@
 )
 
 _T = TypeVar("_T")
+IS_PYPY = platform.python_implementation() == "PyPy"
 
 
 def chained_callable(
@@ -1272,3 +1274,38 @@ def test_convert_multidict_to_cimultidict_eq(
     ) == case_insensitive_multidict_class(
         [("H1", "header1"), ("H2", "header2"), ("H3", "header3")]
     )
+
+
+@pytest.mark.skipif(IS_PYPY, reason="getrefcount is not supported on PyPy")
+def test_extend_does_not_alter_refcount(
+    case_sensitive_multidict_class: type[MultiDict[str]],
+) -> None:
+    """Test that extending a MultiDict with a MultiDict does not alter the refcount of the original."""
+    original = case_sensitive_multidict_class([("h1", "header1")])
+    new = case_sensitive_multidict_class([("h2", "header2")])
+    original_refcount = sys.getrefcount(original)
+    new.extend(original)
+    assert sys.getrefcount(original) == original_refcount
+
+
+@pytest.mark.skipif(IS_PYPY, reason="getrefcount is not supported on PyPy")
+def test_update_does_not_alter_refcount(
+    case_sensitive_multidict_class: type[MultiDict[str]],
+) -> None:
+    """Test that updating a MultiDict with a MultiDict does not alter the refcount of the original."""
+    original = case_sensitive_multidict_class([("h1", "header1")])
+    new = case_sensitive_multidict_class([("h2", "header2")])
+    original_refcount = sys.getrefcount(original)
+    new.update(original)
+    assert sys.getrefcount(original) == original_refcount
+
+
+@pytest.mark.skipif(IS_PYPY, reason="getrefcount is not supported on PyPy")
+def test_init_does_not_alter_refcount(
+    case_sensitive_multidict_class: type[MultiDict[str]],
+) -> None:
+    """Test that initializing a MultiDict with a MultiDict does not alter the refcount of the original."""
+    original = case_sensitive_multidict_class([("h1", "header1")])
+    original_refcount = sys.getrefcount(original)
+    case_sensitive_multidict_class(original)
+    assert sys.getrefcount(original) == original_refcount