| import time |
| import threading |
|
|
|
|
| class TokenBucket: |
| def __init__(self, rate: float, capacity: float): |
| self.rate = float(rate) |
| self.capacity = float(capacity) |
| self._tokens = float(capacity) |
| self._updated_at = time.monotonic() |
| self._lock = threading.Lock() |
|
|
| def acquire(self, tokens: float = 1.0): |
| tokens = float(tokens) |
| while True: |
| wait_s = 0.0 |
| with self._lock: |
| now = time.monotonic() |
| elapsed = now - self._updated_at |
| if elapsed > 0: |
| self._tokens = min(self.capacity, self._tokens + elapsed * self.rate) |
| self._updated_at = now |
| if self._tokens >= tokens: |
| self._tokens -= tokens |
| return |
| if self.rate <= 0: |
| wait_s = 0.2 |
| else: |
| wait_s = (tokens - self._tokens) / self.rate |
| if wait_s <= 0: |
| wait_s = 0.01 |
| time.sleep(wait_s) |
|
|
| def set_rate_capacity(self, rate: float, capacity: float): |
| with self._lock: |
| self.rate = float(rate) |
| self.capacity = float(capacity) |
| self._tokens = min(self._tokens, self.capacity) |
|
|
|
|
| class DomainRateLimiter: |
| def __init__(self, rules: dict[str, tuple[float, float]] | None = None): |
| self._rules = rules or {} |
| self._buckets: dict[str, TokenBucket] = {} |
| self._penalties: dict[str, tuple[float, float]] = {} |
| self._lock = threading.Lock() |
|
|
| def acquire(self, domain: str, tokens: float = 1.0): |
| if not domain: |
| return |
| with self._lock: |
| now = time.monotonic() |
| if domain in self._penalties: |
| until, factor = self._penalties[domain] |
| if now >= until: |
| del self._penalties[domain] |
| if domain not in self._buckets: |
| rate, cap = self._rules.get(domain, (2.0, 2.0)) |
| self._buckets[domain] = TokenBucket(rate=rate, capacity=cap) |
| bucket = self._buckets[domain] |
| if domain in self._penalties: |
| _, factor = self._penalties[domain] |
| base_rate, base_cap = self._rules.get(domain, (2.0, 2.0)) |
| bucket.set_rate_capacity(rate=max(0.1, base_rate * factor), capacity=max(0.5, base_cap * factor)) |
| else: |
| base_rate, base_cap = self._rules.get(domain, (2.0, 2.0)) |
| bucket.set_rate_capacity(rate=base_rate, capacity=base_cap) |
| bucket.acquire(tokens=tokens) |
|
|
| def penalize(self, domain: str, factor: float = 0.25, ttl_s: float = 60.0): |
| if not domain: |
| return |
| factor = max(0.05, min(1.0, float(factor))) |
| ttl_s = max(1.0, float(ttl_s)) |
| with self._lock: |
| until = time.monotonic() + ttl_s |
| self._penalties[domain] = (until, factor) |
|
|
|
|
| def get_default_rate_limiter(): |
| rules = { |
| "edith.xiaohongshu.com": (2.0, 2.0), |
| "www.xiaohongshu.com": (1.5, 2.0), |
| "creator.xiaohongshu.com": (1.5, 2.0), |
| "ros-upload.xiaohongshu.com": (0.8, 1.0), |
| "pgy.xiaohongshu.com": (1.0, 1.0), |
| } |
| return DomainRateLimiter(rules=rules) |
|
|
