Fix: prevent cancellation of the message websocket

src/aleph/toolkit/shield.py

@@ -0,0 +1,13 @@
+import asyncio
+from functools import wraps
+
+
+def shielded(func):
+    """
+    Protects a coroutine from cancellation.
+    """
+    @wraps(func)
+    async def wrapped(*args, **kwargs):
+        return await asyncio.shield(func(*args, **kwargs))
+
+    return wrapped

src/aleph/web/controllers/messages.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from typing import List, Optional, Any, Dict, Iterable
 
@@ -31,6 +32,7 @@
     format_message_dict,
     PostMessage,
 )
+from aleph.toolkit.shield import shielded
 from aleph.types.db_session import DbSessionFactory, DbSession
 from aleph.types.message_status import MessageStatus
 from aleph.types.sort_order import SortOrder, SortBy
@@ -338,6 +340,7 @@ async def _process_message(mq_message: aio_pika.abc.AbstractMessage):
     return consumer_tag
 
 
+@shielded
 async def messages_ws(request: web.Request) -> web.WebSocketResponse:
     ws = web.WebSocketResponse()
     await ws.prepare(request)

src/aleph/web/controllers/p2p.py

@@ -13,6 +13,7 @@
 from aleph.schemas.pending_messages import parse_message, BasePendingMessage
 from aleph.services.ipfs import IpfsService
 from aleph.services.p2p.pubsub import publish as pub_p2p
+from aleph.toolkit.shield import shielded
 from aleph.types.message_status import (
     InvalidMessageException,
     MessageStatus,
@@ -183,6 +184,7 @@ class PubMessageResponse(BaseModel):
     message_status: Optional[MessageStatus]
 
 
+@shielded
 async def pub_message(request: web.Request):
     try:
         request_data = PubMessageRequest.parse_obj(await request.json())