From db88c733b4ea910cd37d97453d49b7b24f767dda Mon Sep 17 00:00:00 2001
From: Andrew Godwin <andrew@aeracode.org>
Date: Sat, 26 Nov 2022 11:54:14 -0700
Subject: [PATCH] The glorious return of the in-view runner

---
 stator/runner.py | 25 ++++++++++++++++---------
 stator/views.py  | 24 ++++++++++++++++++++++++
 takahe/urls.py   |  3 +++
 3 files changed, 43 insertions(+), 9 deletions(-)
 create mode 100644 stator/views.py

diff --git a/stator/runner.py b/stator/runner.py
index cb97f6e..48549bc 100644
--- a/stator/runner.py
+++ b/stator/runner.py
@@ -26,6 +26,7 @@ class StatorRunner:
         liveness_file: Optional[str] = None,
         schedule_interval: int = 30,
         lock_expiry: int = 300,
+        run_for: int = 0,
     ):
         self.models = models
         self.runner_id = uuid.uuid4().hex
@@ -34,9 +35,11 @@ class StatorRunner:
         self.liveness_file = liveness_file
         self.schedule_interval = schedule_interval
         self.lock_expiry = lock_expiry
+        self.run_for = run_for
 
     async def run(self):
         self.handled = 0
+        self.started = time.monotonic()
         self.last_clean = time.monotonic() - self.schedule_interval
         self.tasks = []
         # For the first time period, launch tasks
@@ -71,17 +74,21 @@ class StatorRunner:
                             )
                             self.handled += 1
                             space_remaining -= 1
-                # Prevent busylooping
-                await asyncio.sleep(0.1)
-        except KeyboardInterrupt:
-            # Wait for tasks to finish
-            print("Waiting for tasks to complete")
-            while True:
-                self.remove_completed_tasks()
-                if not self.tasks:
+                # Are we in limited run mode?
+                if self.run_for and (time.monotonic() - self.started) > self.run_for:
                     break
                 # Prevent busylooping
-                await asyncio.sleep(1)
+                await asyncio.sleep(0.5)
+        except KeyboardInterrupt:
+            pass
+        # Wait for tasks to finish
+        print("Waiting for tasks to complete")
+        while True:
+            self.remove_completed_tasks()
+            if not self.tasks:
+                break
+            # Prevent busylooping
+            await asyncio.sleep(0.1)
         print("Complete")
         return self.handled
 
diff --git a/stator/views.py b/stator/views.py
new file mode 100644
index 0000000..df51e3d
--- /dev/null
+++ b/stator/views.py
@@ -0,0 +1,24 @@
+from django.conf import settings
+from django.http import HttpResponse, HttpResponseForbidden
+from django.views import View
+
+from stator.models import StatorModel
+from stator.runner import StatorRunner
+
+
+class RequestRunner(View):
+    """
+    Runs a Stator runner within a HTTP request.
+    For when you're on something serverless.
+    """
+
+    async def get(self, request):
+        # Check the token, if supplied
+        if not settings.STATOR_TOKEN:
+            return HttpResponseForbidden("No token set")
+        if request.GET.get("token") != settings.STATOR_TOKEN:
+            return HttpResponseForbidden("Invalid token")
+        # Run on all models
+        runner = StatorRunner(StatorModel.subclasses, run_for=2)
+        handled = await runner.run()
+        return HttpResponse(f"Handled {handled}")
diff --git a/takahe/urls.py b/takahe/urls.py
index b94f205..98e1050 100644
--- a/takahe/urls.py
+++ b/takahe/urls.py
@@ -5,6 +5,7 @@ from django.views.static import serve
 
 from activities.views import posts, search, timelines
 from core import views as core
+from stator import views as stator
 from users.views import activitypub, admin, auth, follows, identity, settings
 
 urlpatterns = [
@@ -121,6 +122,8 @@ urlpatterns = [
     path(".well-known/nodeinfo", activitypub.NodeInfo.as_view()),
     path("nodeinfo/2.0/", activitypub.NodeInfo2.as_view()),
     path("actor/", activitypub.SystemActorView.as_view()),
+    # Stator
+    path(".stator/", stator.RequestRunner.as_view()),
     # Django admin
     path("djadmin/", djadmin.site.urls),
     # Media files