mirror of
https://github.com/Flowseal/tg-ws-proxy.git
synced 2026-10-21 02:58:45 +00:00
WsPool refactoring, TCP fallback backoff, removed direct connections outside of a pool
This commit is contained in:
Flowseal committed
1 parent
48954ff9ab
commit
50c980e9ea
13 files changed
+563
-294
No files matched your search
+99
-86
@@ -10,12 +10,11 @@ from typing import Dict, List, Optional, Tuple, Set
|
||||
from .raw_websocket import RawWebSocket, WsHandshakeError
|
||||
from .stats import stats
|
||||
from .config import proxy_config
|
||||
from .utils import ws_domains, DC_DEFAULT_IPS
|
||||
from .utils import ws_domains, DC_DEFAULT_IPS, WS_PATH, WS_PATH_TEST
|
||||
|
||||
log = logging.getLogger('tg-mtproto-proxy')
|
||||
|
||||
# TODO: domains handling is broken: wrong is_media flag causes tcp_reset after handshake,
|
||||
# but initial connection is still established no matter what is is_media flag is set to
|
||||
|
||||
class _WsPool:
|
||||
WS_POOL_MAX_AGE = 120.0
|
||||
WS_POOL_CHECK_INTERVAL = 5.0
|
||||
@@ -23,17 +22,20 @@ class _WsPool:
|
||||
REFILL_BACKOFF_MAX = 3600.0
|
||||
|
||||
def __init__(self):
|
||||
self._idle: Dict[Tuple[int, bool], deque] = {}
|
||||
self._refilling: Set[Tuple[int, bool]] = set()
|
||||
self._rotating: Dict[Tuple[int, bool], asyncio.Task] = {}
|
||||
self._refill_failures: Dict[Tuple[int, bool], int] = {}
|
||||
self._refill_after: Dict[Tuple[int, bool], float] = {}
|
||||
self.try_fronting_first = True
|
||||
self._idle: Dict[Tuple[int, bool, bool], deque] = {}
|
||||
self._refilling: Dict[Tuple[int, bool, bool], asyncio.Task] = {}
|
||||
self._rotating: Dict[Tuple[int, bool, bool], asyncio.Task] = {}
|
||||
self._refill_failures: Dict[Tuple[int, bool, bool], int] = {}
|
||||
self._refill_after: Dict[Tuple[int, bool, bool], float] = {}
|
||||
self.try_fronting_first = False
|
||||
|
||||
async def get(self, dc: int, is_media: bool,
|
||||
target_ip: str, domains: List[str]
|
||||
async def get(self, dc: int, is_media: bool, *, is_test_dc: bool = False
|
||||
) -> Optional[RawWebSocket]:
|
||||
key = (dc, is_media)
|
||||
target_ip = proxy_config.dc_redirects.get(dc)
|
||||
if not target_ip or proxy_config.pool_size <= 0:
|
||||
return None
|
||||
key = (dc, is_media, is_test_dc)
|
||||
domains = ws_domains(dc, is_media)
|
||||
now = time.monotonic()
|
||||
|
||||
bucket = self._idle.get(key)
|
||||
@@ -43,14 +45,12 @@ class _WsPool:
|
||||
while bucket:
|
||||
ws, created = bucket.popleft()
|
||||
age = now - created
|
||||
if (age > self.WS_POOL_MAX_AGE or ws._closed
|
||||
or ws.writer.transport.is_closing()):
|
||||
if self._is_stale(ws, created, now):
|
||||
asyncio.create_task(self._quiet_close(ws))
|
||||
continue
|
||||
stats.pool_hits += 1
|
||||
log.debug("WS pool hit DC%d%s (age=%.1fs, left=%d)",
|
||||
dc, 'm' if is_media else '', age, len(bucket))
|
||||
self.report_success(dc, is_media)
|
||||
log.debug("WS pool hit DC%d%s%s (age=%.1fs, left=%d)",
|
||||
dc, 't' if is_test_dc else '', 'm' if is_media else '', age, len(bucket))
|
||||
self._schedule_refill(key, target_ip, domains)
|
||||
return ws
|
||||
|
||||
@@ -58,20 +58,25 @@ class _WsPool:
|
||||
self._schedule_refill(key, target_ip, domains)
|
||||
return None
|
||||
|
||||
def _is_stale(self, ws, created, now):
|
||||
return (now - created >= self.WS_POOL_MAX_AGE or ws._closed
|
||||
or ws.writer.transport.is_closing() or ws.reader.at_eof()
|
||||
or ws.reader.exception() is not None)
|
||||
|
||||
def _schedule_refill(self, key, target_ip, domains):
|
||||
if proxy_config.pool_size <= 0:
|
||||
return
|
||||
self._schedule_rotation(key, target_ip, domains)
|
||||
if (key in self._refilling
|
||||
or time.monotonic() < self._refill_after.get(key, 0)):
|
||||
return
|
||||
self._refilling.add(key)
|
||||
asyncio.create_task(self._refill(key, target_ip, domains))
|
||||
|
||||
def report_success(self, dc: int, is_media: bool) -> None:
|
||||
key = (dc, is_media)
|
||||
self._refill_failures.pop(key, None)
|
||||
self._refill_after.pop(key, None)
|
||||
self._refilling[key] = asyncio.create_task(
|
||||
self._refill(key, target_ip, domains))
|
||||
|
||||
async def _refill(self, key, target_ip, domains):
|
||||
dc, is_media = key
|
||||
dc, is_media, is_test_dc = key
|
||||
tasks = []
|
||||
adopted = set()
|
||||
try:
|
||||
bucket = self._idle.setdefault(key, deque())
|
||||
needed = proxy_config.pool_size - len(bucket)
|
||||
@@ -79,19 +84,23 @@ class _WsPool:
|
||||
return
|
||||
connected = 0
|
||||
tasks = [asyncio.create_task(
|
||||
self._connect_one(target_ip, domains))
|
||||
self._connect_one(target_ip, domains,
|
||||
WS_PATH_TEST if is_test_dc else WS_PATH))
|
||||
for _ in range(needed)]
|
||||
for t in tasks:
|
||||
for t in asyncio.as_completed(tasks):
|
||||
try:
|
||||
ws = await t
|
||||
if ws:
|
||||
if self._refilling.get(key) is not asyncio.current_task():
|
||||
return
|
||||
bucket.append((ws, time.monotonic()))
|
||||
adopted.add(ws)
|
||||
connected += 1
|
||||
self._schedule_rotation(key, target_ip, domains)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as exc:
|
||||
log.debug("WS pool connect failed: %r", exc)
|
||||
if connected:
|
||||
self.report_success(dc, is_media)
|
||||
self._refill_failures.pop(key, None)
|
||||
self._refill_after.pop(key, None)
|
||||
else:
|
||||
failures = self._refill_failures.get(key, 0) + 1
|
||||
self._refill_failures[key] = failures
|
||||
@@ -102,12 +111,20 @@ class _WsPool:
|
||||
)
|
||||
self._refill_after[key] = time.monotonic() + delay
|
||||
log.info(
|
||||
"WS pool refill failed for DC%d%s, retry in %.0fs",
|
||||
dc, 'm' if is_media else '', delay)
|
||||
log.debug("WS pool refilled DC%d%s: %d ready",
|
||||
dc, 'm' if is_media else '', len(bucket))
|
||||
"WS pool refill failed for DC%d%s%s, retry in %.0fs",
|
||||
dc, 't' if is_test_dc else '', 'm' if is_media else '', delay)
|
||||
log.debug("WS pool refilled DC%d%s%s: %d ready",
|
||||
dc, 't' if is_test_dc else '', 'm' if is_media else '', len(bucket))
|
||||
finally:
|
||||
self._refilling.discard(key)
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
for ws in results:
|
||||
if ws is not None and not isinstance(ws, BaseException) and ws not in adopted:
|
||||
await self._quiet_close(ws)
|
||||
if self._refilling.get(key) is asyncio.current_task():
|
||||
self._refilling.pop(key, None)
|
||||
|
||||
def _schedule_rotation(self, key, target_ip, domains):
|
||||
if key in self._rotating:
|
||||
@@ -116,28 +133,16 @@ class _WsPool:
|
||||
self._rotate(key, target_ip, domains))
|
||||
|
||||
async def _rotate(self, key, target_ip, domains):
|
||||
dc, is_media = key
|
||||
dc, is_media, is_test_dc = key
|
||||
try:
|
||||
while True:
|
||||
bucket = self._idle.get(key)
|
||||
if not bucket:
|
||||
return
|
||||
|
||||
expires_at = min(
|
||||
created + self.WS_POOL_MAX_AGE
|
||||
for _, created in bucket)
|
||||
await asyncio.sleep(min(
|
||||
self.WS_POOL_CHECK_INTERVAL,
|
||||
max(0, expires_at - time.monotonic())))
|
||||
|
||||
while proxy_config.pool_size > 0:
|
||||
bucket = self._idle.setdefault(key, deque())
|
||||
now = time.monotonic()
|
||||
expired = []
|
||||
ready = deque()
|
||||
while bucket:
|
||||
ws, created = bucket.popleft()
|
||||
if (now - created >= self.WS_POOL_MAX_AGE
|
||||
or ws._closed
|
||||
or ws.writer.transport.is_closing()):
|
||||
if self._is_stale(ws, created, now):
|
||||
expired.append(ws)
|
||||
else:
|
||||
ready.append((ws, created))
|
||||
@@ -147,49 +152,45 @@ class _WsPool:
|
||||
for ws in expired:
|
||||
asyncio.create_task(self._quiet_close(ws))
|
||||
log.debug(
|
||||
"WS pool rotated DC%d%s: %d stale, %d ready",
|
||||
dc, 'm' if is_media else '', len(expired), len(bucket))
|
||||
"WS pool rotated DC%d%s%s: %d stale, %d ready",
|
||||
dc, 't' if is_test_dc else '', 'm' if is_media else '',
|
||||
len(expired), len(bucket))
|
||||
|
||||
if len(bucket) < proxy_config.pool_size:
|
||||
self._schedule_refill(key, target_ip, domains)
|
||||
|
||||
wake_at = now + self.WS_POOL_CHECK_INTERVAL
|
||||
if bucket:
|
||||
wake_at = min(wake_at, min(
|
||||
created + self.WS_POOL_MAX_AGE for _, created in bucket))
|
||||
if key not in self._refilling and self._refill_after.get(key, 0) > now:
|
||||
wake_at = min(wake_at, self._refill_after[key])
|
||||
await asyncio.sleep(max(0, wake_at - time.monotonic()))
|
||||
finally:
|
||||
if self._rotating.get(key) is asyncio.current_task():
|
||||
self._rotating.pop(key, None)
|
||||
|
||||
async def _connect_one(self, target_ip, domains) -> Optional[RawWebSocket]:
|
||||
async def _connect_one(self, target_ip, domains, path=WS_PATH) -> Optional[RawWebSocket]:
|
||||
for domain in domains:
|
||||
if self.try_fronting_first:
|
||||
ws = await self._connect_fronted(target_ip, domain)
|
||||
if ws:
|
||||
return ws
|
||||
try:
|
||||
ws = await RawWebSocket.connect(
|
||||
target_ip, domain, timeout=8)
|
||||
self.try_fronting_first = False
|
||||
return ws
|
||||
except (asyncio.TimeoutError, ConnectionResetError):
|
||||
if self.try_fronting_first:
|
||||
return None
|
||||
return await self._connect_fronted(target_ip, domain)
|
||||
except WsHandshakeError as exc:
|
||||
if exc.is_redirect:
|
||||
modes = (True, False) if self.try_fronting_first else (False, True)
|
||||
for fronted in modes:
|
||||
try:
|
||||
ws = await RawWebSocket.connect(
|
||||
target_ip, domain, timeout=7 if fronted else 8, path=path,
|
||||
sni="sprinthost.ru" if fronted else None)
|
||||
except Exception as exc:
|
||||
stats.ws_errors += 1
|
||||
log.debug("WS pool connect %s%s via %s (fronting=%s): %r",
|
||||
domain, path, target_ip, fronted, exc)
|
||||
continue
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
self.try_fronting_first = fronted
|
||||
if fronted:
|
||||
stats.connections_fronting += 1
|
||||
log.debug("WS pool connected %s%s via %s (fronting=%s)",
|
||||
domain, path, target_ip, fronted)
|
||||
return ws
|
||||
return None
|
||||
|
||||
async def _connect_fronted(self, target_ip, domain) -> Optional[RawWebSocket]:
|
||||
try:
|
||||
ws = await RawWebSocket.connect(
|
||||
target_ip, domain, timeout=7, sni="sprinthost.ru")
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
stats.connections_fronting += 1
|
||||
self.try_fronting_first = True
|
||||
return ws
|
||||
|
||||
async def _quiet_close(self, ws):
|
||||
try:
|
||||
await ws.close()
|
||||
@@ -202,14 +203,21 @@ class _WsPool:
|
||||
continue
|
||||
for is_media in (False, True):
|
||||
domains = ws_domains(dc, is_media)
|
||||
self._schedule_refill((dc, is_media), target_ip, domains)
|
||||
key = (dc, is_media, proxy_config.force_test_dc)
|
||||
self._schedule_refill(key, target_ip, domains)
|
||||
log.info("WS pool warmup started for %d DC(s)", len(proxy_config.dc_redirects))
|
||||
|
||||
def reset(self):
|
||||
loop = asyncio.get_running_loop()
|
||||
for task in self._rotating.values():
|
||||
for task in list(self._rotating.values()) + list(self._refilling.values()):
|
||||
if not task.done() and task.get_loop() is loop:
|
||||
task.cancel()
|
||||
for bucket in self._idle.values():
|
||||
for ws, _ in bucket:
|
||||
try:
|
||||
ws.writer.close()
|
||||
except Exception as exc:
|
||||
log.debug("WS pool close failed: %r", exc)
|
||||
self._idle.clear()
|
||||
self._refilling.clear()
|
||||
self._rotating.clear()
|
||||
@@ -217,6 +225,11 @@ class _WsPool:
|
||||
self._refill_after.clear()
|
||||
self.try_fronting_first = False
|
||||
|
||||
async def close(self):
|
||||
tasks = list(self._rotating.values()) + list(self._refilling.values())
|
||||
self.reset()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
|
||||
class _CfWorkerPool:
|
||||
WS_POOL_MAX_AGE = 100.0
|
||||
|
||||
Reference in new issue
Block a user