arno/etesync-server
1
0
Fork 0
mirror of https://github.com/etesync/server synced 2025-02-21 04:42:00 +00:00
Code Issues Releases Wiki Activity

Update the code for the new redis lib.

Browse Source
This commit is contained in:
Tom Hacohen 2022-09-04 13:51:17 -04:00
parent 1b90b2ad94
commit b5ebd64961
2 changed files with 16 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
etebase_server/fastapi/redis.py
View File

@ -12,12 +12,11 @@ class RedisWrapper:
async def setup(self):
if self.redis_uri is not None:
self.redis = await aioredis.create_redis_pool(self.redis_uri)
self.redis = await aioredis.from_url(self.redis_uri)
async def close(self):
if hasattr(self, "redis"):
self.redis.close()
await self.redis.wait_closed()
await self.redis.close()
@property
def is_active(self):

24
etebase_server/fastapi/routers/websocket.py
View File

@ -2,6 +2,7 @@ import asyncio
import typing as t
from redis import asyncio as aioredis
from redis.exceptions import ConnectionError
from asgiref.sync import sync_to_async
from django.db.models import QuerySet
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, status
@ -51,7 +52,7 @@ async def get_ticket(
uid = nacl.encoding.URLSafeBase64Encoder.encode(nacl.utils.random(32))
ticket_model = TicketInner(user=user.id, req=ticket_request)
ticket_raw = msgpack_encode(ticket_model.dict())
await redisw.redis.set(uid, ticket_raw, expire=TICKET_VALIDITY_SECONDS * 1000)
await redisw.redis.set(uid, ticket_raw, ex=TICKET_VALIDITY_SECONDS * 1000)
return TicketOut(ticket=uid)
@ -103,9 +104,9 @@ async def send_item_updates(
async def redis_connector(websocket: WebSocket, ticket_model: TicketInner, user: UserType, stoken: t.Optional[str]):
async def producer_handler(r: aioredis.Redis, ws: WebSocket):
pubsub = r.pubsub()
channel_name = f"col.{ticket_model.req.collection}"
(channel,) = await r.psubscribe(channel_name)
assert isinstance(channel, aioredis.Channel)
await pubsub.subscribe(channel_name)
# Send missing items if we are not up to date
queryset: QuerySet[models.Collection] = get_collection_queryset(user)
@ -117,12 +118,20 @@ async def redis_connector(websocket: WebSocket, ticket_model: TicketInner, user:
return
await send_item_updates(websocket, collection, user, stoken)
async def handle_message():
msg = await pubsub.get_message(ignore_subscribe_messages=True, timeout=20)
message_raw = t.cast(t.Optional[t.Tuple[str, bytes]], msg)
if message_raw:
_, message = message_raw
await ws.send_bytes(message)
try:
while True:
# We wait on the websocket so we fail if web sockets fail or get data
receive = asyncio.create_task(websocket.receive())
done, pending = await asyncio.wait(
{receive, channel.wait_message()}, return_when=asyncio.FIRST_COMPLETED
{receive, handle_message()},
return_when=asyncio.FIRST_COMPLETED,
)
for task in pending:
task.cancel()
@ -131,12 +140,7 @@ async def redis_connector(websocket: WebSocket, ticket_model: TicketInner, user:
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
return
message_raw = t.cast(t.Optional[t.Tuple[str, bytes]], await channel.get())
if message_raw:
_, message = message_raw
await ws.send_bytes(message)
except aioredis.errors.ConnectionClosedError:
except ConnectionError:
await websocket.close(code=status.WS_1012_SERVICE_RESTART)
except WebSocketDisconnect:
pass