feat: push Douyin events through shared WebSocket and SSE

2026-10-07 23:45:32 +08:00
parent 297b813b90
commit 360bf96bab
30 changed files with 1625 additions and 140 deletions
@@ -0,0 +1,135 @@
+"""One shared-port WebSocket carrying event subscriptions and acknowledgements."""
+import json
+import logging
+import threading
+import time
+from simple_websocket import Server, ConnectionClosed
+from ..platform.douyin import DouyinError
+
+LOG = logging.getLogger('browser-gateway')
+
+
+def serve_channel(handler):
+    from .http import RequestError
+    if handler.headers.get('Upgrade', '').lower() != 'websocket':
+        raise RequestError('event channel requires WebSocket upgrade', 400)
+    environ = {'werkzeug.socket': handler.connection}
+    # simple-websocket's documented socket environment adapter owns framing and handshake.
+    for name, value in handler.headers.items():
+        environ['HTTP_' + name.upper().replace('-', '_')] = value
+    handler.connection.settimeout(None)
+    handler.close_connection = True
+    handler.websocket_handoff = True
+    ws = Server(environ, ping_interval=15, max_message_size=1024 * 1024)
+    handler.close_connection = True
+    manager = handler.server.gateway.subscriptions
+    gateway = handler.server.gateway
+    stopped = threading.Event()
+    lock = threading.RLock()
+    writing = threading.Lock()
+    subscriptions = {}
+
+    def send(value):
+        with writing:
+            ws.send(json.dumps(value, ensure_ascii=False, separators=(',', ':')))
+
+    def close():
+        try:
+            ws.close()
+        except ConnectionClosed:
+            LOG.debug('event WS already closed')
+
+    def sender():
+        try:
+            while not stopped.is_set():
+                with manager.changed:
+                    version = manager.version
+                with lock:
+                    for alias, state in list(subscriptions.items()):
+                        sub, item, sent = state
+                        try:
+                            current = manager._get(alias)
+                            error = item.error or ('event subscription replaced or stopped' if current is not item or item.stopped else '')
+                        except DouyinError as exc:
+                            error = str(exc)
+                        if error:
+                            send({'type': 'error', 'subscription': sub, 'alias': alias, 'error': error})
+                            subscriptions.pop(alias)
+                            continue
+                        if sent:
+                            continue
+                        pending = [d for d in item.pending() if d['delivery_id'] not in sent][:100]
+                        if pending:
+                            started = time.monotonic()
+                            send({'type': 'deliveries', 'subscription': sub, 'alias': alias, 'deliveries': pending})
+                            sent.update({d['delivery_id']: started for d in pending})
+                            LOG.info('event WS delivered alias=%s subscription=%s count=%s', alias, sub, len(pending))
+                with manager.changed:
+                    manager.changed.wait_for(lambda: stopped.is_set() or manager.version != version, timeout=15)
+        except Exception:
+            if not stopped.is_set():
+                LOG.exception('event WS delivery failed')
+                try:
+                    send({'type': 'error', 'error': 'event channel delivery failed'})
+                except ConnectionClosed:
+                    LOG.info('event WS disconnected during error reporting')
+                close()
+            stopped.set()
+
+    thread = threading.Thread(target=sender, name='event-ws-delivery', daemon=True)
+    thread.start()
+    try:
+        while not stopped.is_set():
+            raw = ws.receive()
+            if not isinstance(raw, str):
+                raise DouyinError('event channel requires text JSON')
+            value = json.loads(raw)
+            if not isinstance(value, dict) or value.get('type') not in ('subscribe', 'ack', 'unsubscribe'):
+                raise DouyinError('event channel message invalid')
+            alias, sub = value.get('alias'), value.get('subscription')
+            if not isinstance(alias, str) or not isinstance(sub, str) or not sub or len(sub) > 128:
+                raise DouyinError('event channel subscription invalid')
+            try:
+                # This is a metadata snapshot check, not a browser operation.
+                # Do not wait on the browser-operation lock in the shared receiver.
+                if value['type'] == 'subscribe':
+                    gateway._require_douyin_generation(alias, value)
+                with lock:
+                    if value['type'] == 'subscribe':
+                        item = manager._get(alias)
+                        if item.uid != value.get('uid') or item.stopped or getattr(item, 'generation', None) != value.get('listener_generation') or item.session_id != value.get('session_id'):
+                            raise DouyinError('event channel identity changed')
+                        subscriptions[alias] = (sub, item, {})
+                        send({'type': 'subscribed', 'alias': alias, 'subscription': sub})
+                    else:
+                        state = subscriptions.get(alias)
+                        if not state or state[0] != sub or manager._get(alias) is not state[1]:
+                            raise DouyinError('stale event channel subscription')
+                        if value['type'] == 'unsubscribe':
+                            subscriptions.pop(alias)
+                        else:
+                            ids = value.get('delivery_ids')
+                            if not isinstance(ids, list) or not ids or any(not isinstance(i, str) or i not in state[2] for i in ids):
+                                raise DouyinError('event acknowledgement invalid')
+                            state[1].ack(ids)
+                            elapsed_ms = round((time.monotonic() - min(state[2][i] for i in ids)) * 1000, 3)
+                            for delivery_id in ids:
+                                state[2].pop(delivery_id)
+                            LOG.info('event WS acknowledged alias=%s count=%s send_to_persist_ack_ms=%s', alias, len(ids), elapsed_ms)
+            except (DouyinError, RequestError) as exc:
+                LOG.warning('event WS subscription rejected alias=%s subscription=%s reason=%s', alias, sub, exc)
+                send({'type': 'error', 'subscription': sub, 'alias': alias, 'error': str(exc)})
+            manager.notify()
+    except ConnectionClosed:
+        LOG.info('event WS disconnected subscriptions=%s', len(subscriptions))
+    except Exception as exc:
+        LOG.exception('event WS request failed')
+        try:
+            send({'type': 'error', 'error': str(exc)})
+        except ConnectionClosed:
+            LOG.info('event WS disconnected during error reporting')
+    finally:
+        stopped.set()
+        manager.notify()
+        thread.join(timeout=2)
+        close()
@@ -512,7 +512,7 @@ class Gateway:
                self._require_douyin_generation(alias, input)
            # Polling-session initialization performs read-only page requests.
            # It must not occupy the manual-send slot while waiting on them.
-            result = self.subscriptions.start(alias, expected_uid, input.get("boundary_at"), input.get("notice_checkpoints", {}))
+            result = self.subscriptions.start(alias, expected_uid, input.get("boundary_at"), input.get("notice_checkpoints", {}), input.get('listener_generation'))
            self._require_douyin_generation(alias, input)
            return result
        except DouyinError as exc:
@@ -522,37 +522,18 @@ class Gateway:
            )
            raise RequestError(f"Douyin notification polling could not start: {exc}") from exc

-    def poll_douyin_events(self, alias: str, input: dict, query: dict) -> list[dict]:
-        if not valid_douyin_generation(input):
-            raise RequestError("invalid Douyin event request", 400)
-        try:
-            limit = int(query.get("limit", ["50"])[0])
-            wait = float(query.get("wait", ["0"])[0])
-        except (IndexError, TypeError, ValueError) as exc:
-            raise RequestError("invalid event poll options", 400) from exc
-        if not 1 <= limit <= 100 or not 0 <= wait <= 30:
-            raise RequestError("invalid event poll options", 400)
-        try:
-            with self._alias_lock(alias):
-                self._require_douyin_generation(alias, input)
-                acknowledgements = query.get("ack", [])
-                delivery_ids = [value for raw in acknowledgements for value in raw.split(",")]
-                self.subscriptions.ack(alias, delivery_ids)
-                return self.subscriptions.poll(alias, limit, 0)
-        except DouyinError as exc:
-            LOG.warning(
-                "Douyin event listener poll failed alias=%s reason=%s",
-                alias,
-                exc,
-            )
-            raise RequestError(f"Douyin notification polling failed: {exc}") from exc
-
    def stop_douyin_events(self, alias: str, input: dict) -> None:
        if not valid_douyin_generation(input):
            raise RequestError("invalid Douyin event request", 400)
        with self._alias_lock(alias):
            self._require_douyin_generation(alias, input)
-            self.subscriptions.stop(alias)
+            if input.get('stop_all') is True:
+                self.subscriptions.stop(alias)
+            else:
+                session_id = input.get('session_id')
+                if not isinstance(session_id, str) or not session_id:
+                    raise RequestError('event listener session_id required', 400)
+                self.subscriptions.stop(alias, input.get('listener_generation'), session_id)

 class GatewayHTTPServer(ThreadingHTTPServer):
    daemon_threads = True
@@ -635,9 +616,11 @@ class GatewayHandler(BaseHTTPRequestHandler):
            self._respond(401, json_bytes({"error": "gateway token rejected"}))
            return
        try:
-            needs_body = method in {"POST", "DELETE"} or (
-                method == "GET" and parsed.path.endswith("/douyin/events")
-            )
+            if parsed.path == '/v1/channel' and method == 'GET':
+                from .event_channel import serve_channel
+                serve_channel(self)
+                return
+            needs_body = method in {"POST", "DELETE"}
            body = self._body() if needs_body else {}
            result = self._route(method, parsed.path, parse_qs(parsed.query), body)
            if result is None:
@@ -653,6 +636,10 @@ class GatewayHandler(BaseHTTPRequestHandler):
            self._handle_exception(parsed.path, exc)

    def _handle_exception(self, path: str, exc: Exception) -> None:
+        if getattr(self, 'websocket_handoff', False):
+            LOG.exception('gateway WebSocket failed after handoff', extra={'path': path})
+            self.close_connection = True
+            return
        if isinstance(exc, RequestError):
            payload = {"error": str(exc)}
            if exc.network_id:
@@ -721,8 +708,6 @@ class GatewayHandler(BaseHTTPRequestHandler):
            if action == "events":
                if method == "POST":
                    return gateway.start_douyin_events(alias, body)
-                if method == "GET":
-                    return gateway.poll_douyin_events(alias, body, query)
                if method == "DELETE":
                    gateway.stop_douyin_events(alias, body)
                    return None