From 4f87f6e42edb696b3c3c3ac2a7868cd2f13efb57 Mon Sep 17 00:00:00 2001 From: zenfyr Date: Thu, 16 Apr 2026 16:49:19 +0700 Subject: [PATCH] fix: add backoff to websockets --- bluesky/input.py | 25 +++--------------------- mastodon/input.py | 28 ++++----------------------- misskey/input.py | 26 +++---------------------- util/util.py | 48 ++++++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 57 insertions(+), 70 deletions(-) diff --git a/bluesky/input.py b/bluesky/input.py index 3d86f71..9e71cc7 100644 --- a/bluesky/input.py +++ b/bluesky/input.py @@ -24,7 +24,7 @@ from cross.media import Blob, download_blob from cross.post import Post, PostRef from cross.service import InputService from database.connection import DatabasePool -from util.util import Result +from util.util import Result, listen_websocket @dataclass(kw_only=True) @@ -302,24 +302,5 @@ class BlueskyJetstreamInputService(BlueskyBaseInputService): url += "&wantedCollections=app.bsky.feed.repost" url += f"&wantedDids={self.did}" - while True: - try: - async with websockets.connect( - url, - ping_interval=20, - ping_timeout=10, - close_timeout=5, - ) as ws: - self.log.info("Listening to '%s'...", env.JETSTREAM_URL) - - async for msg in ws: - self.submitter(lambda: self._accept_msg(msg)) - - except websockets.ConnectionClosedError as e: - self.log.warning("Connection closed: %s", e) - except TimeoutError as e: - self.log.warning("Connection timeout: '%s'", e) - except Exception as e: - self.log.error("Unexpected error: '%s'", e, exc_info=True) - - self.log.info("Reconnecting to '%s'...", env.JETSTREAM_URL) + async for msg in listen_websocket(url, self.log): + self.submitter(lambda: self._accept_msg(msg)) diff --git a/mastodon/input.py b/mastodon/input.py index 3c4e7de..9ada172 100644 --- a/mastodon/input.py +++ b/mastodon/input.py @@ -19,6 +19,7 @@ from cross.post import Post, PostRef from cross.service import InputService from database.connection import DatabasePool from mastodon.info import MastodonService, StatusParser, validate_and_transform +from util.util import listen_websocket ALLOWED_VISIBILITY: list[str] = ["public", "unlisted"] @@ -248,28 +249,7 @@ class MastodonInputService(MastodonService, InputService): @override async def listen(self): url = f"{self.streaming_url}/api/v1/streaming?stream=user" + headers = {"Authorization": f"Bearer {self.options.token}"} - while True: - try: - async with websockets.connect( - url, - additional_headers={ - "Authorization": f"Bearer {self.options.token}" - }, - ping_interval=20, - ping_timeout=10, - close_timeout=5, - ) as ws: - self.log.info("Listening to '%s'...", self.streaming_url) - - async for msg in ws: - self.submitter(lambda: self._accept_msg(msg)) - - except websockets.ConnectionClosedError as e: - self.log.warning("Connection closed: %s", e) - except TimeoutError as e: - self.log.warning("Connection timeout: '%s'", e) - except Exception as e: - self.log.error("Unexpected error: '%s'", e, exc_info=True) - - self.log.info("Reconnecting to '%s'...", self.streaming_url) + async for msg in listen_websocket(url, self.log, headers=headers): + self.submitter(lambda: self._accept_msg(msg)) diff --git a/misskey/input.py b/misskey/input.py index 3d79806..14db900 100644 --- a/misskey/input.py +++ b/misskey/input.py @@ -20,7 +20,7 @@ from cross.service import InputService from database.connection import DatabasePool from misskey.info import MisskeyService from util.markdown import MarkdownParser -from util.util import normalize_service_url +from util.util import listen_websocket, normalize_service_url ALLOWED_VISIBILITY = ["public", "home"] @@ -242,25 +242,5 @@ class MisskeyInputService(MisskeyService, InputService): streaming: str = f"{'wss' if self.url.startswith('https') else 'ws'}://{self.url.split('://', 1)[1]}" url: str = f"{streaming}/streaming?i={self.options.token}" - while True: - try: - async with websockets.connect( - url, - ping_interval=20, - ping_timeout=10, - close_timeout=5, - ) as ws: - self.log.info("Listening to '%s'...", streaming) - await self._subscribe_to_home(ws) - - async for msg in ws: - self.submitter(lambda: self._accept_msg(msg)) - - except websockets.ConnectionClosedError as e: - self.log.warning("Connection closed: %s", e) - except TimeoutError as e: - self.log.warning("Connection timeout: '%s'", e) - except Exception as e: - self.log.error("Unexpected error: '%s'", e, exc_info=True) - - self.log.info("Reconnecting to '%s'...", streaming) + async for msg in listen_websocket(url, self.log): + self.submitter(lambda: self._accept_msg(msg)) diff --git a/util/util.py b/util/util.py index cafb69b..04b6946 100644 --- a/util/util.py +++ b/util/util.py @@ -1,8 +1,12 @@ +import asyncio import logging import os import sys -from collections.abc import Callable +from collections.abc import AsyncIterator, Callable from typing import Any, cast +from urllib.parse import urlparse + +import websockets import env @@ -43,6 +47,48 @@ class Result[V, E]: return cast("Result[V, E]", Result(val, None)) +async def listen_websocket( + url: str, + log: logging.Logger, + *, + headers: dict[str, str] | None = None, + ping_interval: int = 20, + ping_timeout: int = 10, + close_timeout: int = 5, +) -> AsyncIterator[str | bytes]: + label = urlparse(url)._replace(path="", query="", fragment="").geturl() + + backoff = 1.0 + backoff_max: float = 60.0 + backoff_factor: float = 3.0 + + while True: + try: + async with websockets.connect( + url, + additional_headers=headers or {}, + ping_interval=ping_interval, + ping_timeout=ping_timeout, + close_timeout=close_timeout, + ) as ws: + log.info("Listening to '%s'...", label) + backoff = 1.0 + async for msg in ws: + yield msg + + except websockets.ConnectionClosedError as e: + log.warning("Connection closed: %s", e) + except TimeoutError as e: + log.warning("Connection timeout: %s", e) + except Exception as e: + log.error("Unexpected error: %s", e, exc_info=True) + + delay = min(backoff, backoff_max) + log.info("Reconnecting to '%s' in %.1fs...", label, delay) + await asyncio.sleep(delay) + backoff = min(backoff * backoff_factor, backoff_max) + + def normalize_service_url(url: str) -> str: if not url.startswith("https://") and not url.startswith("http://"): raise ValueError(f"Invalid service url {url}! Must start with http/https!") -- 2.51.2