Initial deploy setup

apps/sources/services/scheduling.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import hashlib
+from datetime import timedelta
+from urllib.parse import urlparse
+
+from django.conf import settings
+from django.db import transaction
+from django.db.models import Max
+from django.utils import timezone
+
+from apps.sources.models import Source, SourceLease, SourceOriginState
+
+
+def origin_key(source: Source) -> str:
+    if source.domain:
+        return source.domain.lower()
+    hostname = (urlparse(source.base_url).hostname or "").lower()
+    return hostname.rstrip(".")
+
+
+def max_origin_minimum_interval_seconds(source: Source) -> int:
+    domain = origin_key(source)
+    aggregate = Source.objects.filter(domain=domain).aggregate(
+        maximum_interval=Max("minimum_interval_seconds")
+    )
+    max_interval = aggregate["maximum_interval"]
+    if max_interval is None:
+        return int(source.minimum_interval_seconds)
+    return int(max_interval)
+
+
+def _lease_ttl_seconds() -> int:
+    ttl = settings.SOURCE_LEASE_TTL_SECONDS
+    if ttl <= 0:
+        return 180
+    return int(ttl)
+
+
+def _max_origin_interval_seconds(source: Source) -> int:
+    explicit = settings.SOURCE_ORIGIN_MIN_INTERVAL_SECONDS
+    if explicit > 0:
+        return int(explicit)
+    return max_origin_minimum_interval_seconds(source)
+
+
+def calculate_jitter_seconds(source: Source, *, max_seconds: int) -> int:
+    max_seconds = int(max_seconds)
+    if max_seconds <= 0:
+        return 0
+    jitter_seed = f"{source.pk}:{source.domain}:{source.minimum_interval_seconds}"
+    digest = hashlib.sha256(jitter_seed.encode("utf-8")).digest()
+    jitter_span = max_seconds + 1
+    return int(int.from_bytes(digest[:8], "big") % jitter_span)
+
+
+def calculate_failure_backoff_seconds(
+    source: Source,
+    *,
+    failure_count: int,
+    retry_after_seconds: int | None = None,
+    timeout: bool = False,
+) -> int:
+    failure_count = max(1, failure_count)
+    base_seconds = settings.SOURCE_FAILURE_BACKOFF_BASE_SECONDS
+    if timeout:
+        base_seconds = max(base_seconds, settings.SOURCE_FAILURE_TIMEOUT_BASE_SECONDS)
+    factor = 2 ** min(failure_count - 1, 10)
+    backoff_seconds = base_seconds * factor
+    if retry_after_seconds:
+        backoff_seconds = max(backoff_seconds, retry_after_seconds)
+    max_seconds = settings.SOURCE_FAILURE_BACKOFF_MAX_SECONDS
+    return min(max_seconds, backoff_seconds) + calculate_jitter_seconds(
+        source, max_seconds=settings.SOURCE_FAILURE_JITTER_SECONDS
+    )
+
+
+def calculate_success_jitter_seconds(source: Source) -> int:
+    return calculate_jitter_seconds(
+        source, max_seconds=settings.SOURCE_SUCCESS_JITTER_SECONDS
+    )
+
+
+def acquire_source_lease(
+    *, source_id: int, worker_token: str, now=None
+) -> SourceLease | None:
+    now = now or timezone.now()
+    with transaction.atomic():
+        source = Source.objects.select_for_update().get(pk=source_id)
+        if not source.domain:
+            return None
+
+        domain = origin_key(source)
+        origin_state = SourceOriginState.objects.select_for_update().get_or_create(
+            domain=domain, defaults={"next_allowed_at": now - timedelta(seconds=1)}
+        )[0]
+
+        if origin_state.next_allowed_at and origin_state.next_allowed_at > now:
+            existing_lease = SourceLease.objects.select_for_update().filter(source=source).first()
+            if not existing_lease or existing_lease.token != worker_token:
+                return None
+
+        lease = SourceLease.objects.select_for_update().filter(source=source).first()
+        if lease is not None and not lease.is_expired:
+            if lease.token != worker_token:
+                return None
+
+        active_leases = SourceLease.objects.select_for_update().filter(
+            source__domain=domain, expires_at__gt=now
+        )
+        active_count = active_leases.exclude(source=source).count()
+        max_concurrency = max(1, int(source.max_concurrency))
+        if active_count >= max_concurrency:
+            return None
+
+        if lease is None:
+            lease = SourceLease(source=source)
+        lease.token = worker_token
+        lease.worker_id = worker_token
+        lease.expires_at = now + timedelta(seconds=_lease_ttl_seconds())
+        lease.save(update_fields=["token", "worker_id", "expires_at", "updated_at"])
+        return lease
+
+
+def release_source_lease(*, source_id: int, worker_token: str, now=None) -> bool:
+    now = now or timezone.now()
+    with transaction.atomic():
+        source = Source.objects.select_for_update().get(pk=source_id)
+        if not source.domain:
+            return False
+
+        lease = SourceLease.objects.select_for_update().filter(
+            source=source, token=worker_token
+        ).first()
+        if not lease:
+            return False
+
+        lease.expires_at = now
+        lease.save(update_fields=["expires_at", "updated_at"])
+
+        interval_seconds = _max_origin_interval_seconds(source)
+        domain = origin_key(source)
+        origin_state = SourceOriginState.objects.select_for_update().get_or_create(
+            domain=domain, defaults={"next_allowed_at": now}
+        )[0]
+        origin_state.next_allowed_at = now + timedelta(seconds=max(0, interval_seconds))
+        origin_state.save(update_fields=["next_allowed_at", "updated_at"])
+        return True