From 71ba4407617c29ed35ace3e3b83f3256f02808bd Mon Sep 17 00:00:00 2001
From: Carlton Gibson <carlton.gibson@noumenal.es>
Date: Wed, 6 Jul 2022 12:37:26 +0200
Subject: [PATCH] Added support for ASGI_THREADS max worker limit. (#422)

Closes #319
---
 CHANGELOG.txt     |  6 ++++++
 daphne/server.py  |  8 ++++++++
 tests/test_cli.py | 12 +++++++++++-
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.txt b/CHANGELOG.txt
index 31e4650..8eb5463 100644
--- a/CHANGELOG.txt
+++ b/CHANGELOG.txt
@@ -20,6 +20,12 @@ Unreleased
 
 * Added ``--log-fmt`` CLI argument.
 
+* Added support for ``ASGI_THREADS`` environment variable, setting the maximum
+  number of workers used by a ``SyncToAsync`` thread-pool executor.
+
+  Set e.g. ``ASGI_THREADS=4 daphne ...`` when running to limit the number of
+  workers.
+
 3.0.2 (2021-04-07)
 ------------------
 
diff --git a/daphne/server.py b/daphne/server.py
index 4334217..e5728b3 100755
--- a/daphne/server.py
+++ b/daphne/server.py
@@ -1,10 +1,18 @@
 # This has to be done first as Twisted is import-order-sensitive with reactors
 import asyncio  # isort:skip
+import os  # isort:skip
 import sys  # isort:skip
 import warnings  # isort:skip
+from concurrent.futures import ThreadPoolExecutor  # isort:skip
 from twisted.internet import asyncioreactor  # isort:skip
 
+
 twisted_loop = asyncio.new_event_loop()
+if "ASGI_THREADS" in os.environ:
+    twisted_loop.set_default_executor(
+        ThreadPoolExecutor(max_workers=int(os.environ["ASGI_THREADS"]))
+    )
+
 current_reactor = sys.modules.get("twisted.internet.reactor", None)
 if current_reactor is not None:
     if not isinstance(current_reactor, asyncioreactor.AsyncioSelectorReactor):
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 51eab2e..8368488 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,6 +1,7 @@
 import logging
+import os
 from argparse import ArgumentError
-from unittest import TestCase
+from unittest import TestCase, skipUnless
 
 from daphne.cli import CommandLineInterface
 from daphne.endpoints import build_endpoint_description_strings as build
@@ -255,3 +256,12 @@ class TestCLIInterface(TestCase):
         Passing `--no-server-name` will set server name to '' (empty string)
         """
         self.assertCLI(["--no-server-name"], {"server_name": ""})
+
+
+@skipUnless(os.getenv("ASGI_THREADS"), "ASGI_THREADS environment variable not set.")
+class TestASGIThreads(TestCase):
+    def test_default_executor(self):
+        from daphne.server import twisted_loop
+
+        executor = twisted_loop._default_executor
+        self.assertEqual(executor._max_workers, int(os.getenv("ASGI_THREADS")))