136 lines
6.7 KiB
Python
136 lines
6.7 KiB
Python
"""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()
|