diff --git a/src/s5p/pool.py b/src/s5p/pool.py index adb554f..f1d25d4 100644 --- a/src/s5p/pool.py +++ b/src/s5p/pool.py @@ -88,12 +88,21 @@ class ProxyPool: self._save_state() async def get(self) -> ChainHop | None: - """Return a random alive proxy, or None if pool is empty.""" + """Return a weighted-random alive proxy, or None if pool is empty.""" if not self._alive_keys: return None - key = random.choice(self._alive_keys) - entry = self._proxies.get(key) - return entry.hop if entry else None + now = time.time() + weights = [self._weight(self._proxies[k], now) for k in self._alive_keys] + key = random.choices(self._alive_keys, weights=weights)[0] + return self._proxies[key].hop + + @staticmethod + def _weight(entry: ProxyEntry, now: float) -> float: + """Compute selection weight from recency of last success.""" + if entry.last_ok <= 0: + return 0.01 + age = now - entry.last_ok + return 1.0 / (1.0 + age / 300.0) @property def count(self) -> int: diff --git a/tests/test_pool.py b/tests/test_pool.py index eee9b5c..131ed59 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -72,6 +72,51 @@ class TestProxyPoolGet: assert result.host == "1.2.3.4" +class TestProxyPoolWeight: + """Test weighted proxy selection.""" + + def test_recent_proxy_preferred(self): + import asyncio + from collections import Counter + + cfg = ProxyPoolConfig(sources=[]) + pool = ProxyPool(cfg, [], timeout=10.0) + + now = time.time() + fresh = ChainHop(proto="socks5", host="10.0.0.1", port=1080) + stale = ChainHop(proto="socks5", host="10.0.0.2", port=1080) + pool._proxies["socks5://10.0.0.1:1080"] = ProxyEntry( + hop=fresh, alive=True, last_ok=now, + ) + pool._proxies["socks5://10.0.0.2:1080"] = ProxyEntry( + hop=stale, alive=True, last_ok=now - 3600, + ) + pool._rebuild_alive() + + counts: Counter[str] = Counter() + for _ in range(1000): + hop = asyncio.run(pool.get()) + counts[hop.host] += 1 + + assert counts["10.0.0.1"] > counts["10.0.0.2"] * 3 + + def test_weight_values(self): + hop = ChainHop(proto="socks5", host="1.2.3.4", port=1080) + now = 1000.0 + + # just tested + entry = ProxyEntry(hop=hop, last_ok=now) + assert ProxyPool._weight(entry, now) == pytest.approx(1.0) + + # 5 minutes ago + entry.last_ok = now - 300 + assert ProxyPool._weight(entry, now) == pytest.approx(0.5) + + # never tested + entry.last_ok = 0 + assert ProxyPool._weight(entry, now) == pytest.approx(0.01) + + class TestProxyPoolFetchFile: """Test file source parsing."""