From 2ccf2e683e35c33966ccb6a1109615f72991a59b Mon Sep 17 00:00:00 2001
From: Andrew Godwin <andrew@aeracode.org>
Date: Tue, 20 Dec 2022 08:02:35 +0000
Subject: [PATCH] Allow tuning of stator concurrency

---
 docs/installation.rst |  7 ++-----
 docs/tuning.rst       | 16 ++++++++++++++++
 stator/runner.py      |  7 +++++--
 takahe/settings.py    |  6 ++++++
 4 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/docs/installation.rst b/docs/installation.rst
index 78f3797..ebb8c24 100644
--- a/docs/installation.rst
+++ b/docs/installation.rst
@@ -119,11 +119,8 @@ be provided to the containers from the first boot.
   ``["andrew@aeracode.org"]`` (if you're doing this via shell, be careful
   about escaping!)
 
-In addition, there are some optional variables you can set:
-
-* ``TAKAHE_NGINX_CACHE_SIZE`` allows you to specify the size of the disk cache
-  that is used to cache proxied avatars, profile images and media. See
-  :doc:`tuning` for more.
+There are some other, optional variables you can tweak once the
+system is up and working - see :doc:`tuning` for more.
 
 
 .. _media_configuration:
diff --git a/docs/tuning.rst b/docs/tuning.rst
index ac20597..51fe86f 100644
--- a/docs/tuning.rst
+++ b/docs/tuning.rst
@@ -43,6 +43,22 @@ problems, please get in touch with us and discuss it; Takahē is young enough
 that we need data and insight from those installations to help optimise it more.
 
 
+Stator (Task Processing)
+------------------------
+
+Takahē's background task processing system is called Stator, and it uses
+asynchronous Python to pack loads of tasks at once time into a single process.
+
+By default, it will try to run up to 100 tasks at once, with a maximum of 40
+from any single model (FanOut will usually be the one it's doing most of).
+You can tweak these with the ``TAKAHE_STATOR_CONCURRENCY`` and
+``TAKAHE_STATOR_CONCURRENCY_PER_MODEL`` environment variables.
+
+The only real limits Stator can hit are CPU and memory usage; if you see your
+Stator (worker) containers not using anywhere near all of their CPU or memory,
+you can safely increase these numbers.
+
+
 Federation
 ----------
 
diff --git a/stator/runner.py b/stator/runner.py
index ad3f660..9d0f6aa 100644
--- a/stator/runner.py
+++ b/stator/runner.py
@@ -5,6 +5,7 @@ import traceback
 import uuid
 
 from asgiref.sync import async_to_sync, sync_to_async
+from django.conf import settings
 from django.utils import timezone
 
 from core import exceptions, sentry
@@ -21,8 +22,10 @@ class StatorRunner:
     def __init__(
         self,
         models: list[type[StatorModel]],
-        concurrency: int = 50,
-        concurrency_per_model: int = 10,
+        concurrency: int = getattr(settings, "STATOR_CONCURRENCY", 50),
+        concurrency_per_model: int = getattr(
+            settings, "STATOR_CONCURRENCY_PER_MODEL", 20
+        ),
         liveness_file: str | None = None,
         schedule_interval: int = 30,
         lock_expiry: int = 300,
diff --git a/takahe/settings.py b/takahe/settings.py
index ef28512..2789066 100644
--- a/takahe/settings.py
+++ b/takahe/settings.py
@@ -127,6 +127,10 @@ class Settings(BaseSettings):
     #: Default cache backend
     CACHES_DEFAULT: CacheBackendUrl | None = None
 
+    # Stator tuning
+    STATOR_CONCURRENCY: int = 100
+    STATOR_CONCURRENCY_PER_MODEL: int = 40
+
     PGHOST: str | None = None
     PGPORT: int | None = 5432
     PGNAME: str = "takahe"
@@ -291,6 +295,8 @@ ALLOWED_HOSTS = SETUP.ALLOWED_HOSTS
 AUTO_ADMIN_EMAIL = SETUP.AUTO_ADMIN_EMAIL
 
 STATOR_TOKEN = SETUP.STATOR_TOKEN
+STATOR_CONCURRENCY = SETUP.STATOR_CONCURRENCY
+STATOR_CONCURRENCY_PER_MODEL = SETUP.STATOR_CONCURRENCY_PER_MODEL
 
 CORS_ORIGIN_ALLOW_ALL = True  # Temporary
 CORS_ORIGIN_WHITELIST = SETUP.CORS_HOSTS