69 lines
1.9 KiB
Python
Executable File
69 lines
1.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Example of aiohttp.web.Application.on_startup signal handler"""
|
|
import asyncio
|
|
from contextlib import suppress
|
|
from typing import AsyncIterator, List
|
|
|
|
import valkey.asyncio as valkey
|
|
|
|
from aiohttp import web
|
|
|
|
valkey_listener = web.AppKey("valkey_listener", asyncio.Task[None])
|
|
websockets = web.AppKey("websockets", List[web.WebSocketResponse])
|
|
|
|
|
|
async def websocket_handler(request):
|
|
ws = web.WebSocketResponse()
|
|
await ws.prepare(request)
|
|
request.app[websockets].append(ws)
|
|
try:
|
|
async for msg in ws:
|
|
print(msg)
|
|
await asyncio.sleep(1)
|
|
finally:
|
|
request.app[websockets].remove(ws)
|
|
return ws
|
|
|
|
|
|
async def on_shutdown(app: web.Application) -> None:
|
|
for ws in app[websockets]:
|
|
await ws.close(code=999, message=b"Server shutdown")
|
|
|
|
|
|
async def listen_to_valkey(app: web.Application) -> None:
|
|
r = valkey.Valkey(host="localhost", port=6379, decode_responses=True)
|
|
channel = "news"
|
|
async with r.pubsub() as sub:
|
|
await sub.subscribe(channel)
|
|
async for msg in sub.listen():
|
|
if msg["type"] != "message":
|
|
continue
|
|
# Forward message to all connected websockets:
|
|
for ws in app[websockets]:
|
|
await ws.send_str(f"{channel}: {msg}")
|
|
print(f"message in {channel}: {msg}")
|
|
|
|
|
|
async def background_tasks(app: web.Application) -> AsyncIterator[None]:
|
|
app[valkey_listener] = asyncio.create_task(listen_to_valkey(app))
|
|
|
|
yield
|
|
|
|
print("cleanup background tasks...")
|
|
app[valkey_listener].cancel()
|
|
with suppress(asyncio.CancelledError):
|
|
await app[valkey_listener]
|
|
|
|
|
|
def init():
|
|
app = web.Application()
|
|
l: List[web.WebSocketResponse] = []
|
|
app[websockets] = l
|
|
app.router.add_get("/news", websocket_handler)
|
|
app.cleanup_ctx.append(background_tasks)
|
|
app.on_shutdown.append(on_shutdown)
|
|
return app
|
|
|
|
|
|
web.run_app(init())
|