Replaced threading.Semaphore with asyncio.Semaphore, improved tests and cleaned everything up.

Author: excitoon

django_threaded_sync_to_async/__init__.py

@@ -1,3 +1,4 @@
+import asyncio
 import concurrent.futures
 import contextlib
 import contextvars
@@ -6,70 +7,13 @@
 
 import asgiref.sync
 
-
-_reentrant_patch_lock = threading.Lock()
-
-
-@contextlib.contextmanager
-def reentrant_patch(obj, attr, value):
-    """
-    Makes time-aware patch on the attribute of the object without locking like in `unittest.mock.patch`, the context
-    will leak system-wide.
-    However, if no `await` happens after obtaining the context, and no threads are getting the same attribute,
-    it guarantees that the attribute will have the desired value.
-    Effectively guarantees to restore original value after all contexts are destroyed.
-    No protection from interleaving foreign code doing same.
-    """
-
-    with _reentrant_patch_lock:
-        contexts = getattr(obj, f"__{attr}__contexts__", {})
-        if not contexts:
-            contexts[1] = getattr(obj, attr)
-            setattr(obj, f"__{attr}__contexts__", contexts)
-        context_id = len(contexts) + 1
-        contexts[context_id] = value
-        setattr(obj, attr, value)
-
-    yield
-
-    with _reentrant_patch_lock:
-        last_context_id = next(reversed(contexts))
-        del contexts[context_id]
-        if last_context_id == context_id:
-            setattr(obj, attr, next(reversed(contexts.values())))
-        if len(contexts) == 1:
-            delattr(obj, f"__{attr}__contexts__")
-
-
-_one_time_patch_lock = threading.Lock()
-
-
-@contextlib.contextmanager
-def one_time_patch(obj, attr, value):
-    """
-    More lightweight implementation, only sets the attribute once — in outer context.
-    Effectively guarantees to restore original value after all contexts are destroyed.
-    """
-
-    with _one_time_patch_lock:
-        if not hasattr(obj, f"__{attr}__value__"):
-            setattr(obj, f"__{attr}__value__", getattr(obj, attr))
-            setattr(obj, attr, value)
-        setattr(obj, f"__{attr}__count__", getattr(obj, f"__{attr}__count__", 0) + 1)
-
-    yield
-
-    with _one_time_patch_lock:
-        count = getattr(obj, f"__{attr}__count__") - 1
-        setattr(obj, f"__{attr}__count__", count)
-        if not count:
-            setattr(obj, attr, getattr(obj, f"__{attr}__value__"))
-            delattr(obj, f"__{attr}__value__")
-            delattr(obj, f"__{attr}__count__")
+import django_threaded_sync_to_async.patch
 
 
 _current_executor = contextvars.ContextVar("current_executor", default=None)
 _max_tasks_semaphore = contextvars.ContextVar("max_tasks_semaphore", default=None)
+_shared_executors = {}
+_shared_executors_lock = threading.Lock()
 
 
 async def _sync_to_async_call(self, orig, *args, **kwargs):
@@ -78,14 +22,10 @@ async def _sync_to_async_call(self, orig, *args, **kwargs):
 
     else:
         """
-        The task is called outside of executor's scope (or in different context).
+        The task hit the call outside of executor's scope (or in different context).
         """
 
-    if _max_tasks_semaphore.get() is not None:
-        with _max_tasks_semaphore.get():
-            return await orig(self, *args, **kwargs)
-
-    else:
+    async with _max_tasks_semaphore.get() or contextlib.nullcontext(): # Python 3.10+.
         return await orig(self, *args, **kwargs)
 
 
@@ -98,11 +38,10 @@ def _set_context_variable(variable, value):
 
 @contextlib.contextmanager
 def _use_executor(executor):
-    with _set_context_variable(_current_executor, executor):
-        # Can be replaced by a single call to `setattr(asgiref.sync.SyncToAsync, "__call__", new_call)`
-        # if we don't care about restoring everything back.
-        new_call = functools.partialmethod(_sync_to_async_call, asgiref.sync.SyncToAsync.__call__)
-        with one_time_patch(asgiref.sync.SyncToAsync, "__call__", new_call):
+    # `patch.one_time()` can be replaced with `patch.permanent()` if we don't care about restoring everything back.
+    new_call = functools.partialmethod(_sync_to_async_call, asgiref.sync.SyncToAsync.__call__)
+    with django_threaded_sync_to_async.patch.one_time(asgiref.sync.SyncToAsync, "__call__", new_call):
+        with _set_context_variable(_current_executor, executor):
             yield executor
 
 
@@ -113,10 +52,6 @@ async def Executor(*args, **kwargs):
             yield executor
 
 
-_shared_executors = {}
-_shared_executors_lock = threading.Lock()
-
-
 @contextlib.asynccontextmanager
 async def SharedExecutor(name, *args, max_tasks=None, **kwargs):
     with _shared_executors_lock:
@@ -128,6 +63,6 @@ async def SharedExecutor(name, *args, max_tasks=None, **kwargs):
             kwargs.setdefault("thread_name_prefix", name)
             executor = _shared_executors[name] = concurrent.futures.ThreadPoolExecutor(*args, **kwargs)
 
-    with _set_context_variable(_max_tasks_semaphore, max_tasks and threading.Semaphore(max_tasks)):
+    with _set_context_variable(_max_tasks_semaphore, max_tasks and asyncio.Semaphore(max_tasks)):
         with _use_executor(executor):
             yield executor

django_threaded_sync_to_async/patch.py

@@ -0,0 +1,72 @@
+import contextlib
+import threading
+
+
+_reentrant_lock = threading.Lock()
+_one_time_lock = threading.Lock()
+
+
+@contextlib.contextmanager
+def reentrant(obj, attr, value):
+    """
+    Makes time-aware patch on the attribute of the object without locking like in `unittest.mock.patch`, the context
+    will leak system-wide.
+    However, if no `await` happens after obtaining the context, and no threads are getting the same attribute,
+    it guarantees that the attribute will have the desired value.
+    Effectively guarantees to restore original value after all contexts are destroyed.
+    No protection from interleaving foreign code doing same.
+    """
+
+    with _reentrant_lock:
+        contexts = getattr(obj, f"__{attr}__contexts__", {})
+        if not contexts:
+            contexts[1] = getattr(obj, attr)
+            setattr(obj, f"__{attr}__contexts__", contexts)
+        context_id = len(contexts) + 1
+        contexts[context_id] = value
+        setattr(obj, attr, value)
+
+    yield
+
+    with _reentrant_lock:
+        last_context_id = next(reversed(contexts))
+        del contexts[context_id]
+        if last_context_id == context_id:
+            setattr(obj, attr, next(reversed(contexts.values())))
+        if len(contexts) == 1:
+            delattr(obj, f"__{attr}__contexts__")
+
+
+@contextlib.contextmanager
+def one_time(obj, attr, value):
+    """
+    More lightweight implementation, only sets the attribute once — in outer context.
+    Effectively guarantees to restore original value after all contexts are destroyed.
+    """
+
+    with _one_time_lock:
+        if not hasattr(obj, f"__{attr}__value__"):
+            setattr(obj, f"__{attr}__value__", getattr(obj, attr))
+            setattr(obj, attr, value)
+        setattr(obj, f"__{attr}__count__", getattr(obj, f"__{attr}__count__", 0) + 1)
+
+    yield
+
+    with _one_time_lock:
+        count = getattr(obj, f"__{attr}__count__") - 1
+        setattr(obj, f"__{attr}__count__", count)
+        if not count:
+            setattr(obj, attr, getattr(obj, f"__{attr}__value__"))
+            delattr(obj, f"__{attr}__value__")
+            delattr(obj, f"__{attr}__count__")
+
+
+@contextlib.contextmanager
+def permanent(obj, attr, value):
+    """
+    Most lightweight implementation.
+    """
+
+    setattr(obj, attr, value)
+
+    yield

example.py

@@ -38,7 +38,6 @@ async def test():
         await four_calls()
 
 
-# N.B. `contextlib.asynccontextmanager` only works as decorator since Python 3.10.
 @django_threaded_sync_to_async.Executor(thread_name_prefix="thread", max_workers=3)
 async def test2():
     await four_calls()

setup.py

@@ -21,10 +21,6 @@
                 "License :: OSI Approved :: MIT License",
                 "Operating System :: OS Independent",
                 "Programming Language :: Python :: 3 :: Only",
-                "Programming Language :: Python :: 3.6",
-                "Programming Language :: Python :: 3.7",
-                "Programming Language :: Python :: 3.8",
-                "Programming Language :: Python :: 3.9",
                 "Programming Language :: Python :: 3.10",
                 "Programming Language :: Python :: 3.11",
                 "Programming Language :: Python :: 3.12",

tests/test_executor.py

@@ -39,6 +39,7 @@ async def empty(**kwargs):
                 with self.subTest(parallel=parallel, decorated=decorated):
                     cv = threading.Condition()
                     threads = set()
+
                     async with context(max_workers=workers):
                         tt = [asyncio.create_task(function(cv, threads)) for _ in range(workers)]
                         try:
@@ -52,3 +53,5 @@ async def empty(**kwargs):
                         else:
                             if not parallel:
                                 self.assertEqual("No", "exception")
+
+                    self.assertEqual(len(threads), workers if parallel else 1)

tests/test_one_time_patch.py

@@ -1,6 +1,6 @@
 import unittest
 
-import django_threaded_sync_to_async
+import django_threaded_sync_to_async.patch
 
 
 class TestOneTimePatch(unittest.TestCase):
@@ -13,10 +13,10 @@ class Dummy:
 
         with self.subTest(inside=[]):
             self.assertEqual(o.x, 1)
-        with django_threaded_sync_to_async.one_time_patch(o, "x", 2):
+        with django_threaded_sync_to_async.patch.one_time(o, "x", 2):
             with self.subTest(inside=["c1"]):
                 self.assertEqual(o.x, 2)
-            with django_threaded_sync_to_async.one_time_patch(o, "x", 3):
+            with django_threaded_sync_to_async.patch.one_time(o, "x", 3):
                 with self.subTest(inside=["c1", "c2"]):
                     self.assertEqual(o.x, 2)
             with self.subTest(inside=["c1"]):
@@ -34,8 +34,8 @@ class Dummy:
         o = Dummy()
         fields = dir(o)
 
-        c1 = django_threaded_sync_to_async.one_time_patch(o, "x", 2)
-        c2 = django_threaded_sync_to_async.one_time_patch(o, "x", 3)
+        c1 = django_threaded_sync_to_async.patch.one_time(o, "x", 2)
+        c2 = django_threaded_sync_to_async.patch.one_time(o, "x", 3)
 
         with self.subTest(inside=[]):
             self.assertEqual(o.x, 1)

tests/test_reentrant_patch.py

@@ -1,6 +1,6 @@
 import unittest
 
-import django_threaded_sync_to_async
+import django_threaded_sync_to_async.patch
 
 
 class TestReentrantPatch(unittest.TestCase):
@@ -13,10 +13,10 @@ class Dummy:
 
         with self.subTest(inside=[]):
             self.assertEqual(o.x, 1)
-        with django_threaded_sync_to_async.reentrant_patch(o, "x", 2):
+        with django_threaded_sync_to_async.patch.reentrant(o, "x", 2):
             with self.subTest(inside=["c1"]):
                 self.assertEqual(o.x, 2)
-            with django_threaded_sync_to_async.reentrant_patch(o, "x", 3):
+            with django_threaded_sync_to_async.patch.reentrant(o, "x", 3):
                 with self.subTest(inside=["c1", "c2"]):
                     self.assertEqual(o.x, 3)
             with self.subTest(inside=["c1"]):
@@ -34,8 +34,8 @@ class Dummy:
         o = Dummy()
         fields = dir(o)
 
-        c1 = django_threaded_sync_to_async.reentrant_patch(o, "x", 2)
-        c2 = django_threaded_sync_to_async.reentrant_patch(o, "x", 3)
+        c1 = django_threaded_sync_to_async.patch.reentrant(o, "x", 2)
+        c2 = django_threaded_sync_to_async.patch.reentrant(o, "x", 3)
 
         with self.subTest(inside=[]):
             self.assertEqual(o.x, 1)

tests/test_shared_executor.py

@@ -1,3 +1,7 @@
+import asyncio
+import contextlib
+import functools
+import threading
 import unittest
 
 import asgiref.sync
@@ -7,18 +11,76 @@
 
 class TestSharedExecutor(unittest.IsolatedAsyncioTestCase):
     async def testSimple(self):
-        async with django_threaded_sync_to_async.SharedExecutor("common") as executor:
+        async with django_threaded_sync_to_async.SharedExecutor("simple_common") as executor:
             pass
 
         with self.subTest(same_name=True):
-            async with django_threaded_sync_to_async.SharedExecutor("common") as another_executor:
+            async with django_threaded_sync_to_async.SharedExecutor("simple_common") as another_executor:
                 self.assertIs(executor, another_executor)
 
         with self.subTest(same_name=False):
-            async with django_threaded_sync_to_async.SharedExecutor("specific") as specific_executor:
+            async with django_threaded_sync_to_async.SharedExecutor("simple_specific") as specific_executor:
                 self.assertIsNot(executor, specific_executor)
 
     async def testMaxTasks(self):
-        async with django_threaded_sync_to_async.SharedExecutor("common", max_tasks=2, max_workers=3) as executor:
-            # TODO Make proper test on `max_tasks`.
-            self.assertEqual(await asgiref.sync.sync_to_async(lambda: 42)(), 42)
+        workers = 10
+        timeout = 0.05
+
+        def long_call(cv, threads, another_threads):
+            threads.add(threading.current_thread().name)
+
+            def notify_cv(predicate):
+                with cv:
+                    if predicate():
+                        cv.notify_all()
+
+            with cv:
+                notify = threading.Thread(target=notify_cv, args=(lambda: len(threads) == workers,))
+                notify.start()
+                cv.wait_for(lambda: len(threads) == workers, timeout)
+
+            result = len(threads)
+            another_threads.add(threading.current_thread().name)
+
+            with cv:
+                notify = threading.Thread(target=notify_cv, args=(lambda: len(another_threads) == workers,))
+                notify.start()
+                cv.wait_for(lambda: len(another_threads) == workers, timeout)
+
+            threads.discard(threading.current_thread().name)
+            another_threads.discard(threading.current_thread().name)
+            return result
+
+        @asgiref.sync.sync_to_async
+        def decorated_long_call(*args):
+            return long_call(*args)
+
+        @contextlib.asynccontextmanager
+        async def empty(name, **kwargs):
+            yield
+
+        for parallel, context in ((False, empty), (True, django_threaded_sync_to_async.SharedExecutor)):
+            for decorated, function in ((False, asgiref.sync.sync_to_async(long_call)), (True, decorated_long_call)):
+                for tasks in (workers, workers - 1):
+                    # One or two passes are allowed — `tasks` must not be less than `workers/2`.
+                    with self.subTest(parallel=parallel, decorated=decorated, tasks=tasks):
+                        cv = threading.Condition()
+                        threads = set()
+                        another_threads = set()
+
+                        async with context(
+                            f"max_tasks_{parallel}_{decorated}_{tasks}", max_workers=workers, max_tasks=tasks
+                        ):
+                            tt = [asyncio.create_task(function(cv, threads, another_threads)) for _ in range(workers)]
+                            try:
+                                total_timeout = timeout * (2.5 if tasks == workers else 4.5)
+                                for c in asyncio.as_completed(tt, timeout=total_timeout):
+                                    self.assertIn(await c, (tasks, workers - tasks) if parallel else (1,))
+                            except asyncio.TimeoutError:
+                                for t in tt:
+                                    t.cancel()
+                                if parallel:
+                                    self.assertEqual("Exception", "occurred")
+                            else:
+                                if not parallel:
+                                    self.assertEqual("No", "exception")