summaryrefslogtreecommitdiffstats
path: root/stator
diff options
context:
space:
mode:
Diffstat (limited to 'stator')
-rw-r--r--stator/management/commands/runstator.py11
-rw-r--r--stator/runner.py33
2 files changed, 29 insertions, 15 deletions
diff --git a/stator/management/commands/runstator.py b/stator/management/commands/runstator.py
index 1307fef..a77192e 100644
--- a/stator/management/commands/runstator.py
+++ b/stator/management/commands/runstator.py
@@ -12,9 +12,16 @@ class Command(BaseCommand):
help = "Runs a Stator runner for a short period"
def add_arguments(self, parser):
+ parser.add_argument(
+ "--concurrency",
+ "-c",
+ type=int,
+ default=30,
+ help="How many tasks to run at once",
+ )
parser.add_argument("model_labels", nargs="*", type=str)
- def handle(self, model_labels: List[str], *args, **options):
+ def handle(self, model_labels: List[str], concurrency: int, *args, **options):
# Resolve the models list into names
models = cast(
List[Type[StatorModel]],
@@ -24,5 +31,5 @@ class Command(BaseCommand):
models = StatorModel.subclasses
print("Running for models: " + " ".join(m._meta.label_lower for m in models))
# Run a runner
- runner = StatorRunner(models)
+ runner = StatorRunner(models, concurrency=concurrency)
async_to_sync(runner.run)()
diff --git a/stator/runner.py b/stator/runner.py
index 0b42b27..187aa47 100644
--- a/stator/runner.py
+++ b/stator/runner.py
@@ -16,16 +16,20 @@ class StatorRunner:
Designed to run in a one-shot mode, living inside a request.
"""
- START_TIMEOUT = 30
- TOTAL_TIMEOUT = 60
- LOCK_TIMEOUT = 120
-
- MAX_TASKS = 30
- MAX_TASKS_PER_MODEL = 5
-
- def __init__(self, models: List[Type[StatorModel]]):
+ def __init__(
+ self,
+ models: List[Type[StatorModel]],
+ concurrency: int = 30,
+ concurrency_per_model: int = 5,
+ run_period: int = 30,
+ wait_period: int = 30,
+ ):
self.models = models
self.runner_id = uuid.uuid4().hex
+ self.concurrency = concurrency
+ self.concurrency_per_model = concurrency_per_model
+ self.run_period = run_period
+ self.total_period = run_period + wait_period
async def run(self):
start_time = time.monotonic()
@@ -40,15 +44,18 @@ class StatorRunner:
await asyncio.gather(*initial_tasks)
# For the first time period, launch tasks
print("Running main task loop")
- while (time.monotonic() - start_time) < self.START_TIMEOUT:
+ while (time.monotonic() - start_time) < self.run_period:
self.remove_completed_tasks()
- space_remaining = self.MAX_TASKS - len(self.tasks)
+ space_remaining = self.concurrency - len(self.tasks)
# Fetch new tasks
for model in self.models:
if space_remaining > 0:
for instance in await model.atransition_get_with_lock(
- min(space_remaining, self.MAX_TASKS_PER_MODEL),
- timezone.now() + datetime.timedelta(seconds=self.LOCK_TIMEOUT),
+ number=min(space_remaining, self.concurrency_per_model),
+ lock_expiry=(
+ timezone.now()
+ + datetime.timedelta(seconds=(self.total_period * 2) + 60)
+ ),
):
self.tasks.append(
asyncio.create_task(self.run_transition(instance))
@@ -59,7 +66,7 @@ class StatorRunner:
await asyncio.sleep(0.1)
# Then wait for tasks to finish
print("Waiting for tasks to complete")
- while (time.monotonic() - start_time) < self.TOTAL_TIMEOUT:
+ while (time.monotonic() - start_time) < self.total_period:
self.remove_completed_tasks()
if not self.tasks:
break