karldns.de/router_dyndns/store/rate_limit.py

from __future__ import annotations

import time
from datetime import UTC, datetime, timedelta

from .types import CleanupRunRow

RATE_LIMIT_FLUSH_SECONDS = 10.0


class RateLimitCleanupMixin:
    def list_cleanup_runs(self, limit: int = 20, offset: int = 0) -> list[CleanupRunRow]:
        with self.connect() as conn:
            rows = conn.execute(
                """
                SELECT created_at, domain_challenges_deleted, unused_accounts_deleted
                FROM cleanup_runs
                ORDER BY created_at DESC
                LIMIT ? OFFSET ?
                """,
                (limit, offset),
            ).fetchall()
        return [CleanupRunRow(**dict(row)) for row in rows]

    def count_cleanup_runs(self) -> int:
        with self.connect() as conn:
            row = conn.execute("SELECT COUNT(*) FROM cleanup_runs").fetchone()
        return int(row[0])

    def allow_rate_limit(self, key: str, limit: int) -> bool:
        if limit <= 0:
            return False

        now = time.time()
        window = int(now) // 60
        refill_rate = limit / 60.0
        with self._rate_limit_lock:
            bucket = self._rate_limit_buckets.get(key)
            if bucket is None:
                bucket = {"tokens": float(limit), "updated_at": now, "window": window}
            else:
                elapsed_seconds = max(0.0, now - float(bucket["updated_at"]))
                bucket["tokens"] = min(float(limit), float(bucket["tokens"]) + elapsed_seconds * refill_rate)
                bucket["updated_at"] = now
                bucket["window"] = window

            if float(bucket["tokens"]) < 1.0:
                self._rate_limit_buckets[key] = bucket
                return False

            bucket["tokens"] = float(bucket["tokens"]) - 1.0
            self._rate_limit_buckets[key] = bucket
            dirty_key = (key, window)
            self._rate_limit_dirty_counts[dirty_key] = self._rate_limit_dirty_counts.get(dirty_key, 0) + 1
            self._prune_memory_rate_limits(now)

        self._flush_rate_limits_if_due()
        return True

    def flush_rate_limits(self) -> None:
        with self._rate_limit_lock:
            dirty_counts = self._rate_limit_dirty_counts
            self._rate_limit_dirty_counts = {}
            self._last_rate_limit_flush = time.monotonic()

        if not dirty_counts:
            return

        try:
            with self.connect() as conn:
                conn.executemany(
                    """
                    INSERT INTO rate_limits(key, window, count)
                    VALUES (?, ?, ?)
                    ON CONFLICT(key, window) DO UPDATE SET count=count + excluded.count
                    """,
                    [(key, window, count) for (key, window), count in dirty_counts.items()],
                )
        except Exception:
            with self._rate_limit_lock:
                for dirty_key, count in dirty_counts.items():
                    self._rate_limit_dirty_counts[dirty_key] = self._rate_limit_dirty_counts.get(dirty_key, 0) + count
            raise

    def _flush_rate_limits_if_due(self) -> None:
        if time.monotonic() - self._last_rate_limit_flush < RATE_LIMIT_FLUSH_SECONDS:
            return
        self.flush_rate_limits()

    def _prune_memory_rate_limits(self, now: float) -> None:
        stale_before = now - 120
        stale_keys = [
            key
            for key, bucket in self._rate_limit_buckets.items()
            if float(bucket["updated_at"]) < stale_before and float(bucket["tokens"]) >= 1.0
        ]
        for key in stale_keys:
            del self._rate_limit_buckets[key]

    def prune_rate_limit_storage(self) -> None:
        current_window = int(time.time()) // 60
        should_prune = self._should_prune_rate_limits(current_window)
        if not should_prune:
            return
        with self.connect() as conn:
            conn.execute("DELETE FROM rate_limits WHERE window < ?", (current_window - 2,))

    def _should_prune_rate_limits(self, window: int) -> bool:
        if self._last_rate_limit_prune_window == window:
            return False
        with self._rate_limit_prune_lock:
            if self._last_rate_limit_prune_window == window:
                return False
            self._last_rate_limit_prune_window = window
            return True

    def cleanup(
        self,
        challenge_hours: int,
        unused_account_hours: int,
        event_history_hours: int = 24 * 180,
    ) -> dict[str, int]:
        self.flush_rate_limits()
        challenge_cutoff = (datetime.now(UTC) - timedelta(hours=challenge_hours)).isoformat()
        unused_account_cutoff = (datetime.now(UTC) - timedelta(hours=unused_account_hours)).isoformat()
        event_history_cutoff = (datetime.now(UTC) - timedelta(hours=event_history_hours)).isoformat()
        with self.connect() as conn:
            challenge_cursor = conn.execute(
                "DELETE FROM domain_challenges WHERE verified_at IS NULL AND created_at < ?",
                (challenge_cutoff,),
            )
            account_cursor = conn.execute(
                """
                DELETE FROM accounts
                WHERE created_at < ?
                  AND NOT EXISTS (
                    SELECT 1 FROM records WHERE records.hostname = accounts.hostname
                  )
                """,
                (unused_account_cutoff,),
            )
            rollup_cursor = conn.execute(
                "DELETE FROM update_event_rollups WHERE last_seen_at < ?",
                (unused_account_cutoff,),
            )
            event_cursor = conn.execute(
                "DELETE FROM update_events WHERE created_at < ?",
                (event_history_cutoff,),
            )
            current_window = int(time.time()) // 60
            conn.execute("DELETE FROM rate_limits WHERE window < ?", (current_window - 60,))
            conn.execute(
                """
                INSERT INTO cleanup_runs(created_at, domain_challenges_deleted, unused_accounts_deleted)
                VALUES (?, ?, ?)
                """,
                (datetime.now(UTC).isoformat(), challenge_cursor.rowcount, account_cursor.rowcount),
            )
        return {
            "domain_challenges": challenge_cursor.rowcount,
            "unused_accounts": account_cursor.rowcount,
            "event_rollups": rollup_cursor.rowcount,
            "update_events": event_cursor.rowcount,
        }