diff --git a/bluesky/input.py b/bluesky/input.py --- 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 @@ url += "wantedCollections=app.bsky.feed.post" 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 --- 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 @@ @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 --- 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 @@ async def listen(self): 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 --- 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 @@ -41,6 +45,48 @@ @classmethod def ok(cls, val: V) -> "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: