Skip to content

Commit

Permalink
Fix asgi.py's Task thread safety
Browse files Browse the repository at this point in the history
  • Loading branch information
abersheeran committed Jan 6, 2025
1 parent d19c707 commit f4b94af
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions a2wsgi/asgi.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,14 @@
import threading
from http import HTTPStatus
from io import BytesIO
from typing import Any, Coroutine, Deque, Iterable, Optional
from typing import Any, Coroutine, Deque, Iterable, Optional, TypeVar
from typing import cast as typing_cast

from .asgi_typing import HTTPScope, ASGIApp, ReceiveEvent, SendEvent
from .wsgi_typing import Environ, StartResponse, IterableChunks

T = TypeVar("T")


class defaultdict(dict):
def __init__(self, default_factory, *args, **kwargs) -> None:
Expand Down Expand Up @@ -198,7 +200,7 @@ def asgi_done_callback(self, future: asyncio.Future) -> None:
finally:
self.asgi_done.set()

def start_asgi_app(self, environ: Environ) -> asyncio.Task:
async def start_asgi_app(self, environ: Environ) -> asyncio.Task:
run_asgi: asyncio.Task = self.loop.create_task(
typing_cast(
Coroutine[None, None, None],
Expand All @@ -208,6 +210,9 @@ def start_asgi_app(self, environ: Environ) -> asyncio.Task:
run_asgi.add_done_callback(self.asgi_done_callback)
return run_asgi

def execute_in_loop(self, coro: Coroutine[None, None, T]) -> T:
return asyncio.run_coroutine_threadsafe(coro, self.loop).result()

def __call__(
self, environ: Environ, start_response: StartResponse
) -> IterableChunks:
Expand All @@ -217,7 +222,7 @@ def __call__(
receive_eof = False
body_sent = False

asgi_task = self.start_asgi_app(environ)
asgi_task = self.execute_in_loop(self.start_asgi_app(environ))
# activate loop
self.loop.call_soon_threadsafe(lambda: None)

Expand Down Expand Up @@ -287,10 +292,10 @@ def __call__(
self.receive_event.set({"type": "http.disconnect"})
break

if asgi_task.done():
if self.asgi_done.is_set():
break

# HTTP response ends, wait for run_asgi's background tasks
self.asgi_done.wait(self.wait_time)
asgi_task.cancel()
self.loop.call_soon_threadsafe(asgi_task.cancel)
yield b""

0 comments on commit f4b94af

Please sign in to comment.