import asyncio import datetime import time import uuid from typing import List, Type from asgiref.sync import sync_to_async from django.db import transaction from django.utils import timezone from stator.models import StatorModel, StatorTask class StatorRunner: """ Runs tasks on models that are looking for state changes. Designed to run in a one-shot mode, living inside a request. """ START_TIMEOUT = 30 TOTAL_TIMEOUT = 60 LOCK_TIMEOUT = 120 MAX_TASKS = 30 def __init__(self, models: List[Type[StatorModel]]): self.models = models self.runner_id = uuid.uuid4().hex async def run(self): start_time = time.monotonic() self.handled = 0 self.tasks = [] # Clean up old locks await StatorTask.aclean_old_locks() # Examine what needs scheduling # For the first time period, launch tasks while (time.monotonic() - start_time) < self.START_TIMEOUT: self.remove_completed_tasks() space_remaining = self.MAX_TASKS - len(self.tasks) # Fetch new tasks if space_remaining > 0: for new_task in await StatorTask.aget_for_execution( space_remaining, timezone.now() + datetime.timedelta(seconds=self.LOCK_TIMEOUT), ): self.tasks.append(asyncio.create_task(self.run_task(new_task))) self.handled += 1 # Prevent busylooping await asyncio.sleep(0.01) # Then wait for tasks to finish while (time.monotonic() - start_time) < self.TOTAL_TIMEOUT: self.remove_completed_tasks() if not self.tasks: break # Prevent busylooping await asyncio.sleep(1) return self.handled async def run_task(self, task: StatorTask): # Resolve the model instance model_instance = await task.aget_model_instance() await model_instance.attempt_transition() # Remove ourselves from the database as complete await task.adelete() def remove_completed_tasks(self): self.tasks = [t for t in self.tasks if not t.done()]