Write a WeakValueDictionary with safer key removal

[nbruin]

On ticket #12313 we found that the use of WeakValueDictionaries as caches can cause removal callbacks in rather harsh environments. Normal WeakValueDictionaries remove keys with dead values by looking up the key. This involves Python equality testing on the key, which can cause any kind of operation in Sage. We need a dictionary where we can delete entries without key comparisons. See below for possible strategies.

To the release manager:

Apply attachment: trac13394-weak_value_dictionary.patch

Upstream: None of the above - read trac for reasoning.

CC: @simon-king-jena

Component: memleak

Author: Simon King, Nils Bruin

Reviewer: Nils Bruin, Simon King

Merged: sage-5.13.beta3

Issue created by migration from https://trac.sagemath.org/ticket/13394

[nbruin]

Description changed:

--- 
+++ 
@@ -1 +1 @@
-On ticket #12323 we found that the use of `WeakValueDictionaries` as caches can cause `removal` callbacks in rather harsh environments. Normal `WeakValueDictionaries` remove keys with dead values by looking up the key. This involves Python equality testing on the key, which can cause any kind of operation in Sage. We need a dictionary where we can delete entries without key comparisons. See below for possible strategies.
+On ticket #12313 we found that the use of `WeakValueDictionaries` as caches can cause `removal` callbacks in rather harsh environments. Normal `WeakValueDictionaries` remove keys with dead values by looking up the key. This involves Python equality testing on the key, which can cause any kind of operation in Sage. We need a dictionary where we can delete entries without key comparisons. See below for possible strategies.

[nbruin]

comment:2

The standard Python WeakValueDictionary is built on top of a normal dictionary, built on top of a normal dictionary. Let's say we have

W = WeakValueDictionary()

There is an underlying D = dict. If we do

W[k] = v

then it really executes

D[k] = KeyedRef(v, k, remove_callback)

This is a weak reference to v, so v can still get collected. If it does, the callback can ensure that

del D[k]

gets executed. This locates the entry by computing the hash of h, finding the right bucket and finding an equal key by comparison. When we create the entry, we have better lookup information, though: We already have hash(k) and id(v). These two don't necessarily pin down the entry in the dictionary, but in our applications they do and for a WeakValueDictionary identical v would be the same (dead) KeyedRef anyway, so it wouldn't hurt to remove them all. We need one more method on our dictionary:

D.remove_entry_by_hash_and_value_id(h,i):
    """
    remove all entries { k : v } from the dictionary that have
    hash(k) == h   and  id(v) == i
    """ 

Then the callback can be

function remove(ref):
    D.remove_entry_by_hash_and_value_id(hash(ref.key),id(ref))

and the KeyedRef that gets stored would be

KeyedRef(v,hash(k),remove)

(normal precautions to put the dict reference in a closure or a class should apply here in defining remove)

Or one could just immediately integrate the KeyedRefs in the design. See also #13387 for discussions about dictionary designs.

In either case, the main effect is that removal does not cause any hashing or comparison methods to be called on k anymore.

[simon-king-jena]

Attachment: weak_value_dictionary.pyx.gz

Proof of concept

[simon-king-jena]

comment:5

I have attached a proof of concept. I think the class WVD defined in the attachment is feature-complete. According to my benchmarks, it is faster than weakref.WeakValueDictionary, and it is safer.

Performance

sage: def test_dict(D,L):
....:     for k,v in L:
....:         D[k] = v
....:     for k,v in L:
....:         if k not in D:
....:             raise RuntimeError("containment")
....:     for k,v in L:
....:         D[k] = v
....:     for k in D.iterkeys():
....:         if k not in D:
....:             raise RuntimeError("second containment")
....:     for k,v in L:
....:         assert D.get(k)==v
....:         del D[k]
....:     for k,v in L:
....:         if k in D:
....:             raise RuntimeError("anti-containment")
....:         
sage: L = [(p,GF(p)) for p in prime_range(2*10^4)]
sage: from weakref import WeakValueDictionary
sage: D = WeakValueDictionary()
sage: %timeit test_dict(D,L)
10 loops, best of 3: 43.2 ms per loop
sage: D = WVD()
sage: %timeit test_dict(D,L)
10 loops, best of 3: 29.4 ms per loop

Hence, the WVD is faster than the weakref.WeakValueDictionary.

Safety

sage: class Vals(object): pass
sage: class Keys:                       
....:     def __init__(self, val):
....:         self.val = weakref.ref(val)
....:     def __hash__(self):
....:         return hash(self.val())
....:     def __eq__(self, other):
....:         return self.val() == other.val()
....: 
sage: ValList = [Vals() for _ in range(10)]
sage: D = WeakValueDictionary()
sage: for v in ValList:
....:     D[Keys(v)] = v
....:     
sage: len(D)
10
sage: del ValList
Exception KeyError: (<__main__.Keys instance at 0xdcfc96c>,) in <function remove at 0xdcb9df4> ignored
Exception KeyError: (<__main__.Keys instance at 0xdcfc90c>,) in <function remove at 0xdcb9df4> ignored
Exception KeyError: (<__main__.Keys instance at 0xdcfcc2c>,) in <function remove at 0xdcb9df4> ignored
Exception KeyError: (<__main__.Keys instance at 0xdcfc9ec>,) in <function remove at 0xdcb9df4> ignored
Exception KeyError: (<__main__.Keys instance at 0xdcfca4c>,) in <function remove at 0xdcb9df4> ignored
Exception KeyError: (<__main__.Keys instance at 0xdcfc9cc>,) in <function remove at 0xdcb9df4> ignored
Exception KeyError: (<__main__.Keys instance at 0xdcfc98c>,) in <function remove at 0xdcb9df4> ignored
Exception KeyError: (<__main__.Keys instance at 0xdcfccac>,) in <function remove at 0xdcb9df4> ignored
Exception KeyError: (<__main__.Keys instance at 0xdcfca0c>,) in <function remove at 0xdcb9df4> ignored
sage: len(D)
10

Hence, the weakref.WeakValueDictionary is unsafe.

sage: ValList = [Vals() for _ in range(10)]
sage: D = WVD()
sage: for v in ValList:
....:     D[Keys(v)] = v
....:     
sage: len(D)
10
sage: del ValList
sage: len(D)
1
sage: del v
sage: len(D)
0

Hence, the WVD is safe (or at least: safer...)

I suppose I should add tests and documentation to the proof of concept, put into sage.misc and rename WVD into sage.misc.weak_dict.WeakValueDictionary. And then we should see if we can really smoothly replace weakref.WeakValueDictionary in Sage.

[nbruin]

comment:6

Excellent work! I like how you found a way to reuse python dictionaries. I think there is one fairly easy optimization that you can make that is fairly similar to how we ended up implementing the buckets for MonoDict and TripleDict: Presently, your bucket is a list of tuples. That provides an extra layer of indirection, meaning both slower access and more memory use and allocation.

If instead you "inline" the tuples by making the bucket a list with the layout [key1, weakref_to_value1, key2, weakref_to_value2,...] you can save yourself some overhead. Normally these lists should only contain one key-value pair anyway.

It may be worth browsing through MonoDict anyway. There may more more little tricks that we used there that may apply here and that I don't remember right now.

[simon-king-jena]

comment:7

Replying to @nbruin:

If instead you "inline" the tuples by making the bucket a list with the layout [key1, weakref_to_value1, key2, weakref_to_value2,...] you can save yourself some overhead.

Good point.

In addition, I think that weakref_to_value_n.key used for the callback should be improved: Storing hash(key_n) should be enough. Namely, it allows to find the correct hash bucket, and then the hash bucket is searched to find weakref_to_value_n that is subject to the callback. The latter is simply comparing memory addresses.

[simon-king-jena]

Branch: u/SimonKing/ticket/13394

[simon-king-jena]

comment:9

I have posted an initial version, that still lacks documentation and tests (and with the next commit I will also fix some trailing whitespace).

I am afraid that the current version of a weak value dictionary is not faster than what I have described in the proof of concept, although I used the improved bucket layout and am using C-API functions rather thoroughly.

Anyway. I guess part of the documentation will be benchmarks for each separate task. Then we will see what should be further improved.