The glorious return of the in-view runner

2 files changed, 40 insertions, 9 deletions

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}")