Added SharedExecutor.

Author: excitoon

README.md

@@ -1,6 +1,10 @@
-### `django_threaded_sync_to_async`
+## `django_threaded_sync_to_async`
 
-Under executor context, replaces `sync_to_async` calls to `sync_to_async(thread_sensitive=None, executor=...)`, effectively allowing Django to make calls to database concurrently:
+FIXME add description
+
+### `Executor`
+
+Under executor context, `Executor` replaces `sync_to_async` calls to `sync_to_async(thread_sensitive=None, executor=...)`, effectively allowing Django to make calls to database concurrently:
 
 ```python3
 async with django_threaded_sync_to_async.Executor(thread_name_prefix="thread", max_workers=3) as executor:
@@ -10,3 +14,17 @@ async with django_threaded_sync_to_async.Executor(thread_name_prefix="thread", m
     d = asgiref.sync.sync_to_async(long_call)(4)
     await asyncio.gather(a, b, c, d)
 ```
+
+### `SharedExecutor`
+
+Maintains global dictionary of executors (`concurrent.futures.ThreadPoolExecutor`) accessed by name and allows to limit utilization of executor for a single context.
+
+```python3
+@django_threaded_sync_to_async.SharedExecutor("common", max_workers=3, max_tasks=2)
+def operations():
+    a = asgiref.sync.sync_to_async(long_call)(1)
+    b = asgiref.sync.sync_to_async(long_call)(2)
+    c = asgiref.sync.sync_to_async(long_call)(3)
+    d = asgiref.sync.sync_to_async(long_call)(4)
+    await asyncio.gather(a, b, c, d)
+```

django_threaded_sync_to_async/__init__.py

@@ -68,37 +68,66 @@ def one_time_patch(obj, attr, value):
             delattr(obj, f"__{attr}__count__")
 
 
+_current_executor = contextvars.ContextVar("current_executor", default=None)
+_max_tasks_semaphore = contextvars.ContextVar("max_tasks_semaphore", default=None)
+
+
 async def _sync_to_async_call(self, orig, *args, **kwargs):
-    if (executor := _get_current_executor()) is not None:
+    if (executor := _current_executor.get()) is not None:
         self = asgiref.sync.SyncToAsync(self.func, thread_sensitive=False, executor=executor)
 
     else:
         """
         The task is called outside of executor's scope (or in different context).
         """
 
-    return await orig(self, *args, **kwargs)
+    if _max_tasks_semaphore.get() is not None:
+        with _max_tasks_semaphore.get():
+            return await orig(self, *args, **kwargs)
 
-
-_current_executor = contextvars.ContextVar("current_executor", default=None)
+    else:
+        return await orig(self, *args, **kwargs)
 
 
 @contextlib.contextmanager
-def _set_current_executor(value):
-    token = _current_executor.set(value)
+def _set_context_variable(variable, value):
+    token = variable.set(value)
     yield
-    _current_executor.reset(token)
+    variable.reset(token)
 
 
-def _get_current_executor():
-    return _current_executor.get()
+@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):
+            yield executor
 
 
 @contextlib.asynccontextmanager
 async def Executor(*args, **kwargs):
     with concurrent.futures.ThreadPoolExecutor(*args, **kwargs) as executor:
-        with _set_current_executor(executor):
-            # It can be replaced by a single call to `setattr(obj, attr, value)` 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):
-                yield executor
+        with _use_executor(executor):
+            yield executor
+
+
+_shared_executors = {}
+_shared_executors_lock = threading.Lock()
+
+
+@contextlib.asynccontextmanager
+async def SharedExecutor(name, *args, max_tasks=None, **kwargs):
+    with _shared_executors_lock:
+        if name in _shared_executors:
+            executor = _shared_executors[name]
+            if "max_workers" in kwargs:
+                executor._max_workers = max(kwargs["max_workers"], executor._max_workers)
+        else:
+            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 _use_executor(executor):
+            yield executor

tests/test_shared_executor.py

@@ -0,0 +1,24 @@
+import unittest
+
+import asgiref.sync
+
+import django_threaded_sync_to_async
+
+
+class TestSharedExecutor(unittest.IsolatedAsyncioTestCase):
+    async def testSimple(self):
+        async with django_threaded_sync_to_async.SharedExecutor("common") as executor:
+            pass
+
+        with self.subTest(same_name=True):
+            async with django_threaded_sync_to_async.SharedExecutor("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:
+                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)