mirror of
https://git.ianrenton.com/ian/spothole.git
synced 2026-08-06 02:21:42 +00:00
99 lines
3.9 KiB
Python
99 lines
3.9 KiB
Python
import logging
|
|
from datetime import datetime
|
|
from threading import Event, Lock, Thread
|
|
|
|
import pytz
|
|
from requests_sse import EventSource
|
|
|
|
from core.constants import HTTP_HEADERS
|
|
from providers.spot.spot_provider import SpotProvider
|
|
|
|
|
|
class SSESpotProvider(SpotProvider):
|
|
"""Spot provider using Server-Sent Events."""
|
|
|
|
def __init__(self, name, provider_config, url):
|
|
super().__init__(name, provider_config)
|
|
self._url = url
|
|
self._thread = None
|
|
self._last_event_id = None
|
|
self._stop_event = Event()
|
|
self._event_source_lock = Lock()
|
|
self._event_source = None
|
|
|
|
def start(self):
|
|
logging.info("Set up SSE connection to " + self.name + " spot API.")
|
|
self._stop_event.clear()
|
|
self._thread = Thread(target=self._run, name=f"SSESpotProvider-{self.name}")
|
|
self._thread.daemon = True
|
|
self._thread.start()
|
|
|
|
def stop(self):
|
|
self._stop_event.set()
|
|
|
|
with self._event_source_lock:
|
|
event_source = self._event_source
|
|
if event_source:
|
|
try:
|
|
event_source.close()
|
|
except Exception:
|
|
logging.exception(
|
|
"Exception closing SSE connection for " + self.name + " during stop()")
|
|
|
|
if self._thread:
|
|
self._thread.join(timeout=15)
|
|
if self._thread.is_alive():
|
|
logging.warning(self.name + " SSE worker thread did not exit on time and will be killed.")
|
|
|
|
def _on_open(self):
|
|
self.status = "Waiting for Data"
|
|
|
|
def _on_error(self):
|
|
self.status = "Connecting"
|
|
|
|
def _set_event_source(self, event_source):
|
|
with self._event_source_lock:
|
|
self._event_source = event_source
|
|
|
|
def _run(self):
|
|
while not self._stop_event.is_set():
|
|
try:
|
|
logging.debug("Connecting to " + self.name + " spot API...")
|
|
self.status = "Connecting"
|
|
with EventSource(self._url, headers=HTTP_HEADERS, latest_event_id=self._last_event_id, timeout=10,
|
|
on_open=self._on_open, on_error=self._on_error) as event_source:
|
|
self._set_event_source(event_source)
|
|
try:
|
|
for event in event_source:
|
|
if self._stop_event.is_set():
|
|
break
|
|
if event.type == 'message':
|
|
try:
|
|
self._last_event_id = event.last_event_id
|
|
new_spot = self._sse_message_to_spot(event.data)
|
|
if new_spot:
|
|
self._submit(new_spot)
|
|
|
|
self.status = "OK"
|
|
self.last_update_time = datetime.now(pytz.UTC)
|
|
logging.debug("Received data from " + self.name + " spot API.")
|
|
|
|
except Exception:
|
|
logging.exception(
|
|
"Exception processing message from SSE Spot Provider (" + self.name + ")")
|
|
finally:
|
|
self._set_event_source(None)
|
|
|
|
except Exception:
|
|
self.status = "Error"
|
|
logging.exception("Exception in SSE Spot Provider (" + self.name + ")")
|
|
else:
|
|
self.status = "Disconnected"
|
|
self._stop_event.wait(timeout=5) # Wait before trying to reconnect
|
|
|
|
def _sse_message_to_spot(self, message_data):
|
|
"""Convert an SSE message received from the API into a spot. The whole message data is provided here so the subclass
|
|
implementations can handle the message as JSON, XML, text, whatever the API actually provides."""
|
|
|
|
raise NotImplementedError("Subclasses must implement this method")
|