Ensure to send ASGI websocket.disconnect after server-initiated close (#801)

* Fix websocket.disconnect not delivered after server-initiated close

Signed-off-by: JaeHyuck Sa <jaehyuck.sa.dev@gmail.com>

* Refactor code

* Reduce tests flakyness

---------

Signed-off-by: JaeHyuck Sa <jaehyuck.sa.dev@gmail.com>
Co-authored-by: Giovanni Barillari <giovanni.barillari@sentry.io>

Author: JaeHyuckSa

src/asgi/io.rs

@@ -558,6 +558,11 @@ impl ASGIWebsocketProtocol {
                     }
                 }
             }
+
+            if closed.load(atomic::Ordering::Acquire) {
+                return FutureResultToPy::ASGIWSMessage(Message::Close(None));
+            }
+
             FutureResultToPy::Err(error_flow!("Transport not initialized or closed"))
         })
     }

tests/apps/asgi.py

@@ -138,6 +138,17 @@ async def ws_push(scope, receive, send):
         pass
 
 
+async def ws_close(scope, receive, send):
+    await receive()
+    await send({'type': 'websocket.accept'})
+    msg = await receive()
+    fpath = msg.get('text') or msg['bytes'].decode('utf8')
+    await send({'type': 'websocket.close', 'code': 1000})
+    msg = await receive()
+    if msg['type'] == 'websocket.disconnect':
+        pathlib.Path(fpath).touch()
+
+
 async def err_app(scope, receive, send):
     1 / 0
 
@@ -208,6 +219,7 @@ def app(scope, receive, send):
         '/ws_rejectc': ws_reject_custom,
         '/ws_info': ws_info,
         '/ws_echo': ws_echo,
+        '/ws_close': ws_close,
         '/ws_push': ws_push,
         '/err_app': err_app,
         '/err_proto/type': err_proto_msg,

tests/test_ws.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import os
 
@@ -32,6 +33,25 @@ async def test_reject(server, runtime_mode):
     assert exc.value.response.status_code == 403
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize('runtime_mode', ['mt', 'st'])
+async def test_asgi_server_close(asgi_server, runtime_mode, tmp_path):
+    target = tmp_path / 'ws_result'
+
+    async with asgi_server(runtime_mode) as port:
+        async with websockets.connect(f'ws://localhost:{port}/ws_close') as ws:
+            await ws.send(str(target.resolve()))
+            try:
+                await ws.recv()
+            except Exception:
+                pass
+
+        # reduce flakyness
+        await asyncio.sleep(0.1)
+
+    assert target.exists()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize('runtime_mode', ['mt', 'st'])
 async def test_asgi_reject_custom(asgi_server, runtime_mode):