fastapi-modular 0.1.0__py3-none-any.whl

pymodular/core/websocket/namespace.py

@@ -0,0 +1,148 @@
+"""Namespace — sổ đăng ký các kết nối của MỘT gateway.
+
+Mỗi gateway (mỗi đường dẫn WebSocket) có namespace riêng, nên phòng "alerts"
+của `/ws/chat` và phòng "alerts" của `/ws/telemetry` là hai phòng khác nhau.
+Giống namespace của socket.io mà NestJS dùng.
+
+Ba chỉ mục được giữ song song, tất cả đều là O(1):
+
+    socket_id -> Socket        gửi thẳng cho một kết nối
+    room      -> {socket_id}   gửi cho một phòng
+    user_id   -> {socket_id}   gửi cho một người, kể cả khi họ mở nhiều tab
+
+Chỉ mục thứ ba là thứ hay bị bỏ sót. Không có nó thì "gửi thông báo cho user
+X" phải quét toàn bộ kết nối; với vài nghìn kết nối, mỗi lần gửi là một vòng
+lặp toàn bộ.
+
+Mọi thao tác ở đây là ĐỒNG BỘ và chỉ chạy trong event loop của tiến trình này
+nên không cần khoá. Nhiều worker thì mỗi worker có sổ riêng — xem adapter.py.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from typing import TYPE_CHECKING, Any
+
+from pymodular.core.logging import get_logger
+
+if TYPE_CHECKING:
+    from pymodular.core.websocket.socket import Socket
+
+log = get_logger(__name__)
+
+
+class Namespace:
+    __slots__ = ("_rooms", "_sockets", "_users", "path")
+
+    def __init__(self, path: str) -> None:
+        self.path = path
+        self._sockets: dict[str, Socket] = {}
+        self._rooms: dict[str, set[str]] = {}
+        self._users: dict[str, set[str]] = {}
+
+    # --------------------------------------------------------- vào/ra sổ
+    def add(self, socket: Socket) -> None:
+        self._sockets[socket.id] = socket
+        if socket.user_id:
+            self._users.setdefault(socket.user_id, set()).add(socket.id)
+
+    def remove(self, socket: Socket) -> None:
+        self._sockets.pop(socket.id, None)
+
+        for room in list(socket.rooms):
+            self.leave(socket, room)
+
+        if socket.user_id:
+            ids = self._users.get(socket.user_id)
+            if ids is not None:
+                ids.discard(socket.id)
+                # Xoá khoá rỗng, nếu không dict phình theo số user từng vào.
+                if not ids:
+                    del self._users[socket.user_id]
+
+    # ------------------------------------------------------------- phòng
+    def join(self, socket: Socket, room: str) -> None:
+        self._rooms.setdefault(room, set()).add(socket.id)
+        socket.rooms.add(room)
+
+    def leave(self, socket: Socket, room: str) -> None:
+        members = self._rooms.get(room)
+        if members is not None:
+            members.discard(socket.id)
+            if not members:
+                del self._rooms[room]
+        socket.rooms.discard(room)
+
+    def room_size(self, room: str) -> int:
+        return len(self._rooms.get(room, ()))
+
+    def sockets_in(self, room: str) -> list[Socket]:
+        return [s for sid in self._rooms.get(room, ()) if (s := self._sockets.get(sid))]
+
+    def sockets_of(self, user_id: str) -> list[Socket]:
+        return [s for sid in self._users.get(user_id, ()) if (s := self._sockets.get(sid))]
+
+    def get(self, socket_id: str) -> Socket | None:
+        return self._sockets.get(socket_id)
+
+    # -------------------------------------------------------------- gửi
+    def deliver(
+        self,
+        event: str,
+        data: Any = None,
+        *,
+        room: str | None = None,
+        user: str | None = None,
+        socket: str | None = None,
+        exclude: Iterable[str] = (),
+    ) -> int:
+        """Gửi cho các kết nối ĐANG Ở TIẾN TRÌNH NÀY. Trả về số tin đã xếp hàng.
+
+        Không có tiêu chí nào (`room`/`user`/`socket` đều None) nghĩa là gửi
+        cho toàn bộ namespace.
+        """
+        skip = set(exclude)
+        count = 0
+        for target in self._targets(room=room, user=user, socket=socket):
+            if target.id in skip or target.closing:
+                continue
+            if target.emit(event, data):
+                count += 1
+        return count
+
+    def _targets(
+        self, *, room: str | None, user: str | None, socket: str | None
+    ) -> Iterator[Socket]:
+        if socket is not None:
+            found = self._sockets.get(socket)
+            if found is not None:
+                yield found
+            return
+        if user is not None:
+            yield from self.sockets_of(user)
+            return
+        if room is not None:
+            yield from self.sockets_in(room)
+            return
+        yield from list(self._sockets.values())
+
+    # ------------------------------------------------------------ thống kê
+    def __len__(self) -> int:
+        return len(self._sockets)
+
+    @property
+    def sockets(self) -> list[Socket]:
+        return list(self._sockets.values())
+
+    @property
+    def rooms(self) -> dict[str, int]:
+        return {room: len(ids) for room, ids in self._rooms.items()}
+
+    def stats(self) -> dict[str, Any]:
+        return {
+            "path": self.path,
+            "sockets": len(self._sockets),
+            "users": len(self._users),
+            "rooms": len(self._rooms),
+            "pending": sum(s.pending for s in self._sockets.values()),
+        }

pymodular/core/websocket/protocol.py

@@ -0,0 +1,157 @@
+"""Khuôn tin nhắn và mã đóng kết nối.
+
+Một kết nối WebSocket là một ống byte trần: không có URL, không có status code,
+không có content-type. Muốn nhiều loại nghiệp vụ đi chung một kết nối thì phải
+tự quy ước, và quy ước đó chính là "giao thức" dưới đây — tương đương cặp
+`event` + `data` của socket.io mà NestJS dùng.
+
+Mọi khung tin là JSON một tầng:
+
+    {"event": "room.join", "data": {"room": "alerts"}, "id": "c1"}
+
+- `event` : tên nghiệp vụ, quyết định handler nào chạy. Bắt buộc.
+- `data`  : tham số, tuỳ handler. Thiếu thì hiểu là `null`.
+- `id`    : mã do CLIENT tự đặt để ghép câu trả lời với câu hỏi. Tuỳ chọn.
+
+Server trả lời bằng cùng khuôn, thêm `ack` = `id` của khung hỏi:
+
+    {"event": "room.join", "data": {"room": "alerts", "size": 3}, "ack": "c1"}
+    {"event": "error", "data": {"code": "forbidden", "message": "..."}, "ack": "c1"}
+
+Server đẩy chủ động thì không có `ack`:
+
+    {"event": "alert.created", "data": {...}}
+
+Khung có thêm ô `meta` tuỳ chọn, nội dung do bên gửi tự quyết. Lõi không đọc
+tới nó; nó dành cho các lớp bên ngoài muốn kèm nguồn gốc của tin, ví dụ:
+
+    {"event": "alert.created", "data": {...}, "meta": {"nguon": "..."}}
+
+Vì sao cần `id`/`ack`: WebSocket không ghép cặp request-response như HTTP. Gửi
+ba lệnh liên tiếp rồi nhận ba câu trả lời thì không có cách nào biết cái nào
+của cái nào, trừ khi tự đánh số.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Any
+
+from pymodular.core.exceptions import AppError
+
+
+class CloseCode(IntEnum):
+    """Mã đóng kết nối.
+
+    1000–2999 do chuẩn WebSocket định nghĩa; 4000–4999 dành cho ứng dụng tự
+    đặt. Ở đây dải 44xx cố ý ánh xạ 1-1 với HTTP status quen thuộc (4401 ~ 401,
+    4403 ~ 403...) để đọc log không phải tra bảng.
+    """
+
+    NORMAL = 1000
+    GOING_AWAY = 1001            # server tắt/khởi động lại
+    PROTOCOL_ERROR = 1002
+    TOO_BIG = 1009               # khung tin vượt ngưỡng
+    TRY_AGAIN_LATER = 1013       # client đọc quá chậm, hàng đợi gửi đã đầy
+
+    BAD_REQUEST = 4400
+    UNAUTHORIZED = 4401
+    FORBIDDEN = 4403
+    NOT_FOUND = 4404
+    IDLE_TIMEOUT = 4408          # im lặng quá lâu, coi như đã chết
+    CONFLICT = 4409
+    TOO_MANY = 4429              # vượt hạn mức kết nối hoặc tần suất gửi
+    INTERNAL = 4500
+
+
+# AppError mang status_code HTTP; đổi sang mã đóng tương ứng để client dùng
+# chung một bảng lỗi cho cả REST lẫn WebSocket.
+_HTTP_TO_CLOSE: dict[int, CloseCode] = {
+    400: CloseCode.BAD_REQUEST,
+    401: CloseCode.UNAUTHORIZED,
+    403: CloseCode.FORBIDDEN,
+    404: CloseCode.NOT_FOUND,
+    409: CloseCode.CONFLICT,
+    422: CloseCode.BAD_REQUEST,
+    429: CloseCode.TOO_MANY,
+}
+
+
+def close_code_for(exc: BaseException) -> CloseCode:
+    if isinstance(exc, AppError):
+        return _HTTP_TO_CLOSE.get(exc.status_code, CloseCode.INTERNAL)
+    return CloseCode.INTERNAL
+
+
+# Tên sự kiện do framework giữ chỗ; gateway không được đăng ký trùng.
+EVENT_CONNECTED = "connected"
+EVENT_ERROR = "error"
+EVENT_PING = "ping"
+EVENT_PONG = "pong"
+EVENT_ROOM_JOIN = "room.join"
+EVENT_ROOM_LEAVE = "room.leave"
+
+RESERVED_EVENTS = frozenset(
+    {EVENT_CONNECTED, EVENT_ERROR, EVENT_PONG, EVENT_ROOM_JOIN, EVENT_ROOM_LEAVE}
+)
+
+
+class ProtocolError(AppError):
+    """Khung tin sai khuôn — không phải JSON, thiếu `event`, hoặc quá dài."""
+
+    status_code = 400
+    error_code = "ws_protocol_error"
+    message = "Khung tin không hợp lệ"
+
+
+@dataclass(slots=True)
+class Frame:
+    event: str
+    data: Any = None
+    id: str | None = None
+    ack: str | None = None
+    meta: dict[str, Any] | None = None
+
+    def to_json(self) -> str:
+        payload: dict[str, Any] = {"event": self.event, "data": self.data}
+        if self.ack is not None:
+            payload["ack"] = self.ack
+        if self.meta is not None:
+            payload["meta"] = self.meta
+        # default=str để datetime/UUID/Decimal trong data không làm chết vòng
+        # gửi. Mất kiểu còn hơn mất kết nối.
+        return json.dumps(payload, ensure_ascii=False, default=str)
+
+
+def parse_frame(raw: str) -> Frame:
+    """Đọc khung tin từ client. Ném ProtocolError nếu sai khuôn."""
+    try:
+        payload = json.loads(raw)
+    except (json.JSONDecodeError, ValueError) as exc:
+        raise ProtocolError(f"Không phải JSON hợp lệ: {exc}") from exc
+
+    if not isinstance(payload, dict):
+        raise ProtocolError("Khung tin phải là một object JSON")
+
+    event = payload.get("event")
+    if not isinstance(event, str) or not event:
+        raise ProtocolError("Thiếu trường 'event'")
+
+    frame_id = payload.get("id")
+    if frame_id is not None and not isinstance(frame_id, str):
+        raise ProtocolError("Trường 'id' phải là chuỗi")
+
+    return Frame(event=event, data=payload.get("data"), id=frame_id)
+
+
+def error_frame(exc: BaseException, *, ack: str | None = None, debug: bool = False) -> Frame:
+    """Đổi exception thành khung `error` cùng hình dạng với lỗi HTTP."""
+    if isinstance(exc, AppError):
+        return Frame(event=EVENT_ERROR, data=exc.to_dict(), ack=ack)
+
+    data: dict[str, Any] = {"code": "internal_error", "message": "Internal server error"}
+    if debug:
+        data["details"] = f"{type(exc).__name__}: {exc}"
+    return Frame(event=EVENT_ERROR, data=data, ack=ack)

pymodular/core/websocket/server.py

@@ -0,0 +1,175 @@
+"""`WebSocketServer` — điểm vào duy nhất để đẩy tin, dùng được từ mọi nơi.
+
+Tương đương `@WebSocketServer() server: Server` của NestJS, nhưng lấy qua DI
+như mọi provider khác:
+
+    @injectable
+    class AlertService:
+        def __init__(self, ws: WebSocketServer) -> None:
+            self._ws = ws
+
+        async def raise_alert(self, alert):
+            await self._ws.to_room("alerts", "alert.created", {"id": alert.id})
+
+Gửi được từ controller HTTP, từ service, từ tác vụ nền — không cần đang ở
+trong một kết nối WebSocket nào. Server tự lo hai việc: gửi cho kết nối trong
+tiến trình này, và đăng tin lên adapter cho các worker khác.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+from pymodular.core.config import Settings
+from pymodular.core.container import injectable
+from pymodular.core.logging import get_logger
+from pymodular.core.metrics import ws_connections
+from pymodular.core.websocket.adapter import (
+    BroadcastAdapter,
+    build_adapter,
+    envelope,
+    new_origin_id,
+)
+from pymodular.core.websocket.namespace import Namespace
+from pymodular.core.websocket.protocol import CloseCode
+
+log = get_logger(__name__)
+
+
+@injectable
+class WebSocketServer:
+    def __init__(self, settings: Settings) -> None:
+        self._config = settings.ws
+        self._namespaces: dict[str, Namespace] = {}
+        self._origin = new_origin_id()
+        self._adapter: BroadcastAdapter = build_adapter(
+            self._config.adapter,
+            url=self._config.redis_url,
+            channel=self._config.channel,
+            origin=self._origin,
+        )
+        self._started = False
+
+    # ------------------------------------------------------------- namespace
+    def namespace(self, path: str) -> Namespace:
+        """Lấy (hoặc tạo) namespace của một đường dẫn WebSocket."""
+        found = self._namespaces.get(path)
+        if found is None:
+            found = self._namespaces[path] = Namespace(path)
+        return found
+
+    @property
+    def namespaces(self) -> dict[str, Namespace]:
+        return dict(self._namespaces)
+
+    def _resolve(self, path: str | None) -> Namespace:
+        if path is not None:
+            return self.namespace(path)
+        if len(self._namespaces) == 1:
+            return next(iter(self._namespaces.values()))
+        raise RuntimeError(
+            "Có nhiều gateway nên phải nói rõ gửi vào namespace nào: "
+            f"emit(..., namespace='/ws/chat'). Đang có: {sorted(self._namespaces)}"
+        )
+
+    # ------------------------------------------------------------------ gửi
+    async def emit(
+        self,
+        event: str,
+        data: Any = None,
+        *,
+        namespace: str | None = None,
+        room: str | None = None,
+        user: str | None = None,
+        socket: str | None = None,
+        exclude: Iterable[str] = (),
+        local_only: bool = False,
+    ) -> int:
+        """Đẩy một tin. Trả về số kết nối TẠI TIẾN TRÌNH NÀY đã nhận.
+
+        Số trả về cố ý không tính các worker khác: adapter là fire-and-forget,
+        biết chắc bên kia gửi được bao nhiêu thì phải chờ phản hồi, mà chờ thì
+        mất hết ý nghĩa "gửi không chặn".
+        """
+        target = self._resolve(namespace)
+        delivered = target.deliver(
+            event, data, room=room, user=user, socket=socket, exclude=exclude
+        )
+
+        if not local_only and self._adapter.name != "local":
+            await self._adapter.publish(
+                envelope(
+                    self._origin,
+                    target.path,
+                    event,
+                    data,
+                    room=room,
+                    user=user,
+                    socket=socket,
+                    exclude=exclude,
+                )
+            )
+        return delivered
+
+    async def to_room(self, room: str, event: str, data: Any = None, **kwargs: Any) -> int:
+        return await self.emit(event, data, room=room, **kwargs)
+
+    async def to_user(self, user_id: str, event: str, data: Any = None, **kwargs: Any) -> int:
+        """Gửi cho MỘT người — tới mọi kết nối của họ (nhiều tab, nhiều thiết bị)."""
+        return await self.emit(event, data, user=user_id, **kwargs)
+
+    async def to_socket(self, socket_id: str, event: str, data: Any = None, **kwargs: Any) -> int:
+        return await self.emit(event, data, socket=socket_id, **kwargs)
+
+    async def broadcast(self, event: str, data: Any = None, **kwargs: Any) -> int:
+        """Gửi cho mọi kết nối trong namespace."""
+        return await self.emit(event, data, **kwargs)
+
+    # ------------------------------------------------------- tin từ worker khác
+    def _on_remote(self, payload: dict[str, Any]) -> None:
+        path = payload.get("ns")
+        if not path or path not in self._namespaces:
+            return   # worker này không phục vụ namespace đó
+        self._namespaces[path].deliver(
+            payload.get("event", ""),
+            payload.get("data"),
+            room=payload.get("room"),
+            user=payload.get("user"),
+            socket=payload.get("socket"),
+            exclude=payload.get("exclude") or (),
+        )
+
+    # ------------------------------------------------------------ vòng đời
+    async def startup(self) -> None:
+        if self._started:
+            return
+        await self._adapter.start(self._on_remote)
+        self._started = True
+
+    async def shutdown(self) -> None:
+        """Đóng mọi kết nối tử tế trước khi tiến trình thoát.
+
+        Mã 1001 (going away) báo cho client biết đây là restart chứ không phải
+        lỗi, để phía client nối lại ngay thay vì chờ hết backoff.
+        """
+        for namespace in self._namespaces.values():
+            for socket in namespace.sockets:
+                await socket.close(CloseCode.GOING_AWAY, "server đang tắt")
+            ws_connections.set(0, namespace=namespace.path)
+        await self._adapter.stop()
+        self._namespaces.clear()
+        self._started = False
+
+    # ------------------------------------------------------------ thống kê
+    @property
+    def adapter_name(self) -> str:
+        return self._adapter.name
+
+    def stats(self) -> dict[str, Any]:
+        return {
+            "adapter": self._adapter.name,
+            "origin": self._origin,
+            "namespaces": [ns.stats() for ns in self._namespaces.values()],
+            "connections": sum(len(ns) for ns in self._namespaces.values()),
+        }