Optimize gevent.local to create fewer weakrefs. Fix leakcheck on the new local test.

gevent/local.py

@@ -136,6 +136,10 @@ from gevent.lock import RLock
 __all__ = ["local"]
 
 
+class _wrefdict(dict):
+    """A dict that can be weak referenced"""
+
+
 class _localimpl(object):
     """A class managing thread-local dicts"""
     __slots__ = 'key', 'dicts', 'localargs', 'locallock', '__weakref__'
@@ -146,7 +150,7 @@ class _localimpl(object):
         # a "real" attribute.
         self.key = '_threading_local._localimpl.' + str(id(self))
         # { id(Thread) -> (ref(Thread), thread-local dict) }
-        self.dicts = {}
+        self.dicts = _wrefdict()
 
     def get_dict(self):
         """Return the dict for the current thread. Raises KeyError if none
@@ -161,6 +165,13 @@ class _localimpl(object):
         thread = getcurrent()
         idt = id(thread)
 
+        # If we are working with a gevent.greenlet.Greenlet, we can
+        # pro-actively clear out with a link. Use rawlink to avoid
+        # spawning any more greenlets
+        try:
+            rawlink = thread.rawlink
+        except AttributeError:
+            # Otherwise we need to do it with weak refs
             def local_deleted(_, key=key):
                 # When the localimpl is deleted, remove the thread attribute.
                 thread = wrthread()
@@ -178,19 +189,17 @@ class _localimpl(object):
             wrlocal = ref(self, local_deleted)
             wrthread = ref(thread, thread_deleted)
             thread.__dict__[key] = wrlocal
-        self.dicts[idt] = wrthread, localdict
-
-        # If we are working with a gevent.greenlet.Greenlet, we can
-        # pro-actively clear out with a link. Use rawlink to avoid
-        # spawning any more greenlets
-        try:
-            rawlink = thread.rawlink
-        except AttributeError:
-            pass
         else:
-            def greenlet_dead(_):
-                thread_deleted(_)
-            rawlink(greenlet_dead)
+            wrdicts = ref(self.dicts)
+
+            def clear(_):
+                dicts = wrdicts()
+                if dicts:
+                    dicts.pop(idt, None)
+            rawlink(clear)
+            wrthread = None
+
+        self.dicts[idt] = wrthread, localdict
         return localdict
 
 

greentest/test__local.py

@@ -21,6 +21,22 @@ class A(local):
 class Obj(object):
     pass
 
+# These next two classes have to be global to avoid the leakchecks
+deleted_sentinels = []
+created_sentinels = []
+
+
+class Sentinel(object):
+    def __del__(self):
+        deleted_sentinels.append(id(self))
+
+
+class MyLocal(local):
+    def __init__(self):
+        local.__init__(self)
+        self.sentinel = Sentinel()
+        created_sentinels.append(id(self.sentinel))
+
 
 class GeventLocalTestCase(greentest.TestCase):
 
@@ -61,27 +77,14 @@ class GeventLocalTestCase(greentest.TestCase):
     def test_locals_collected_when_greenlet_dead_but_still_referenced(self):
         # https://github.com/gevent/gevent/issues/387
         import gevent
-        deleted_sentinels = []
-        created_sentinels = []
-
-        class Sentinel(object):
-            def __del__(self):
-                deleted_sentinels.append(id(self))
-
-        class MyLocal(local):
-            def __init__(self):
-                local.__init__(self)
-                self.sentinel = Sentinel()
-                created_sentinels.append(id(self.sentinel))
 
         my_local = MyLocal()
         my_local.sentinel = None
         if greentest.PYPY:
             import gc
             gc.collect()
-        # drop the original
-        created_sentinels.pop()
-        deleted_sentinels.pop()
+        del created_sentinels[:]
+        del deleted_sentinels[:]
 
         def demonstrate_my_local():
             # Get the important parts