74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
"""Smooth weighted round-robin scheduler (nginx algorithm).
|
|
|
|
Each backend has a static weight. Selection is deterministic and smooth:
|
|
no backend is starved; over a full cycle each backend is picked `weight`
|
|
times, and consecutive picks of the same backend are spaced as evenly as
|
|
possible.
|
|
|
|
Correctness notes:
|
|
- Zero/negative weights are rejected (use weight 0 to disable? we reject <=0
|
|
to keep invariants; callers filter disabled backends).
|
|
- Empty pool raises.
|
|
- Thread-safe via a lock; the algorithm mutates current_weight in place.
|
|
- `peak`/`reset` helpers for introspection and reuse.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Generic, Iterable, List, Optional, TypeVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
@dataclass
|
|
class _Node(Generic[T]):
|
|
item: T
|
|
weight: int
|
|
current: int = field(default=0, init=False)
|
|
|
|
|
|
class SmoothWeightedRoundRobin(Generic[T]):
|
|
def __init__(self, items: Iterable[tuple[T, int]]):
|
|
nodes: List[_Node[T]] = []
|
|
total = 0
|
|
for item, w in items:
|
|
if not isinstance(w, int) or isinstance(w, bool):
|
|
raise TypeError("weight must be int")
|
|
if w <= 0:
|
|
raise ValueError("weight must be > 0")
|
|
nodes.append(_Node(item=item, weight=w))
|
|
total += w
|
|
if not nodes:
|
|
raise ValueError("pool must be non-empty")
|
|
self._nodes = nodes
|
|
self._total = total
|
|
self._lock = threading.Lock()
|
|
|
|
@property
|
|
def total_weight(self) -> int:
|
|
return self._total
|
|
|
|
def size(self) -> int:
|
|
return len(self._nodes)
|
|
|
|
def next(self) -> T:
|
|
with self._lock:
|
|
best = None
|
|
for n in self._nodes:
|
|
n.current += n.weight
|
|
if best is None or n.current > best.current:
|
|
best = n
|
|
assert best is not None
|
|
best.current -= self._total
|
|
return best.item
|
|
|
|
def peek_weights(self) -> list[int]:
|
|
with self._lock:
|
|
return [n.current for n in self._nodes]
|
|
|
|
def reset(self) -> None:
|
|
with self._lock:
|
|
for n in self._nodes:
|
|
n.current = 0
|