mirror of
https://git.ianrenton.com/ian/spothole.git
synced 2026-09-21 06:47:42 +00:00
96 lines
3.7 KiB
Python
96 lines
3.7 KiB
Python
import logging
|
|
from datetime import datetime
|
|
from threading import Event, Thread
|
|
from typing import Any
|
|
|
|
import pytz
|
|
from websocket import WebSocket, create_connection
|
|
|
|
from core.constants import HTTP_HEADERS
|
|
from data.spot import Spot
|
|
from providers.spot.spot_provider import SpotProvider
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class WebsocketSpotProvider(SpotProvider):
|
|
"""Spot provider using websockets."""
|
|
|
|
def __init__(self, name: str, provider_config: dict[str, Any], url: str) -> None:
|
|
super().__init__(name, provider_config)
|
|
self._url: str = url
|
|
self._ws: WebSocket | None = None
|
|
self._thread: Thread | None = None
|
|
self._stop_event: Event = Event()
|
|
self._last_event_id: str | None = None
|
|
|
|
def start(self) -> None:
|
|
logger.info(f"Set up websocket connection to {self.name} spot API.")
|
|
self._stop_event.clear()
|
|
self._thread = Thread(target=self._run, name=f"WebsocketSpotProvider-{self.name}")
|
|
self._thread.daemon = True
|
|
self._thread.start()
|
|
|
|
def stop(self) -> None:
|
|
self._stop_event.set()
|
|
if self._ws:
|
|
self._ws.close()
|
|
if self._thread:
|
|
self._thread.join(timeout=5)
|
|
if self._thread.is_alive():
|
|
logger.warning(f"{self.name} websocket worker thread did not exit on time and will be killed.")
|
|
|
|
def _on_open(self) -> None:
|
|
self.status = "Waiting for Data"
|
|
|
|
def _on_error(self) -> None:
|
|
self.status = "Connecting"
|
|
|
|
def _run(self) -> None:
|
|
while not self._stop_event.is_set():
|
|
try:
|
|
logger.debug(f"Connecting to {self.name} spot API...")
|
|
self.status = "Connecting"
|
|
self._ws = create_connection(self._url, header=HTTP_HEADERS)
|
|
self.status = "Connected"
|
|
|
|
# Keep reading from this same connection until it drops or we're asked to stop, rather than
|
|
# reconnecting for every message.
|
|
while not self._stop_event.is_set():
|
|
data = self._ws.recv()
|
|
if not data:
|
|
break
|
|
try:
|
|
new_spot = self._ws_message_to_spot(data)
|
|
if new_spot:
|
|
self._submit(new_spot)
|
|
|
|
self.status = "OK"
|
|
self.last_update_time = datetime.now(pytz.UTC)
|
|
logger.debug(f"Received data from {self.name} spot API.")
|
|
|
|
except Exception:
|
|
logger.exception(f"Exception processing message from Websocket Spot Provider ({self.name})")
|
|
|
|
except Exception:
|
|
self.status = "Error"
|
|
logger.exception(f"Exception in Websocket Spot Provider ({self.name})")
|
|
else:
|
|
self.status = "Disconnected"
|
|
finally:
|
|
if self._ws:
|
|
try:
|
|
self._ws.close()
|
|
except Exception:
|
|
# No problem, we were getting rid of this object anyway.
|
|
logger.debug(f"Exception while closing socket in {self.name}", exc_info=True)
|
|
self._ws = None
|
|
if not self._stop_event.is_set():
|
|
self._stop_event.wait(timeout=5) # Wait before trying to reconnect
|
|
|
|
def _ws_message_to_spot(self, b: str | bytes) -> Spot | None:
|
|
"""Convert a WS message received from the API into a spot. The exact message data (in bytes) is provided here so the
|
|
subclass implementations can handle the message as string, JSON, XML, whatever the API actually provides."""
|
|
|
|
raise NotImplementedError("Subclasses must implement this method")
|