Issue #9707: Rewritten reference implementation of threading.local which

is friendlier towards reference cycles.  This change is not normally
visible since an optimized C implementation (_thread._local) is used
instead.

3 files changed, 85 insertions, 85 deletions

diff --git a/Lib/_threading_local.py b/Lib/_threading_local.py
index 3af780715d..4ec4828144 100644
--- a/Lib/_threading_local.py
+++ b/Lib/_threading_local.py
@@ -132,6 +132,9 @@ affects what we see:
 >>> del mydata
 """
 
+from weakref import ref
+from contextlib import contextmanager
+
 __all__ = ["local"]
 
 # We need to use objects from the threading module, but the threading
@@ -139,112 +142,105 @@ __all__ = ["local"]
 # isn't compiled in to the `thread` module.  This creates potential problems
 # with circular imports.  For that reason, we don't import `threading`
 # until the bottom of this file (a hack sufficient to worm around the
-# potential problems).  Note that almost all platforms do have support for
-# locals in the `thread` module, and there is no circular import problem
+# potential problems).  Note that all platforms on CPython do have support
+# for locals in the `thread` module, and there is no circular import problem
 # then, so problems introduced by fiddling the order of imports here won't
-# manifest on most boxes.
+# manifest.
+
+class _localimpl:
+    """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 = current_thread()
+        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 = current_thread()
+        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:
+                dct = local.dicts.pop(idt)
+        wrlocal = ref(self, local_deleted)
+        wrthread = ref(thread, thread_deleted)
+        thread.__dict__[key] = wrlocal
+        self.dicts[idt] = wrthread, localdict
+        return localdict
+
+
+@contextmanager
+def _patch(self):
+    impl = object.__getattribute__(self, '_local__impl')
+    try:
+        dct = impl.get_dict()
+    except KeyError:
+        dct = impl.create_dict()
+        args, kw = impl.localargs
+        self.__init__(*args, **kw)
+    with impl.locallock:
+        object.__setattr__(self, '__dict__', dct)
+        yield
 
-class _localbase(object):
-    __slots__ = '_local__key', '_local__args', '_local__lock'
 
-    def __new__(cls, *args, **kw):
-        self = object.__new__(cls)
-        key = '_local__key', 'thread.local.' + str(id(self))
-        object.__setattr__(self, '_local__key', key)
-        object.__setattr__(self, '_local__args', (args, kw))
-        object.__setattr__(self, '_local__lock', RLock())
+class local:
+    __slots__ = '_local__impl', '__dict__'
 
+    def __new__(cls, *args, **kw):
         if (args or kw) and (cls.__init__ is object.__init__):
             raise TypeError("Initialization arguments are not supported")
-
+        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.
-        dict = object.__getattribute__(self, '__dict__')
-        current_thread().__dict__[key] = dict
-
+        impl.create_dict()
         return self
 
-def _patch(self):
-    key = object.__getattribute__(self, '_local__key')
-    d = current_thread().__dict__.get(key)
-    if d is None:
-        d = {}
-        current_thread().__dict__[key] = 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)
-    else:
-        object.__setattr__(self, '__dict__', d)
-
-class local(_localbase):
-
     def __getattribute__(self, name):
-        lock = object.__getattribute__(self, '_local__lock')
-        lock.acquire()
-        try:
-            _patch(self)
+        with _patch(self):
             return object.__getattribute__(self, name)
-        finally:
-            lock.release()
 
     def __setattr__(self, name, value):
         if name == '__dict__':
             raise AttributeError(
                 "%r object attribute '__dict__' is read-only"
                 % self.__class__.__name__)
-        lock = object.__getattribute__(self, '_local__lock')
-        lock.acquire()
-        try:
-            _patch(self)
+        with _patch(self):
             return object.__setattr__(self, name, value)
-        finally:
-            lock.release()
 
     def __delattr__(self, name):
         if name == '__dict__':
             raise AttributeError(
                 "%r object attribute '__dict__' is read-only"
                 % self.__class__.__name__)
-        lock = object.__getattribute__(self, '_local__lock')
-        lock.acquire()
-        try:
-            _patch(self)
+        with _patch(self):
             return object.__delattr__(self, name)
-        finally:
-            lock.release()
-
-    def __del__(self):
-        import threading
-
-        key = object.__getattribute__(self, '_local__key')
-
-        try:
-            # We use the non-locking API since we might already hold the lock
-            # (__del__ can be called at any point by the cyclic GC).
-            threads = threading._enumerate()
-        except:
-            # If enumerating the current threads fails, as it seems to do
-            # during shutdown, we'll skip cleanup under the assumption
-            # that there is nothing to clean up.
-            return
-
-        for thread in threads:
-            try:
-                __dict__ = thread.__dict__
-            except AttributeError:
-                # Thread is dying, rest in peace.
-                continue
-
-            if key in __dict__:
-                try:
-                    del __dict__[key]
-                except KeyError:
-                    pass # didn't have anything in this thread
+
 
 from threading import current_thread, RLock
diff --git a/Lib/test/test_threading_local.py b/Lib/test/test_threading_local.py
index acf37a038f..c886a25d8a 100644
--- a/Lib/test/test_threading_local.py
+++ b/Lib/test/test_threading_local.py
@@ -184,11 +184,6 @@ class BaseLocalTest:
             """To test that subclasses behave properly."""
         self._test_dict_attribute(LocalSubclass)
 
-
-class ThreadLocalTest(unittest.TestCase, BaseLocalTest):
-    _local = _thread._local
-
-    # Fails for the pure Python implementation
     def test_cycle_collection(self):
         class X:
             pass
@@ -201,6 +196,10 @@ class ThreadLocalTest(unittest.TestCase, BaseLocalTest):
         gc.collect()
         self.assertIs(wr(), None)
 
+
+class ThreadLocalTest(unittest.TestCase, BaseLocalTest):
+    _local = _thread._local
+
 class PyThreadingLocalTest(unittest.TestCase, BaseLocalTest):
     _local = _threading_local.local
 
diff --git a/Misc/NEWS b/Misc/NEWS
index f3c5592196..f4d2bb7d31 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -13,6 +13,11 @@ Core and Builtins
 Library
 -------
 
+- Issue #9707: Rewritten reference implementation of threading.local which
+  is friendlier towards reference cycles.  This change is not normally
+  visible since an optimized C implementation (_thread._local) is used
+  instead.
+
 - Issue #6394: os.getppid() is now supported on Windows.  Note that it will
   still return the id of the parent process after it has exited.  This process
   id may even have been reused by another unrelated process.