149 lines
5.2 KiB
Python
149 lines
5.2 KiB
Python
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 and 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())
|
|
if lease.pk is None:
|
|
lease.save()
|
|
else:
|
|
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
|