py3: adapt support of multitenant in threads (#44021)

This commit is contained in:
Benjamin Dauvergne 2020-06-15 11:56:12 +02:00
parent 6361be78c9
commit c448546dbe
2 changed files with 144 additions and 114 deletions

View File

@ -14,13 +14,142 @@
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import sys
import threading
from django.utils import six
class TenantAwareThread(threading.Thread):
def __init__(self, *args, **kwargs):
if six.PY2:
class TenantAwareThread(threading.Thread):
def __init__(self, *args, **kwargs):
from django.db import connection
try:
if hasattr(connection, 'get_tenant'):
self.tenant = connection.get_tenant()
else:
self.tenant = None
except RuntimeError:
# this happens when ImportError is raised at startup; ignore
# the error to let the real one be displayed.
self.tenant = None
super(TenantAwareThread, self).__init__(*args, **kwargs)
def run(self):
if self.tenant is not None:
from django.db import connection
old_tenant = connection.get_tenant()
connection.set_tenant(self.tenant)
try:
super(TenantAwareThread, self).run()
finally:
connection.set_tenant(old_tenant)
connection.close()
else:
super(TenantAwareThread, self).run()
class _Timer(TenantAwareThread):
"""Call a function after a specified number of seconds:
t = Timer(30.0, f, args=[], kwargs={})
t.start()
t.cancel() # stop the timer's action if it's still waiting
"""
def __init__(self, interval, function, args=[], kwargs={}):
super(_Timer, self).__init__()
self.interval = interval
self.function = function
self.args = args
self.kwargs = kwargs
self.finished = threading.Event()
def cancel(self):
"""Stop the timer if it hasn't finished yet"""
self.finished.set()
def run(self):
self.finished.wait(self.interval)
if not self.finished.is_set():
self.function(*self.args, **self.kwargs)
self.finished.set()
class _MainThread(TenantAwareThread):
def __init__(self):
if six.PY3:
super(_MainThread, self).__init__(name="MainThread", daemon=False)
self._set_tstate_lock()
self._started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[self._ident] = self
else:
super(_MainThread, self).__init__(name="MainThread")
self._Thread__started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[threading._get_ident()] = self
def _set_daemon(self):
return False
def _exitfunc(self):
self._Thread__stop()
t = threading._pickSomeNonDaemonThread()
if t:
if __debug__:
self._note("%s: waiting for other threads", self)
while t:
t.join()
t = threading._pickSomeNonDaemonThread()
if __debug__:
self._note("%s: exiting", self)
self._Thread__delete()
class _DummyThread(TenantAwareThread):
def __init__(self):
if six.PY3:
super(_DummyThread, self).__init__(name=threading._newname("Dummy-%d"), daemon=True)
self._started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[self._ident] = self
return
super(_DummyThread, self).__init__(name=threading._newname("Dummy-%d"))
# Thread.__block consumes an OS-level locking primitive, which
# can never be used by a _DummyThread. Since a _DummyThread
# instance is immortal, that's bad, so release this resource.
del self._Thread__block
self._Thread__started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[threading._get_ident()] = self
def _set_daemon(self):
return True
def join(self, timeout=None):
assert False, "cannot join a dummy thread"
def install_tenant_aware_threads():
if getattr(threading, 'multitenant', False):
return
# Install tenant aware Thread class
threading.Thread = TenantAwareThread
threading._DummyThread = _DummyThread
threading._MainThread = _MainThread
threading._Timer = _Timer
threading.multitenant = True
else:
_Thread_bootstrap_inner = threading.Thread._bootstrap_inner
_Thread__init__ = threading.Thread.__init__
def _new__init__(self, *args, **kwargs):
from django.db import connection
try:
if hasattr(connection, 'get_tenant'):
@ -31,120 +160,21 @@ class TenantAwareThread(threading.Thread):
# this happens when ImportError is raised at startup; ignore
# the error to let the real one be displayed.
self.tenant = None
super(TenantAwareThread, self).__init__(*args, **kwargs)
_Thread__init__(self, *args, **kwargs)
def run(self):
def _new_bootstrap_inner(self):
if self.tenant is not None:
from django.db import connection
old_tenant = connection.get_tenant()
connection.set_tenant(self.tenant)
try:
super(TenantAwareThread, self).run()
_Thread_bootstrap_inner(self)
finally:
connection.set_tenant(old_tenant)
connection.close()
else:
super(TenantAwareThread, self).run()
_Thread_bootstrap_inner(self)
class _Timer(TenantAwareThread):
"""Call a function after a specified number of seconds:
t = Timer(30.0, f, args=[], kwargs={})
t.start()
t.cancel() # stop the timer's action if it's still waiting
"""
def __init__(self, interval, function, args=[], kwargs={}):
super(_Timer, self).__init__()
self.interval = interval
self.function = function
self.args = args
self.kwargs = kwargs
self.finished = threading.Event()
def cancel(self):
"""Stop the timer if it hasn't finished yet"""
self.finished.set()
def run(self):
self.finished.wait(self.interval)
if not self.finished.is_set():
self.function(*self.args, **self.kwargs)
self.finished.set()
class _MainThread(TenantAwareThread):
def __init__(self):
if six.PY3:
super(_MainThread, self).__init__(name="MainThread", daemon=False)
self._set_tstate_lock()
self._started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[self._ident] = self
else:
super(_MainThread, self).__init__(name="MainThread")
self._Thread__started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[threading._get_ident()] = self
def _set_daemon(self):
return False
def _exitfunc(self):
self._Thread__stop()
t = threading._pickSomeNonDaemonThread()
if t:
if __debug__:
self._note("%s: waiting for other threads", self)
while t:
t.join()
t = threading._pickSomeNonDaemonThread()
if __debug__:
self._note("%s: exiting", self)
self._Thread__delete()
class _DummyThread(TenantAwareThread):
def __init__(self):
if six.PY3:
super(_DummyThread, self).__init__(name=threading._newname("Dummy-%d"), daemon=True)
self._started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[self._ident] = self
return
super(_DummyThread, self).__init__(name=threading._newname("Dummy-%d"))
# Thread.__block consumes an OS-level locking primitive, which
# can never be used by a _DummyThread. Since a _DummyThread
# instance is immortal, that's bad, so release this resource.
del self._Thread__block
self._Thread__started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[threading._get_ident()] = self
def _set_daemon(self):
return True
def join(self, timeout=None):
assert False, "cannot join a dummy thread"
def install_tenant_aware_threads():
if getattr(threading, 'multitenant', False):
return
# Install tenant aware Thread class
threading.Thread = TenantAwareThread
threading._DummyThread = _DummyThread
threading._MainThread = _MainThread
threading._Timer = _Timer
threading.multitenant = True
def install_tenant_aware_threads():
threading.Thread.__init__ = _new__init__
threading.Thread._bootstrap_inner = _new_bootstrap_inner

View File

@ -46,8 +46,8 @@ def test_thread(tenants, settings, client):
with tenant_context(tenant):
assert hasattr(settings, 'TEMPLATE_VARS')
t2 = threading.Thread(target=f, args=(tenant,))
t2.start()
t2.join()
t2.start()
t2.join()
assert not hasattr(django.conf.settings, 'TEMPLATE_VARS')
t3 = threading.Thread(target=f)
@ -73,8 +73,8 @@ def test_cache(tenants, client):
def f():
assert cache.get('coin') == tenant.domain_url
t1 = threading.Thread(target=f)
t1.start()
t1.join()
t1.start()
t1.join()
def g():
assert cache.get('coin') == 1
@ -104,8 +104,8 @@ def test_timer_thread(tenants, settings, client):
with tenant_context(tenant):
assert hasattr(settings, 'TEMPLATE_VARS')
t2 = threading.Timer(0.0, f, args=(tenant,))
t2.start()
t2.join()
t2.start()
t2.join()
assert not hasattr(django.conf.settings, 'TEMPLATE_VARS')
t3 = threading.Timer(0.0, f)