stroblme
e37aa2bdd7
Per the Microcopy rule: size and duration errors, the flavor errors, run parameter errors, the Resources field descriptions, 401/409/400 details, secret errors, the websocket node description and its on_connect setting. The generated client follows.
432 lines
16 KiB
Python
432 lines
16 KiB
Python
"""WebSocket nodes: a socket the engine dials or serves, listening or sending."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import hmac
|
|
import json
|
|
import logging
|
|
from collections.abc import Awaitable, Callable, Iterable
|
|
from enum import Enum
|
|
from typing import TYPE_CHECKING, Any
|
|
from urllib.parse import urlsplit
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from fluksio.flow.messages import MessageSpec
|
|
from fluksio.flow.nodes.base import Node
|
|
from fluksio.flow.nodes.http import HttpNode, _first_catch_all
|
|
|
|
if TYPE_CHECKING:
|
|
from fastapi import FastAPI
|
|
from starlette.websockets import WebSocket
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class WebSocketNode(Node):
|
|
"""
|
|
A long-lived WebSocket, as a client or a server.
|
|
|
|
The URL picks the role: ``ws://`` or ``wss://`` dials out and holds the
|
|
connection, a path such as ``/ticker`` is served at
|
|
``/hooks/<flow>/<path>[/<secret>]`` while the flow runs. The ports pick the
|
|
direction, as on the mqtt node: outputs only make a **listener**, inputs a
|
|
**sender**.
|
|
|
|
- client listener: frames the server sends become outputs
|
|
- client sender: inputs go to the server, one frame per message
|
|
- server listener: frames any connected client sends become outputs
|
|
- server sender: inputs go to every connected client
|
|
|
|
Frames are text. An incoming frame is parsed as JSON; an object feeds the
|
|
ports whose names (or ``json_key``) it carries, and a node with one output
|
|
takes anything else whole. Outgoing, a single input is sent as it stands —
|
|
a string raw, anything else as JSON — and several go as one object keyed
|
|
by port.
|
|
"""
|
|
|
|
class Mode(Enum):
|
|
"""Which way the frames go."""
|
|
|
|
LISTENER = "listener"
|
|
SENDER = "sender"
|
|
|
|
# A frame sent twice is two commands to whatever reads it.
|
|
idempotent = False
|
|
|
|
__slots__ = (
|
|
"url",
|
|
"mode",
|
|
"serves",
|
|
"headers",
|
|
"on_connect",
|
|
"secret",
|
|
"timeout",
|
|
"send_queue_size",
|
|
"json_keys",
|
|
"_clients",
|
|
"_route_registered",
|
|
"_task",
|
|
"_queue",
|
|
"_loop",
|
|
)
|
|
|
|
class Params(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
url: str = Field(
|
|
description=(
|
|
"ws:// or wss:// to connect to a server, or a path such as "
|
|
"/ticker to serve one at /hooks/<flow>/<path>."
|
|
)
|
|
)
|
|
headers: dict[str, str] = Field(
|
|
default_factory=dict,
|
|
description=(
|
|
"Handshake headers, client only. A value may be a secret reference."
|
|
),
|
|
)
|
|
on_connect: list[Any] = Field(
|
|
default_factory=list,
|
|
description=(
|
|
"Frames sent after every (re)connect, client only: the "
|
|
"subscribe or auth message many APIs expect. A string goes "
|
|
"raw, anything else as JSON; an item may be a secret reference."
|
|
),
|
|
)
|
|
secret: str = Field(
|
|
default="",
|
|
description=(
|
|
"Shared secret for a served socket, appended to its URL: "
|
|
"/hooks/<flow>/<url>/<secret>. Empty leaves it open to anyone."
|
|
),
|
|
json_schema_extra={"x-secret": True},
|
|
)
|
|
timeout: float = Field(
|
|
default=10.0,
|
|
gt=0,
|
|
description=(
|
|
"Seconds allowed for the handshake, and for a send to one "
|
|
"client of a served socket."
|
|
),
|
|
)
|
|
send_queue_size: int = Field(
|
|
default=256,
|
|
gt=0,
|
|
description=(
|
|
"How many messages may wait to be sent. Past this the oldest "
|
|
"is dropped and the node reports degraded."
|
|
),
|
|
)
|
|
# Same as the mqtt node: a device that wraps its reading in an object.
|
|
json_key: str | dict[str, str] = Field(
|
|
default="",
|
|
description=(
|
|
"Key to lift out of a JSON object frame, for every port or as "
|
|
'{"port": "key"}.'
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def target_key(cls, params: dict[str, Any]) -> str | None:
|
|
# A URL, credentials stripped, exactly as for the http node.
|
|
return HttpNode.target_key(params)
|
|
|
|
def __init__(
|
|
self,
|
|
requires: MessageSpec | Iterable[MessageSpec] = (),
|
|
provides: MessageSpec | Iterable[MessageSpec] = (),
|
|
params: dict[str, Any] | None = None,
|
|
name: str | None = None,
|
|
):
|
|
cfg = self.Params.model_validate(params or {})
|
|
|
|
requires = Node._normalize_ports(requires)
|
|
provides = Node._normalize_ports(provides)
|
|
if not requires and not provides:
|
|
raise ValueError(
|
|
"A WebSocket node needs either inputs (to send) or outputs (to listen)"
|
|
)
|
|
# ponytail: one direction per node, as mqtt and http. Inputs and
|
|
# outputs over one socket is not supported yet.
|
|
self.mode = (
|
|
WebSocketNode.Mode.SENDER if requires else WebSocketNode.Mode.LISTENER
|
|
)
|
|
|
|
self.serves = cfg.url.startswith("/")
|
|
if not self.serves and urlsplit(cfg.url).scheme not in ("ws", "wss"):
|
|
raise ValueError(
|
|
"A WebSocket url is ws://…, wss://… or a path to serve, "
|
|
f"not '{cfg.url}'"
|
|
)
|
|
|
|
ports = requires or provides
|
|
if isinstance(cfg.json_key, dict):
|
|
self.json_keys: dict[str, str] = dict(cfg.json_key)
|
|
else:
|
|
self.json_keys = (
|
|
{spec.port: cfg.json_key for spec in ports} if cfg.json_key else {}
|
|
)
|
|
|
|
self.url = cfg.url
|
|
self.headers = cfg.headers
|
|
self.on_connect = [_frame(item) for item in cfg.on_connect]
|
|
self.secret = cfg.secret
|
|
self.timeout = cfg.timeout
|
|
self.send_queue_size = cfg.send_queue_size
|
|
|
|
self._clients: set[WebSocket] = set()
|
|
self._route_registered = False
|
|
self._task: asyncio.Task[None] | None = None
|
|
self._queue: asyncio.Queue[dict[str, Any]] | None = None
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|
|
|
if name is None:
|
|
safe_url = (
|
|
self.url.split("://", 1)[-1].replace("/", "_").replace(":", "")
|
|
).strip("_")
|
|
name = f"websocket_{self.mode.value}_{safe_url}"
|
|
|
|
super().__init__(
|
|
f=(
|
|
self._sender_handler
|
|
if self.mode is WebSocketNode.Mode.SENDER
|
|
else self._noop_listener
|
|
),
|
|
requires=requires,
|
|
provides=provides,
|
|
params=params,
|
|
name=name,
|
|
)
|
|
|
|
@staticmethod
|
|
def _noop_listener(params: dict[str, Any], **kwargs: Any) -> dict[str, Any] | None:
|
|
"""A listener injects what arrives; it is never called with inputs."""
|
|
return None
|
|
|
|
# -- frames ---------------------------------------------------------------
|
|
|
|
def _encode(self, data: dict[str, Any]) -> str:
|
|
"""One message as one frame."""
|
|
if len(data) == 1:
|
|
return _frame(next(iter(data.values())))
|
|
return json.dumps(data, default=str)
|
|
|
|
def _decode(self, text: str) -> dict[str, Any]:
|
|
"""One frame as outputs, keyed by port. Empty when nothing matched."""
|
|
try:
|
|
parsed = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
parsed = text
|
|
outputs: dict[str, Any] = {}
|
|
if isinstance(parsed, dict):
|
|
for spec in self.output_ports:
|
|
key = self.json_keys.get(spec.port, spec.port)
|
|
if key in parsed:
|
|
outputs[spec.port] = spec.coerce(parsed[key])
|
|
if not outputs and len(self.output_ports) == 1:
|
|
spec = self.output_ports[0]
|
|
outputs[spec.port] = spec.coerce(parsed)
|
|
return outputs
|
|
|
|
async def _receive(self, text: str) -> None:
|
|
"""Inject one frame. A bad one is logged, never fatal to the socket."""
|
|
try:
|
|
outputs = self._decode(text)
|
|
if outputs:
|
|
await asyncio.to_thread(self.inject, outputs)
|
|
else:
|
|
logger.debug("[%s] Frame matched no port: %s", self.name, text)
|
|
except Exception as exc:
|
|
logger.error("[%s] Error processing frame: %s", self.name, exc)
|
|
|
|
# -- sending --------------------------------------------------------------
|
|
|
|
def _sender_handler(
|
|
self, params: dict[str, Any], **kwargs: Any
|
|
) -> dict[str, Any] | None:
|
|
"""Hand the inputs to the node's socket task, which owns the send."""
|
|
loop, queue = self._loop, self._queue
|
|
if loop is not None and queue is not None:
|
|
loop.call_soon_threadsafe(self._enqueue, queue, dict(kwargs))
|
|
return None
|
|
|
|
# Not started: a draft preview or a test. A client connects once and
|
|
# sends; a served socket has nobody to send to.
|
|
if not self.serves:
|
|
from websockets.sync.client import connect
|
|
|
|
with connect(
|
|
self.url,
|
|
additional_headers=self.headers,
|
|
open_timeout=self.timeout,
|
|
) as ws:
|
|
for frame in self.on_connect:
|
|
ws.send(frame)
|
|
ws.send(self._encode(dict(kwargs)))
|
|
return None
|
|
|
|
def _enqueue(
|
|
self, queue: asyncio.Queue[dict[str, Any]], data: dict[str, Any]
|
|
) -> None:
|
|
"""Queue a message, dropping the oldest when the socket cannot keep up."""
|
|
if queue.full():
|
|
with contextlib.suppress(asyncio.QueueEmpty):
|
|
queue.get_nowait()
|
|
logger.warning("Send queue full for node '%s', dropped", self.name)
|
|
self.report_health("degraded", "send queue full")
|
|
queue.put_nowait(data)
|
|
|
|
async def _send_queued(self, send: Callable[[str], Awaitable[None]]) -> None:
|
|
"""Send whatever is queued, in order, for as long as the socket lives."""
|
|
queue = self._queue
|
|
if queue is None:
|
|
return
|
|
while True:
|
|
await send(self._encode(await queue.get()))
|
|
|
|
async def _broadcast(self, text: str) -> None:
|
|
"""One frame to every connected client.
|
|
|
|
ponytail: sequential, so a stalled client holds the others up for
|
|
`timeout`. Per-client queues if a served socket ever has many readers.
|
|
"""
|
|
for ws in list(self._clients):
|
|
try:
|
|
await asyncio.wait_for(ws.send_text(text), self.timeout)
|
|
except Exception:
|
|
self._clients.discard(ws)
|
|
with contextlib.suppress(Exception):
|
|
await asyncio.wait_for(ws.close(code=1011), self.timeout)
|
|
|
|
# -- client ---------------------------------------------------------------
|
|
|
|
async def _client_loop(self) -> None:
|
|
"""Hold one connection. A drop raises, and the supervisor reconnects."""
|
|
from websockets.asyncio.client import connect
|
|
|
|
try:
|
|
async with connect(
|
|
self.url,
|
|
additional_headers=self.headers,
|
|
open_timeout=self.timeout,
|
|
) as ws:
|
|
for frame in self.on_connect:
|
|
await ws.send(frame)
|
|
self.report_health("ok")
|
|
|
|
if self.mode is WebSocketNode.Mode.LISTENER:
|
|
async for received in ws:
|
|
if isinstance(received, str):
|
|
await self._receive(received)
|
|
else:
|
|
sender = asyncio.create_task(self._send_queued(ws.send))
|
|
try:
|
|
# Replies are not a sender's to publish, but they are
|
|
# read: an unread buffer stops the pings being answered.
|
|
async for _ in ws:
|
|
pass
|
|
finally:
|
|
await self._cancel_task(sender)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
self.report_health("down", str(exc))
|
|
raise
|
|
self.report_health("down", "connection closed")
|
|
raise ConnectionError(f"WebSocket for '{self.name}' closed")
|
|
|
|
# -- server ---------------------------------------------------------------
|
|
|
|
async def _serve(self, websocket: WebSocket) -> None:
|
|
"""One client of a served socket, for as long as it stays."""
|
|
from fluksio.cloud.config import via_portal
|
|
|
|
if (
|
|
self.secret
|
|
and not hmac.compare_digest(
|
|
str(websocket.path_params.get("secret", "")).encode(),
|
|
self.secret.encode(),
|
|
)
|
|
) or (not self.secret and via_portal(websocket.headers)):
|
|
# Closed before accepting: the same 403 as a path with no route.
|
|
# Through the portal a socket needs a secret, as a webhook does.
|
|
await websocket.close(code=1008)
|
|
return
|
|
await websocket.accept()
|
|
self._clients.add(websocket)
|
|
try:
|
|
while True:
|
|
message = await websocket.receive()
|
|
if message["type"] == "websocket.disconnect":
|
|
break
|
|
text = message.get("text")
|
|
# A sender's clients are read only to notice them leave.
|
|
if text is not None and self.mode is WebSocketNode.Mode.LISTENER:
|
|
await self._receive(text)
|
|
finally:
|
|
self._clients.discard(websocket)
|
|
|
|
def register_route(self, app: FastAPI) -> None:
|
|
"""Mount the served socket, ahead of any catch-all, as webhooks are."""
|
|
if self._route_registered:
|
|
return
|
|
from starlette.routing import WebSocketRoute
|
|
|
|
# The secret is a path parameter, so the route table never holds it.
|
|
path = f"{self.url}/{{secret}}" if self.secret else self.url
|
|
app.routes.insert(
|
|
_first_catch_all(app), WebSocketRoute(path, self._serve, name=self.id)
|
|
)
|
|
self._route_registered = True
|
|
logger.info("Registered websocket route for node '%s': %s", self.name, path)
|
|
|
|
def unregister_route(self, app: FastAPI) -> None:
|
|
if not self._route_registered:
|
|
return
|
|
app.routes[:] = [
|
|
route for route in app.routes if getattr(route, "name", None) != self.id
|
|
]
|
|
self._route_registered = False
|
|
|
|
# -- lifecycle ------------------------------------------------------------
|
|
|
|
async def start(self, app: FastAPI | None = None) -> None:
|
|
if self.serves and app is not None:
|
|
self.register_route(app)
|
|
if self.mode is WebSocketNode.Mode.SENDER and self._queue is None:
|
|
# Outlives a reconnect, so what was queued meanwhile still goes.
|
|
self._queue = asyncio.Queue(maxsize=self.send_queue_size)
|
|
self._loop = asyncio.get_running_loop()
|
|
# Spawned on every start: enabling a flow cancels its supervised loops
|
|
# and starts the nodes again, and the supervisor ignores a duplicate.
|
|
if self._task is not None and not self._task.done():
|
|
return
|
|
if not self.serves:
|
|
self._task = self._run_supervised("websocket", self._client_loop)
|
|
elif self.mode is WebSocketNode.Mode.SENDER:
|
|
self._task = self._run_supervised(
|
|
"websocket-out", lambda: self._send_queued(self._broadcast)
|
|
)
|
|
|
|
async def stop(self, app: FastAPI | None = None) -> None:
|
|
if self.serves and app is not None:
|
|
self.unregister_route(app)
|
|
for ws in list(self._clients):
|
|
with contextlib.suppress(Exception):
|
|
await ws.close(code=1001)
|
|
self._clients.clear()
|
|
if self._task is not None:
|
|
await self._cancel_task(self._task)
|
|
self._task = None
|
|
self._queue = None
|
|
self._loop = None
|
|
|
|
|
|
def _frame(value: Any) -> str:
|
|
"""A value as frame text: a string as it stands, anything else as JSON."""
|
|
# default=str: a value JSON cannot hold must not kill the send loop.
|
|
return value if isinstance(value, str) else json.dumps(value, default=str)
|