solar/solar/dblayer/gevent_local.py

# flake8: noqa
# Except when otherwise stated (look at the beginning of each file) the
# software and the documentation in this project are copyrighted by:
#
#  Denis Bilenko and the contributors, http://www.gevent.org
#
# and licensed under the MIT license:
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

# Portions copyright 2015 by Mirantis, Inc.


from contextlib import contextmanager
from copy import copy

from gevent.hub import getcurrent
from gevent.hub import PYPY
from gevent.lock import RLock

from weakref import ref

__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__'

    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 = _wrefdict()

    def find_parent(self):
        """Iterate to top most parent, and use it as a base"""
        c = getcurrent()
        while 1:
            tmp_c = getattr(c, '_nested_parent', None)
            if not tmp_c:
                return c
            c = tmp_c

    def get_dict(self):
        """Return the dict for the current thread.

        Raises KeyError if none defined.
        """
        # thread = getcurrent()
        thread = self.find_parent()
        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 = self.find_parent()
        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()
                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
        else:
            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


@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):
        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")
        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 __getattribute__(self, name):
        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__)
        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__)
        with _patch(self):
            return object.__delattr__(self, name)

    def __copy__(self):
        impl = object.__getattribute__(self, '_local__impl')
        current = impl.find_parent()
        currentId = id(current)
        d = impl.get_dict()
        duplicate = copy(d)

        cls = type(self)
        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()

        new_impl = object.__getattribute__(instance, '_local__impl')
        tpl = new_impl.dicts[currentId]
        new_impl.dicts[currentId] = (tpl[0], duplicate)

        return instance