Make gevent.local.local eligible for GC as soon as the greenlet is dead,...

Make gevent.local.local eligible for GC as soon as the greenlet is dead, matching the native threading.local behaviour.

Also update the implementation to match Python 3.4. It appears to trade an extra weakref for better locality of reference.

changelog.rst

@@ -54,6 +54,10 @@ Unreleased
 - ``gevent.killall`` accepts an arbitrary iterable for the greenlets
   to kill. Reported in :issue:`404` by Martin Bachwerk; seen in
   combination with older versions of simple-requests.
+- ``gevent.local.local`` objects are now eligible for garbage
+  collection as soon as the greenlet finishes running, matching the
+  behaviour of the built-in ``threading.local`` (when implemented
+  natively). Reported in :issue:`387` by AusIV.
 
 1.1a1 (Jun 29, 2015)
 ====================

gevent/local.py

@@ -126,112 +126,147 @@ affects what we see:
 
 >>> del mydata
 """
-from weakref import WeakKeyDictionary
+
 from copy import copy
+from weakref import ref
+from contextlib import contextmanager
 from gevent.hub import getcurrent, PYPY
 from gevent.lock import RLock
 
 __all__ = ["local"]
 
 
-class _localbase(object):
-    __slots__ = '_local__args', '_local__lock', '_local__dicts'
+class _localimpl(object):
+    """A class managing thread-local dicts"""
+    __slots__ = 'key', 'dicts', 'localargs', 'locallock', '__weakref__'
+
+    def __init__(self):
+        # The key used in the Thread objects' attribute dicts.
+        # We keep it a string for speed but make it unlikely to clash with
+        # a "real" attribute.
+        self.key = '_threading_local._localimpl.' + str(id(self))
+        # { id(Thread) -> (ref(Thread), thread-local dict) }
+        self.dicts = {}
+
+    def get_dict(self):
+        """Return the dict for the current thread. Raises KeyError if none
+        defined."""
+        thread = getcurrent()
+        return self.dicts[id(thread)][1]
+
+    def create_dict(self):
+        """Create a new dict for the current thread, and return it."""
+        localdict = {}
+        key = self.key
+        thread = getcurrent()
+        idt = id(thread)
+
+        def local_deleted(_, key=key):
+            # When the localimpl is deleted, remove the thread attribute.
+            thread = wrthread()
+            if thread is not None:
+                del thread.__dict__[key]
+
+        def thread_deleted(_, idt=idt):
+            # When the thread is deleted, remove the local dict.
+            # Note that this is suboptimal if the thread object gets
+            # caught in a reference loop. We would like to be called
+            # as soon as the OS-level thread ends instead.
+            _local = wrlocal()
+            if _local is not None:
+                _local.dicts.pop(idt, None)
+        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)
+        return localdict
+
+
+@contextmanager
+def _patch(self):
+    impl = object.__getattribute__(self, '_local__impl')
+    orig_dct = object.__getattribute__(self, '__dict__')
+    try:
+        dct = impl.get_dict()
+    except KeyError:
+        # it's OK to acquire the lock here and not earlier, because the above code won't switch out
+        # however, subclassed __init__ might switch, so we do need to acquire the lock here
+        dct = impl.create_dict()
+        args, kw = impl.localargs
+        with impl.locallock:
+            self.__init__(*args, **kw)
+    with impl.locallock:
+        object.__setattr__(self, '__dict__', dct)
+        yield
+        object.__setattr__(self, '__dict__', orig_dct)
+
+
+class local(object):
+    __slots__ = '_local__impl', '__dict__'
 
     def __new__(cls, *args, **kw):
-        self = object.__new__(cls)
-        object.__setattr__(self, '_local__args', (args, kw))
-        object.__setattr__(self, '_local__lock', RLock())
-        dicts = WeakKeyDictionary()
-        object.__setattr__(self, '_local__dicts', dicts)
-
         if args or kw:
             if (PYPY and cls.__init__ == object.__init__) or (not PYPY and cls.__init__ is object.__init__):
                 raise TypeError("Initialization arguments are not supported")
-
-        # We need to create the greenlet dict in anticipation of
-        # __init__ being called, to make sure we don't call it again ourselves.
-        dict = object.__getattribute__(self, '__dict__')
-        dicts[getcurrent()] = dict
+        self = object.__new__(cls)
+        impl = _localimpl()
+        impl.localargs = (args, kw)
+        impl.locallock = RLock()
+        object.__setattr__(self, '_local__impl', impl)
+        # We need to create the thread dict in anticipation of
+        # __init__ being called, to make sure we don't call it
+        # again ourselves.
+        impl.create_dict()
         return self
 
-
-def _init_locals(self):
-    d = {}
-    dicts = object.__getattribute__(self, '_local__dicts')
-    dicts[getcurrent()] = d
-    object.__setattr__(self, '__dict__', d)
-
-    # we have a new instance dict, so call out __init__ if we have one
-    cls = type(self)
-    if cls.__init__ is not object.__init__:
-        args, kw = object.__getattribute__(self, '_local__args')
-        cls.__init__(self, *args, **kw)
-
-
-class local(_localbase):
-
     def __getattribute__(self, name):
-        d = object.__getattribute__(self, '_local__dicts').get(getcurrent())
-        if d is None:
-            # it's OK to acquire the lock here and not earlier, because the above code won't switch out
-            # however, subclassed __init__ might switch, so we do need to acquire the lock here
-            lock = object.__getattribute__(self, '_local__lock')
-            lock.acquire()
-            try:
-                _init_locals(self)
-                return object.__getattribute__(self, name)
-            finally:
-                lock.release()
-        else:
-            object.__setattr__(self, '__dict__', d)
+        with _patch(self):
             return object.__getattribute__(self, name)
 
     def __setattr__(self, name, value):
         if name == '__dict__':
-            raise AttributeError("%r object attribute '__dict__' is read-only" % self.__class__.__name__)
-        d = object.__getattribute__(self, '_local__dicts').get(getcurrent())
-        if d is None:
-            lock = object.__getattribute__(self, '_local__lock')
-            lock.acquire()
-            try:
-                _init_locals(self)
-                return object.__setattr__(self, name, value)
-            finally:
-                lock.release()
-        else:
-            object.__setattr__(self, '__dict__', d)
+            raise AttributeError(
+                "%r object attribute '__dict__' is read-only"
+                % self.__class__.__name__)
+        with _patch(self):
             return object.__setattr__(self, name, value)
 
     def __delattr__(self, name):
         if name == '__dict__':
-            raise AttributeError("%r object attribute '__dict__' is read-only" % self.__class__.__name__)
-        d = object.__getattribute__(self, '_local__dicts').get(getcurrent())
-        if d is None:
-            lock = object.__getattribute__(self, '_local__lock')
-            lock.acquire()
-            try:
-                _init_locals(self)
-                return object.__delattr__(self, name)
-            finally:
-                lock.release()
-        else:
-            object.__setattr__(self, '__dict__', d)
+            raise AttributeError(
+                "%r object attribute '__dict__' is read-only"
+                % self.__class__.__name__)
+        with _patch(self):
             return object.__delattr__(self, name)
 
     def __copy__(self):
-        currentId = getcurrent()
-        d = object.__getattribute__(self, '_local__dicts').get(currentId)
+        impl = object.__getattribute__(self, '_local__impl')
+        current = getcurrent()
+        currentId = id(current)
+        d = impl.get_dict()
         duplicate = copy(d)
 
         cls = type(self)
-        if cls.__init__ is not object.__init__:
-            args, kw = object.__getattribute__(self, '_local__args')
+        if (PYPY and cls.__init__ != object.__init__) or (not PYPY and cls.__init__ is not object.__init__):
+            args, kw = impl.localargs
             instance = cls(*args, **kw)
         else:
             instance = cls()
 
-        object.__setattr__(instance, '_local__dicts', {
-            currentId: duplicate
-        })
+        new_impl = object.__getattribute__(instance, '_local__impl')
+        tpl = new_impl.dicts[currentId]
+        new_impl.dicts[currentId] = (tpl[0], duplicate)
 
         return instance

greentest/greentest.py

@@ -257,7 +257,7 @@ class TestCaseMetaClass(type):
 
 
 class TestCase(TestCaseMetaClass("NewBase", (BaseTestCase,), {})):
-    __timeout__ = 1
+    __timeout__ = 1 if not os.environ.get('TRAVIS') else 2 # Travis is slow and overloaded
     switch_expected = 'default'
     error_fatal = True
 

greentest/test__local.py

@@ -58,5 +58,49 @@ class GeventLocalTestCase(greentest.TestCase):
 
         self.assertNotEqual(a.path, b.path, 'The values in the two objects must be different')
 
+    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()
+
+        def demonstrate_my_local():
+            # Get the important parts
+            getattr(my_local, 'sentinel')
+
+        # Create and reference greenlets
+        greenlets = [gevent.spawn(demonstrate_my_local) for _ in range(5)]
+        gevent.sleep()
+
+        self.assertEqual(len(created_sentinels), len(greenlets))
+
+        for g in greenlets:
+            assert g.dead
+        gevent.sleep() # let the callbacks run
+        if greentest.PYPY:
+            gc.collect()
+
+        # The sentinels should be gone too
+        self.assertEqual(len(deleted_sentinels), len(greenlets))
+
 if __name__ == '__main__':
     greentest.main()