requestflow 0.1.0__py3-none-any.whl

requestflow/priorities.py

@@ -0,0 +1,49 @@
+from enum import IntEnum
+from typing import Union
+
+
+class Priority(IntEnum):
+    CRITICAL = 0
+    HIGH = 1
+    NORMAL = 2
+    LOW = 3
+
+    @classmethod
+    def from_value(cls, value: Union['Priority', str, int]) -> 'Priority':
+        if isinstance(value, Priority):
+            return value
+        if isinstance(value, int):
+            try:
+                return cls(value)
+            except ValueError:
+                raise ValueError(
+                    f'invalid priority integer: {value}. must be 0, 1, 2, or 3.')
+        if isinstance(value, str):
+            normalized = value.strip().upper()
+            try:
+                return cls[normalized]
+            except KeyError:
+                raise ValueError(
+                    f"unknown priority name: '{value}'. valid options: {[p.name.lower() for p in cls]}")
+        raise TypeError(
+            f'priority must be priority, str, or int, not {type(value).__name__}')
+
+    def __lt__(self, other: object) -> bool:
+        if isinstance(other, Priority):
+            return self.value < other.value
+        return NotImplemented
+
+    def __le__(self, other: object) -> bool:
+        if isinstance(other, Priority):
+            return self.value <= other.value
+        return NotImplemented
+
+    def __gt__(self, other: object) -> bool:
+        if isinstance(other, Priority):
+            return self.value > other.value
+        return NotImplemented
+
+    def __ge__(self, other: object) -> bool:
+        if isinstance(other, Priority):
+            return self.value >= other.value
+        return NotImplemented

requestflow/py.typed

@@ -0,0 +1 @@
+

requestflow/queue.py

@@ -0,0 +1,236 @@
+import asyncio
+import collections
+import time
+from typing import Callable, Deque, Dict, List, Optional, Set
+from requestflow.exceptions import QueueFullError
+from requestflow.models import RequestItem
+from requestflow.policies import OverloadPolicy
+from requestflow.priorities import Priority
+
+
+class RequestQueue:
+    def __init__(self, max_pending: int = 5000, overload_policy: OverloadPolicy = OverloadPolicy.REJECT, aging_threshold_seconds: float = 2.0, guild_quantum: int = 3) -> None:
+        if max_pending <= 0:
+            raise ValueError(f'max_pending must be > 0, got {max_pending}')
+        self._max_pending: int = max_pending
+        self._overload_policy: OverloadPolicy = overload_policy
+        self._aging_threshold: float = aging_threshold_seconds
+        self._guild_quantum: int = guild_quantum
+        self._size: int = 0
+        self._sequence_counter: int = 0
+        self._high_water_mark: int = 0
+        self._queues: Dict[Priority, Dict[Optional[int], Deque[RequestItem]]] = {
+            p: collections.defaultdict(collections.deque) for p in Priority}
+        self._active_guild_keys: Dict[Priority, List[Optional[int]]] = {
+            p: [] for p in Priority}
+        self._guild_indices: Dict[Priority, int] = {p: 0 for p in Priority}
+        self._guild_deficits: Dict[Optional[int],
+                                   int] = collections.defaultdict(int)
+        self._not_empty: asyncio.Event = asyncio.Event()
+        self._not_full: asyncio.Event = asyncio.Event()
+        self._not_full.set()
+        self._lock: asyncio.Lock = asyncio.Lock()
+
+    @property
+    def size(self) -> int:
+        return self._size
+
+    @property
+    def max_pending(self) -> int:
+        return self._max_pending
+
+    @property
+    def high_water_mark(self) -> int:
+        return self._high_water_mark
+
+    @property
+    def overload_policy(self) -> OverloadPolicy:
+        return self._overload_policy
+
+    def is_empty(self) -> bool:
+        return self._size == 0
+
+    def is_full(self) -> bool:
+        return self._size >= self._max_pending
+
+    def get_priority_counts(self) -> Dict[str, int]:
+        counts = {}
+        for p in Priority:
+            total = sum((len(q) for q in self._queues[p].values()))
+            counts[p.name.lower()] = total
+        return counts
+
+    def get_active_guild_count(self) -> int:
+        seen: Set[Optional[int]] = set()
+        for p in Priority:
+            for guild_id, q in self._queues[p].items():
+                if len(q) > 0:
+                    seen.add(guild_id)
+        return len(seen)
+
+    async def put(self, item: RequestItem, timeout: Optional[float] = None) -> None:
+        async with self._lock:
+            if self._size >= self._max_pending:
+                if self._overload_policy == OverloadPolicy.REJECT:
+                    raise QueueFullError(
+                        f'queue is full ({self._size}/{self._max_pending}). workload rejected.')
+                elif self._overload_policy == OverloadPolicy.DROP_OLDEST:
+                    self._drop_oldest_locked()
+                elif self._overload_policy == OverloadPolicy.WAIT:
+                    self._not_full.clear()
+        if self._overload_policy == OverloadPolicy.WAIT and self._size >= self._max_pending:
+            start_wait = time.monotonic()
+            while True:
+                remaining_timeout = None
+                if timeout is not None:
+                    elapsed = time.monotonic() - start_wait
+                    remaining_timeout = timeout - elapsed
+                    if remaining_timeout <= 0:
+                        raise QueueFullError(
+                            'timed out waiting for queue capacity.')
+                try:
+                    await asyncio.wait_for(self._not_full.wait(), timeout=remaining_timeout)
+                except asyncio.TimeoutError:
+                    raise QueueFullError(
+                        'timed out waiting for queue capacity.')
+                async with self._lock:
+                    if self._size < self._max_pending:
+                        break
+        async with self._lock:
+            self._enqueue_item_locked(item)
+
+    def _enqueue_item_locked(self, item: RequestItem) -> None:
+        self._sequence_counter += 1
+        item.sequence = self._sequence_counter
+        p = item.priority
+        g = item.guild_id
+        guild_dict = self._queues[p]
+        if g not in guild_dict or len(guild_dict[g]) == 0:
+            if g not in self._active_guild_keys[p]:
+                self._active_guild_keys[p].append(g)
+        guild_dict[g].append(item)
+        self._size += 1
+        if self._size > self._high_water_mark:
+            self._high_water_mark = self._size
+        self._not_empty.set()
+        if self._size >= self._max_pending:
+            self._not_full.clear()
+
+    def _drop_oldest_locked(self) -> None:
+        for p in reversed(Priority):
+            guild_dict = self._queues[p]
+            oldest_item: Optional[RequestItem] = None
+            target_guild: Optional[int] = None
+            for g, q in guild_dict.items():
+                if q:
+                    head = q[0]
+                    if oldest_item is None or head.sequence < oldest_item.sequence:
+                        oldest_item = head
+                        target_guild = g
+            if oldest_item is not None and target_guild is not None:
+                dropped = guild_dict[target_guild].popleft()
+                if len(guild_dict[target_guild]) == 0:
+                    if target_guild in self._active_guild_keys[p]:
+                        self._active_guild_keys[p].remove(target_guild)
+                self._size -= 1
+                dropped.cancelled = True
+                if dropped.future and (not dropped.future.done()):
+                    dropped.future.set_exception(QueueFullError(
+                        'dropped due to queue overload under drop_oldest policy'))
+                return
+
+    def remove_item(self, item: RequestItem) -> bool:
+        p = item.effective_priority
+        g = item.guild_id
+        guild_dict = self._queues[p]
+        q = guild_dict.get(g)
+        if q:
+            try:
+                q.remove(item)
+                self._size -= 1
+                if len(q) == 0 and g in self._active_guild_keys[p]:
+                    self._active_guild_keys[p].remove(g)
+                if self._size < self._max_pending:
+                    self._not_full.set()
+                if self._size == 0:
+                    self._not_empty.clear()
+                return True
+            except ValueError:
+                pass
+        return False
+
+    async def get(self, is_route_limited: Optional[Callable[[Optional[str]], bool]] = None) -> RequestItem:
+        while True:
+            await self._not_empty.wait()
+            async with self._lock:
+                item = self._extract_next_item_locked(is_route_limited)
+                if item is not None:
+                    self._size -= 1
+                    if self._size < self._max_pending:
+                        self._not_full.set()
+                    if self._size == 0:
+                        self._not_empty.clear()
+                    return item
+                self._not_empty.clear()
+            await asyncio.sleep(0.01)
+
+    def _extract_next_item_locked(self, is_route_limited: Optional[Callable[[Optional[str]], bool]] = None) -> Optional[RequestItem]:
+        now = time.monotonic()
+        self._apply_aging_locked(now)
+        for p in Priority:
+            active_guilds = self._active_guild_keys[p]
+            if not active_guilds:
+                continue
+            num_guilds = len(active_guilds)
+            start_index = self._guild_indices[p] % num_guilds
+            for i in range(num_guilds):
+                idx = (start_index + i) % num_guilds
+                g = active_guilds[idx]
+                q = self._queues[p][g]
+                if not q:
+                    continue
+                head = q[0]
+                if head.cancelled or (head.future and head.future.cancelled()):
+                    q.popleft()
+                    self._size -= 1
+                    continue
+                if is_route_limited and is_route_limited(head.route_key):
+                    continue
+                if self._guild_deficits[g] <= 0:
+                    self._guild_deficits[g] += self._guild_quantum
+                self._guild_deficits[g] -= 1
+                item = q.popleft()
+                if len(q) == 0:
+                    self._guild_deficits[g] = 0
+                    active_guilds.remove(g)
+                    if active_guilds:
+                        self._guild_indices[p] = idx % len(active_guilds)
+                    else:
+                        self._guild_indices[p] = 0
+                elif self._guild_deficits[g] <= 0:
+                    self._guild_deficits[g] = 0
+                    self._guild_indices[p] = (idx + 1) % num_guilds
+                else:
+                    self._guild_indices[p] = idx
+                return item
+        return None
+
+    def _apply_aging_locked(self, now: float) -> None:
+        if self._aging_threshold <= 0:
+            return
+        for p in [Priority.LOW, Priority.NORMAL, Priority.HIGH]:
+            guild_dict = self._queues[p]
+            higher_priority = Priority(p.value - 1)
+            for g, q in list(guild_dict.items()):
+                stale_items = []
+                for item in list(q):
+                    if now - item.enqueue_time >= self._aging_threshold:
+                        stale_items.append(item)
+                for item in stale_items:
+                    q.remove(item)
+                    item.effective_priority = higher_priority
+                    self._queues[higher_priority][g].append(item)
+                    if g not in self._active_guild_keys[higher_priority]:
+                        self._active_guild_keys[higher_priority].append(g)
+                if len(q) == 0 and g in self._active_guild_keys[p]:
+                    self._active_guild_keys[p].remove(g)

requestflow/ratelimit.py

@@ -0,0 +1,188 @@
+import asyncio
+import time
+from typing import Dict, Optional, Tuple
+
+
+class GlobalRateLimiter:
+    def __init__(self, rate: float = 40.0, ceiling: float = 50.0, burst_allowance: float = 5.0, allow_custom_ceiling: bool = False) -> None:
+        if rate <= 0:
+            raise ValueError(f'global rate limit must be positive, got {rate}')
+        if ceiling <= 0:
+            raise ValueError(
+                f'rate limit ceiling must be positive, got {ceiling}')
+        if not allow_custom_ceiling and rate > ceiling:
+            raise ValueError(
+                f'configured rate ({rate}) exceeds standard discord limit ceiling ({ceiling}). discord standard limit is 50 req/s. if you have an officially approved higher limit, set allow_custom_ceiling=true.')
+        self._rate: float = float(rate)
+        self._ceiling: float = float(ceiling)
+        self._burst_tokens: float = float(burst_allowance)
+        self._allow_custom_ceiling: bool = allow_custom_ceiling
+        self._emission_interval: float = 1.0 / self._rate
+        self._tau: float = self._burst_tokens * self._emission_interval
+        self._tat: float = 0.0
+        self._lock: asyncio.Lock = asyncio.Lock()
+        self._global_cooldown_until: float = 0.0
+        self._global_cooldown_event: asyncio.Event = asyncio.Event()
+        self._global_cooldown_event.set()
+        self._total_acquisitions: int = 0
+        self._total_wait_time: float = 0.0
+        self._cooldown_429_count: int = 0
+
+    @property
+    def rate(self) -> float:
+        return self._rate
+
+    @property
+    def ceiling(self) -> float:
+        return self._ceiling
+
+    @property
+    def is_cooling_down(self) -> bool:
+        return time.monotonic() < self._global_cooldown_until
+
+    @property
+    def cooldown_remaining(self) -> float:
+        now = time.monotonic()
+        return max(0.0, self._global_cooldown_until - now)
+
+    async def update_rate(self, new_rate: float) -> None:
+        if new_rate <= 0:
+            raise ValueError(
+                f'global rate limit must be positive, got {new_rate}')
+        if not self._allow_custom_ceiling and new_rate > self._ceiling:
+            raise ValueError(
+                f'new rate ({new_rate}) exceeds ceiling ({self._ceiling}).')
+        async with self._lock:
+            self._rate = float(new_rate)
+            self._emission_interval = 1.0 / self._rate
+            self._tau = self._burst_tokens * self._emission_interval
+            self._tat = max(self._tat, time.monotonic())
+
+    async def acquire(self, *, is_interaction: bool = False) -> float:
+        if is_interaction:
+            return 0.0
+        wait_seconds: float = 0.0
+        while True:
+            if not self._global_cooldown_event.is_set():
+                await self._global_cooldown_event.wait()
+            now = time.monotonic()
+            if now < self._global_cooldown_until:
+                sleep_duration = self._global_cooldown_until - now
+                if sleep_duration > 0:
+                    await asyncio.sleep(sleep_duration)
+                    continue
+            async with self._lock:
+                now = time.monotonic()
+                if now > self._tat:
+                    self._tat = now
+                dispatch_time = self._tat
+                self._tat = dispatch_time + self._emission_interval
+                delay = max(0.0, dispatch_time - now)
+            if delay > 0:
+                try:
+                    await asyncio.sleep(delay)
+                    wait_seconds += delay
+                except asyncio.CancelledError:
+                    async with self._lock:
+                        now_cancel = time.monotonic()
+                        if self._tat > now_cancel:
+                            self._tat = max(
+                                now_cancel, self._tat - self._emission_interval)
+                    raise
+            self._total_acquisitions += 1
+            self._total_wait_time += wait_seconds
+            return wait_seconds
+
+    def handle_global_429(self, retry_after: float) -> None:
+        now = time.monotonic()
+        cooldown_end = now + max(0.05, retry_after)
+        if cooldown_end > self._global_cooldown_until:
+            self._global_cooldown_until = cooldown_end
+            self._cooldown_429_count += 1
+            self._global_cooldown_event.clear()
+            loop = asyncio.get_event_loop()
+            loop.call_later(max(0.05, retry_after),
+                            self._clear_cooldown, cooldown_end)
+
+    def _clear_cooldown(self, target_end: float) -> None:
+        if time.monotonic() >= self._global_cooldown_until or target_end == self._global_cooldown_until:
+            self._global_cooldown_event.set()
+
+
+class RouteBucketCoordinator:
+    def __init__(self) -> None:
+        self._cooldowns: Dict[str, float] = {}
+        self._remaining: Dict[str, int] = {}
+        self._lock: asyncio.Lock = asyncio.Lock()
+        self._route_429_count: int = 0
+
+    def is_route_cooling_down(self, route_key: Optional[str]) -> bool:
+        if not route_key:
+            return False
+        expiry = self._cooldowns.get(route_key)
+        if expiry is None:
+            return False
+        if time.monotonic() >= expiry:
+            self._cooldowns.pop(route_key, None)
+            return False
+        return True
+
+    def get_route_delay(self, route_key: Optional[str]) -> float:
+        if not route_key:
+            return 0.0
+        expiry = self._cooldowns.get(route_key)
+        if expiry is None:
+            return 0.0
+        remaining = expiry - time.monotonic()
+        if remaining <= 0:
+            self._cooldowns.pop(route_key, None)
+            return 0.0
+        return remaining
+
+    def update_route_cooldown(self, route_key: str, retry_after: float) -> None:
+        if not route_key or retry_after <= 0:
+            return
+        cooldown_until = time.monotonic() + retry_after
+        existing = self._cooldowns.get(route_key, 0.0)
+        if cooldown_until > existing:
+            self._cooldowns[route_key] = cooldown_until
+            self._route_429_count += 1
+
+    def update_from_headers(self, route_key: str, headers: Dict[str, str]) -> None:
+        if not route_key:
+            return
+        remaining = headers.get('X-RateLimit-Remaining')
+        reset_after = headers.get('X-RateLimit-Reset-After')
+        if remaining is not None:
+            try:
+                rem_val = int(remaining)
+                self._remaining[route_key] = rem_val
+                if rem_val == 0 and reset_after is not None:
+                    try:
+                        reset_val = float(reset_after)
+                        self.update_route_cooldown(route_key, reset_val)
+                    except ValueError:
+                        pass
+            except ValueError:
+                pass
+
+
+class RateLimitCoordinator:
+    def __init__(self, global_rate_limit: float = 40.0, rate_limit_ceiling: float = 50.0, burst_allowance: float = 5.0, allow_custom_ceiling: bool = False) -> None:
+        self.global_limiter = GlobalRateLimiter(
+            rate=global_rate_limit, ceiling=rate_limit_ceiling, burst_allowance=burst_allowance, allow_custom_ceiling=allow_custom_ceiling)
+        self.route_coordinator = RouteBucketCoordinator()
+
+    async def acquire(self, route_key: Optional[str] = None, *, is_interaction: bool = False) -> float:
+        if route_key and self.route_coordinator.is_route_cooling_down(route_key):
+            delay = self.route_coordinator.get_route_delay(route_key)
+            if delay > 0:
+                await asyncio.sleep(delay)
+        return await self.global_limiter.acquire(is_interaction=is_interaction)
+
+    def handle_429(self, retry_after: float, *, is_global: bool = False, route_key: Optional[str] = None) -> None:
+        if is_global:
+            self.global_limiter.handle_global_429(retry_after)
+        elif route_key:
+            self.route_coordinator.update_route_cooldown(
+                route_key, retry_after)

requestflow/scheduler.py

@@ -0,0 +1,131 @@
+import asyncio
+import time
+from typing import Any, Callable, Coroutine, List, Optional, Set, TypeVar
+from requestflow.exceptions import OperationCancelledError, QueueFullError, SchedulerClosedError, SchedulerNotStartedError
+from requestflow.executor import RequestExecutor
+from requestflow.metrics import RequestMetrics
+from requestflow.models import RequestItem
+from requestflow.policies import OverloadPolicy, RetryPolicy
+from requestflow.priorities import Priority
+from requestflow.queue import RequestQueue
+from requestflow.ratelimit import RateLimitCoordinator
+T = TypeVar('t')
+
+
+class RequestScheduler:
+    def __init__(self, max_pending: int = 5000, max_concurrency: int = 16, global_rate_limit: float = 40.0, rate_limit_ceiling: float = 50.0, burst_allowance: float = 5.0, allow_custom_ceiling: bool = False, overload_policy: OverloadPolicy = OverloadPolicy.REJECT, aging_threshold_seconds: float = 2.0, guild_quantum: int = 3, default_timeout: Optional[float] = None, default_retry_policy: Optional[RetryPolicy] = None) -> None:
+        self.queue = RequestQueue(max_pending=max_pending, overload_policy=overload_policy,
+                                  aging_threshold_seconds=aging_threshold_seconds, guild_quantum=guild_quantum)
+        self.executor = RequestExecutor(
+            max_concurrency=max_concurrency, default_timeout=default_timeout, default_retry_policy=default_retry_policy)
+        self.rate_coordinator = RateLimitCoordinator(global_rate_limit=global_rate_limit, rate_limit_ceiling=rate_limit_ceiling,
+                                                     burst_allowance=burst_allowance, allow_custom_ceiling=allow_custom_ceiling)
+        self.metrics = RequestMetrics()
+        self.metrics.register_providers(get_currently_queued=lambda: self.queue.size, get_currently_executing=lambda: self.executor.currently_executing,
+                                        get_high_water_mark=lambda: self.queue.high_water_mark, get_per_priority=self.queue.get_priority_counts, get_active_guilds=self.queue.get_active_guild_count)
+        self._is_started: bool = False
+        self._is_closed: bool = False
+        self._accepting_requests: bool = False
+        self._dispatch_tasks: List[asyncio.Task] = []
+        self._in_flight_tasks: Set[asyncio.Task] = set()
+        self._drain_event: asyncio.Event = asyncio.Event()
+
+    @property
+    def is_started(self) -> bool:
+        return self._is_started
+
+    @property
+    def is_closed(self) -> bool:
+        return self._is_closed
+
+    async def start(self) -> None:
+        if self._is_closed:
+            raise SchedulerClosedError('cannot start a closed scheduler.')
+        if self._is_started:
+            return
+        self._is_started = True
+        self._accepting_requests = True
+        self._dispatch_tasks = [asyncio.create_task(
+            self._dispatch_loop(), name='requestflow_dispatcher')]
+
+    async def close(self, drain: bool = True, timeout: float = 5.0) -> None:
+        if self._is_closed:
+            return
+        self._accepting_requests = False
+        if drain and self._is_started and (self.queue.size > 0 or len(self._in_flight_tasks) > 0):
+            start_drain = time.monotonic()
+            while time.monotonic() - start_drain < timeout:
+                if self.queue.size == 0 and len(self._in_flight_tasks) == 0:
+                    break
+                await asyncio.sleep(0.01)
+        self._is_closed = True
+        for t in self._dispatch_tasks:
+            t.cancel()
+        for t in self._dispatch_tasks:
+            try:
+                await t
+            except (asyncio.CancelledError, Exception):
+                pass
+        self._dispatch_tasks.clear()
+        for t in list(self._in_flight_tasks):
+            t.cancel()
+        self._in_flight_tasks.clear()
+        while not self.queue.is_empty():
+            try:
+                item = await self.queue.get()
+                if item.future and (not item.future.done()):
+                    item.future.set_exception(SchedulerClosedError(
+                        'scheduler closed before item could be executed.'))
+            except Exception:
+                break
+
+    async def submit(self, operation: Callable[..., Coroutine[Any, Any, T]], *args: Any, priority: Priority = Priority.NORMAL, guild_id: Optional[int] = None, route_key: Optional[str] = None, is_interaction: bool = False, dedup_key: Optional[str] = None, timeout: Optional[float] = None, enqueue_timeout: Optional[float] = None, **kwargs: Any) -> T:
+        if not self._is_started:
+            raise SchedulerNotStartedError(
+                "scheduler has not been started. call 'await flow.start()'.")
+        if not self._accepting_requests or self._is_closed:
+            raise SchedulerClosedError(
+                'scheduler is closed. cannot accept new submissions.')
+        loop = asyncio.get_running_loop()
+        future = loop.create_future()
+        item = RequestItem(sequence=0, operation=operation, args=args, kwargs=kwargs, priority=priority, effective_priority=priority, guild_id=guild_id,
+                           route_key=route_key, is_interaction=is_interaction, dedup_key=dedup_key, timeout=timeout, future=future, enqueue_time=time.monotonic())
+        self.metrics.record_submit()
+        try:
+            await self.queue.put(item, timeout=enqueue_timeout)
+        except QueueFullError:
+            self.metrics.record_reject()
+            raise
+        try:
+            return await future
+        except asyncio.CancelledError:
+            item.cancelled = True
+            self.metrics.record_cancel()
+            self.queue.remove_item(item)
+            raise
+
+    async def _dispatch_loop(self) -> None:
+        while self._is_started and (not self._is_closed):
+            try:
+                item = await self.queue.get(is_route_limited=self.rate_coordinator.route_coordinator.is_route_cooling_down)
+                if item.cancelled or (item.future and item.future.cancelled()):
+                    continue
+                wait_time = await self.rate_coordinator.acquire(item.route_key, is_interaction=item.is_interaction)
+                self.metrics.record_rate_limit_wait(wait_time)
+                task = asyncio.create_task(self._execute_item_task(
+                    item), name=f'requestflow_exec_{item.sequence}')
+                self._in_flight_tasks.add(task)
+                task.add_done_callback(self._in_flight_tasks.discard)
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                await asyncio.sleep(0.01)
+
+    async def _execute_item_task(self, item: RequestItem) -> None:
+        def on_complete(it: RequestItem, err: Optional[BaseException]) -> None:
+            self.metrics.record_completion(
+                queue_wait_seconds=it.queue_wait_time, exec_duration_seconds=it.execution_duration, error=err)
+        try:
+            await self.executor.execute(item, on_complete=on_complete)
+        except asyncio.CancelledError:
+            pass