Issue #28427: old keys should not remove new values from

pitrou · pitrou · commit f939b3c0f76a · 2016-12-27T15:08:27.000+01:00
WeakValueDictionary when collecting from another thread.
diff --git a/Include/dictobject.h b/Include/dictobject.h
@@ -111,6 +111,9 @@ PyAPI_FUNC(PyObject *) PyDict_GetItem(PyObject *mp, PyObject *key);
 PyAPI_FUNC(PyObject *) _PyDict_GetItemWithError(PyObject *mp, PyObject *key);
 PyAPI_FUNC(int) PyDict_SetItem(PyObject *mp, PyObject *key, PyObject *item);
 PyAPI_FUNC(int) PyDict_DelItem(PyObject *mp, PyObject *key);
+PyAPI_FUNC(int) _PyDict_DelItemIf(PyObject *mp, PyObject *key,
+                                  int (*predicate)(PyObject *value));
+
 PyAPI_FUNC(void) PyDict_Clear(PyObject *mp);
 PyAPI_FUNC(int) PyDict_Next(
     PyObject *mp, Py_ssize_t *pos, PyObject **key, PyObject **value);
diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py
@@ -1437,6 +1437,18 @@ def test_threaded_weak_valued_pop(self):
                 x = d.pop(10, 10)
                 self.assertIsNot(x, None)  # we never put None in there!
 
+    def test_threaded_weak_valued_consistency(self):
+        # Issue #28427: old keys should not remove new values from
+        # WeakValueDictionary when collecting from another thread.
+        d = weakref.WeakValueDictionary()
+        with collect_in_thread():
+            for i in range(200000):
+                o = RefCycle()
+                d[10] = o
+                # o is still alive, so the dict can't be empty
+                self.assertEqual(len(d), 1)
+                o = None  # lose ref
+
 
 from test import mapping_tests
 
diff --git a/Lib/weakref.py b/Lib/weakref.py
@@ -18,7 +18,8 @@
      proxy,
      CallableProxyType,
      ProxyType,
-     ReferenceType)
+     ReferenceType,
+     _remove_dead_weakref)
 
 from _weakrefset import WeakSet, _IterationGuard
 
@@ -58,7 +59,9 @@ def remove(wr, selfref=ref(self)):
                 if self._iterating:
                     self._pending_removals.append(wr.key)
                 else:
-                    del self.data[wr.key]
+                    # Atomic removal is necessary since this function
+                    # can be called asynchronously by the GC
+                    _remove_dead_weakref(self.data, wr.key)
         self._remove = remove
         # A list of keys to be removed
         self._pending_removals = []
@@ -71,9 +74,12 @@ def _commit_removals(self):
         # We shouldn't encounter any KeyError, because this method should
         # always be called *before* mutating the dict.
         while l:
-            del d[l.pop()]
+            key = l.pop()
+            _remove_dead_weakref(d, key)
 
     def __getitem__(self, key):
+        if self._pending_removals:
+            self._commit_removals()
         o = self.data[key]()
         if o is None:
             raise KeyError, key
@@ -86,13 +92,17 @@ def __delitem__(self, key):
         del self.data[key]
 
     def __contains__(self, key):
+        if self._pending_removals:
+            self._commit_removals()
         try:
             o = self.data[key]()
         except KeyError:
             return False
         return o is not None
 
     def has_key(self, key):
+        if self._pending_removals:
+            self._commit_removals()
         try:
             o = self.data[key]()
         except KeyError:
@@ -113,6 +123,8 @@ def clear(self):
         self.data.clear()
 
     def copy(self):
+        if self._pending_removals:
+            self._commit_removals()
         new = WeakValueDictionary()
         for key, wr in self.data.items():
             o = wr()
@@ -124,6 +136,8 @@ def copy(self):
 
     def __deepcopy__(self, memo):
         from copy import deepcopy
+        if self._pending_removals:
+            self._commit_removals()
         new = self.__class__()
         for key, wr in self.data.items():
             o = wr()
@@ -132,6 +146,8 @@ def __deepcopy__(self, memo):
         return new
 
     def get(self, key, default=None):
+        if self._pending_removals:
+            self._commit_removals()
         try:
             wr = self.data[key]
         except KeyError:
@@ -145,6 +161,8 @@ def get(self, key, default=None):
                 return o
 
     def items(self):
+        if self._pending_removals:
+            self._commit_removals()
         L = []
         for key, wr in self.data.items():
             o = wr()
@@ -153,13 +171,17 @@ def items(self):
         return L
 
     def iteritems(self):
+        if self._pending_removals:
+            self._commit_removals()
         with _IterationGuard(self):
             for wr in self.data.itervalues():
                 value = wr()
                 if value is not None:
                     yield wr.key, value
 
     def iterkeys(self):
+        if self._pending_removals:
+            self._commit_removals()
         with _IterationGuard(self):
             for k in self.data.iterkeys():
                 yield k
@@ -176,11 +198,15 @@ def itervaluerefs(self):
         keep the values around longer than needed.
 
         """
+        if self._pending_removals:
+            self._commit_removals()
         with _IterationGuard(self):
             for wr in self.data.itervalues():
                 yield wr
 
     def itervalues(self):
+        if self._pending_removals:
+            self._commit_removals()
         with _IterationGuard(self):
             for wr in self.data.itervalues():
                 obj = wr()
@@ -212,13 +238,13 @@ def pop(self, key, *args):
             return o
 
     def setdefault(self, key, default=None):
+        if self._pending_removals:
+            self._commit_removals()
         try:
             o = self.data[key]()
         except KeyError:
             o = None
         if o is None:
-            if self._pending_removals:
-                self._commit_removals()
             self.data[key] = KeyedRef(default, self._remove, key)
             return default
         else:
@@ -254,9 +280,13 @@ def valuerefs(self):
         keep the values around longer than needed.
 
         """
+        if self._pending_removals:
+            self._commit_removals()
         return self.data.values()
 
     def values(self):
+        if self._pending_removals:
+            self._commit_removals()
         L = []
         for wr in self.data.values():
             o = wr()
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -15,6 +15,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #28427: old keys should not remove new values from
+  WeakValueDictionary when collecting from another thread.
+
 - Issue #28998: More APIs now support longs as well as ints.
 
 - Issue 28923: Remove editor artifacts from Tix.py,
diff --git a/Modules/_weakref.c b/Modules/_weakref.c
@@ -5,6 +5,43 @@
         ((PyWeakReference **) PyObject_GET_WEAKREFS_LISTPTR(o))
 
 
+static int
+is_dead_weakref(PyObject *value)
+{
+    if (!PyWeakref_Check(value)) {
+        PyErr_SetString(PyExc_TypeError, "not a weakref");
+        return -1;
+    }
+    return PyWeakref_GET_OBJECT(value) == Py_None;
+}
+
+PyDoc_STRVAR(remove_dead_weakref__doc__,
+"_remove_dead_weakref(dict, key) -- atomically remove key from dict\n"
+"if it points to a dead weakref.");
+
+static PyObject *
+remove_dead_weakref(PyObject *self, PyObject *args)
+{
+    PyObject *dct, *key;
+
+    if (!PyArg_ParseTuple(args, "O!O:_remove_dead_weakref",
+                          &PyDict_Type, &dct, &key)) {
+        return NULL;
+    }
+    if (_PyDict_DelItemIf(dct, key, is_dead_weakref) < 0) {
+        if (PyErr_ExceptionMatches(PyExc_KeyError))
+            /* This function is meant to allow safe weak-value dicts
+               with GC in another thread (see issue #28427), so it's
+               ok if the key doesn't exist anymore.
+               */
+            PyErr_Clear();
+        else
+            return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
 PyDoc_STRVAR(weakref_getweakrefcount__doc__,
 "getweakrefcount(object) -- return the number of weak references\n"
 "to 'object'.");
@@ -84,6 +121,8 @@ weakref_functions[] =  {
      weakref_getweakrefs__doc__},
     {"proxy",           weakref_proxy,                  METH_VARARGS,
      weakref_proxy__doc__},
+    {"_remove_dead_weakref", remove_dead_weakref,       METH_VARARGS,
+     remove_dead_weakref__doc__},
     {NULL, NULL, 0, NULL}
 };
 
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
@@ -848,13 +848,28 @@ PyDict_SetItem(register PyObject *op, PyObject *key, PyObject *value)
     return dict_set_item_by_hash_or_entry(op, key, hash, NULL, value);
 }
 
+static int
+delitem_common(PyDictObject *mp, PyDictEntry *ep)
+{
+    PyObject *old_value, *old_key;
+
+    old_key = ep->me_key;
+    Py_INCREF(dummy);
+    ep->me_key = dummy;
+    old_value = ep->me_value;
+    ep->me_value = NULL;
+    mp->ma_used--;
+    Py_DECREF(old_value);
+    Py_DECREF(old_key);
+    return 0;
+}
+
 int
 PyDict_DelItem(PyObject *op, PyObject *key)
 {
     register PyDictObject *mp;
     register long hash;
     register PyDictEntry *ep;
-    PyObject *old_value, *old_key;
 
     if (!PyDict_Check(op)) {
         PyErr_BadInternalCall();
@@ -875,15 +890,45 @@ PyDict_DelItem(PyObject *op, PyObject *key)
         set_key_error(key);
         return -1;
     }
-    old_key = ep->me_key;
-    Py_INCREF(dummy);
-    ep->me_key = dummy;
-    old_value = ep->me_value;
-    ep->me_value = NULL;
-    mp->ma_used--;
-    Py_DECREF(old_value);
-    Py_DECREF(old_key);
-    return 0;
+
+    return delitem_common(mp, ep);
+}
+
+int
+_PyDict_DelItemIf(PyObject *op, PyObject *key,
+                  int (*predicate)(PyObject *value))
+{
+    register PyDictObject *mp;
+    register long hash;
+    register PyDictEntry *ep;
+    int res;
+
+    if (!PyDict_Check(op)) {
+        PyErr_BadInternalCall();
+        return -1;
+    }
+    assert(key);
+    if (!PyString_CheckExact(key) ||
+        (hash = ((PyStringObject *) key)->ob_shash) == -1) {
+        hash = PyObject_Hash(key);
+        if (hash == -1)
+            return -1;
+    }
+    mp = (PyDictObject *)op;
+    ep = (mp->ma_lookup)(mp, key, hash);
+    if (ep == NULL)
+        return -1;
+    if (ep->me_value == NULL) {
+        set_key_error(key);
+        return -1;
+    }
+    res = predicate(ep->me_value);
+    if (res == -1)
+        return -1;
+    if (res > 0)
+        return delitem_common(mp, ep);
+    else
+        return 0;
 }
 
 void
