Draft of improved H3 for hypercorn.

Author: rthalley

src/hypercorn/asyncio/task_group.py

@@ -6,7 +6,7 @@
 from typing import Any, Awaitable, Callable, Optional
 
 from ..config import Config
-from ..typing import AppWrapper, ASGIReceiveCallable, ASGIReceiveEvent, ASGISendEvent, Scope
+from ..typing import AppWrapper, ASGIReceiveCallable, ASGIReceiveEvent, ASGISendEvent, Scope, Timer
 
 try:
     from asyncio import TaskGroup as AsyncioTaskGroup
@@ -33,6 +33,44 @@ async def _handle(
         await send(None)
 
 
+LONG_SLEEP = 86400.0
+
+class AsyncioTimer(Timer):
+    def __init__(self, action: Callable) -> None:
+        self._action = action
+        self._done = False
+        self._wake_up = asyncio.Condition()
+        self._when: Optional[float] = None
+
+    async def schedule(self, when: Optional[float]) -> None:
+        self._when = when
+        async with self._wake_up:
+            self._wake_up.notify()
+
+    async def stop(self) -> None:
+        self._done = True
+        async with self._wake_up:
+            self._wake_up.notify()
+
+    async def _wait_for_wake_up(self) -> None:
+        async with self._wake_up:
+            await self._wake_up.wait()
+
+    async def run(self) -> None:
+        while not self._done:
+            if self._when is not None and asyncio.get_event_loop().time() >= self._when:
+                self._when = None
+                await self._action()
+            if self._when is not None:
+                timeout = max(self._when - asyncio.get_event_loop().time(), 0.0)
+            else:
+                timeout = LONG_SLEEP
+            if not self._done:
+                try:
+                    await asyncio.wait_for(self._wait_for_wake_up(), timeout)
+                except TimeoutError:
+                    pass
+
 class TaskGroup:
     def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
         self._loop = loop
@@ -66,6 +104,11 @@ def _call_soon(func: Callable, *args: Any) -> Any:
     def spawn(self, func: Callable, *args: Any) -> None:
         self._task_group.create_task(func(*args))
 
+    def create_timer(self, action: Callable) -> Timer:
+        timer = AsyncioTimer(action)
+        self._task_group.create_task(timer.run())
+        return timer
+
     async def __aenter__(self) -> "TaskGroup":
         await self._task_group.__aenter__()
         return self

src/hypercorn/protocol/quic.py

@@ -22,7 +22,7 @@
 from .h3 import H3Protocol
 from ..config import Config
 from ..events import Closed, Event, RawData
-from ..typing import AppWrapper, TaskGroup, WorkerContext
+from ..typing import AppWrapper, TaskGroup, WorkerContext, Timer
 
 
 class QuicProtocol:
@@ -40,6 +40,7 @@ def __init__(
         self.context = context
         self.connections: Dict[bytes, QuicConnection] = {}
         self.http_connections: Dict[QuicConnection, H3Protocol] = {}
+        self.timers: Dict[QuicConnection, Timer] = {}
         self.send = send
         self.server = server
         self.task_group = task_group
@@ -82,10 +83,12 @@ async def handle(self, event: Event) -> None:
                 )
                 self.connections[header.destination_cid] = connection
                 self.connections[connection.host_cid] = connection
+                # This partial() needs python >= 3.8
+                self.timers[connection] = self.task_group.create_timer(partial(self._timeout, connection))
 
             if connection is not None:
                 connection.receive_datagram(event.data, event.address, now=self.context.time())
-                await self._handle_events(connection, event.address)
+                await self._wake_up_timer(connection)
         elif isinstance(event, Closed):
             pass
 
@@ -99,7 +102,16 @@ async def _handle_events(
         event = connection.next_event()
         while event is not None:
             if isinstance(event, ConnectionTerminated):
-                pass
+                await self.timers[connection].stop()
+                del self.timers[connection]
+                # XXXRTH This is not the speediest!  Better would be tracking
+                # assigned ids in a set.
+                prune = []
+                for tcid, tconn in self.connections.items():
+                    if tconn == connection:
+                        prune.append(tcid)
+                for tcid in prune:
+                    del self.connections[tcid]
             elif isinstance(event, ProtocolNegotiated):
                 self.http_connections[connection] = H3Protocol(
                     self.app,
@@ -109,7 +121,7 @@ async def _handle_events(
                     client,
                     self.server,
                     connection,
-                    partial(self.send_all, connection),
+                    partial(self._wake_up_timer, connection),
                 )
             elif isinstance(event, ConnectionIdIssued):
                 self.connections[event.connection_id] = connection
@@ -121,15 +133,20 @@ async def _handle_events(
 
             event = connection.next_event()
 
+    async def _wake_up_timer(self, connection: QuicConnection):
+        # When new output is send, or new input is received, we
+        # fire the timer right away so we update our state.
+        timer = self.timers.get(connection)
+        if timer is not None:
+            await timer.schedule(0.0)
+
+    async def _timeout(self, connection: QuicConnection):
+        now = self.context.time()
+        when = connection.get_timer()
+        if when is not None and now > when:
+            connection.handle_timer(now)
+        await self._handle_events(connection, None)
         await self.send_all(connection)
-
-        timer = connection.get_timer()
+        timer = self.timers.get(connection)
         if timer is not None:
-            self.task_group.spawn(self._handle_timer, timer, connection)
-
-    async def _handle_timer(self, timer: float, connection: QuicConnection) -> None:
-        wait = max(0, timer - self.context.time())
-        await self.context.sleep(wait)
-        if connection._close_at is not None:
-            connection.handle_timer(now=self.context.time())
-            await self._handle_events(connection, None)
+            await timer.schedule(connection.get_timer())

src/hypercorn/trio/task_group.py

@@ -7,7 +7,7 @@
 import trio
 
 from ..config import Config
-from ..typing import AppWrapper, ASGIReceiveCallable, ASGIReceiveEvent, ASGISendEvent, Scope
+from ..typing import AppWrapper, ASGIReceiveCallable, ASGIReceiveEvent, ASGISendEvent, Scope, Timer
 
 if sys.version_info < (3, 11):
     from exceptiongroup import BaseExceptionGroup
@@ -39,6 +39,40 @@ async def _handle(
         await send(None)
 
 
+LONG_SLEEP = 86400.0
+
+class TrioTimer(Timer):
+    def __init__(self, action: Callable) -> None:
+        self._action = action
+        self._done = False
+        self._wake_up = trio.Condition()
+        self._when: Optional[float] = None
+
+    async def schedule(self, when: Optional[float]) -> None:
+        self._when = when
+        async with self._wake_up:
+            self._wake_up.notify()
+
+    async def stop(self) -> None:
+        self._done = True
+        async with self._wake_up:
+            self._wake_up.notify()
+
+    async def run(self) -> None:
+        while not self._done:
+            if self._when is not None and trio.current_time() >= self._when:
+                self._when = None
+                await self._action()
+            if self._when is not None:
+                timeout = max(self._when - trio.current_time(), 0.0)
+            else:
+                timeout = LONG_SLEEP
+            if not self._done:
+                with trio.move_on_after(timeout):
+                    async with self._wake_up:
+                        await self._wake_up.wait()
+
+
 class TaskGroup:
     def __init__(self) -> None:
         self._nursery: Optional[trio._core._run.Nursery] = None
@@ -67,6 +101,11 @@ async def spawn_app(
     def spawn(self, func: Callable, *args: Any) -> None:
         self._nursery.start_soon(func, *args)
 
+    def create_timer(self, action: Callable) -> Timer:
+        timer = TrioTimer(action)
+        self._nursery.start_soon(timer.run)
+        return timer
+
     async def __aenter__(self) -> TaskGroup:
         self._nursery_manager = trio.open_nursery()
         self._nursery = await self._nursery_manager.__aenter__()

src/hypercorn/trio/worker_context.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Optional, Type, Union
+from typing import Awaitable, Optional, Type, Union
 
 import trio
 

src/hypercorn/typing.py

@@ -288,6 +288,20 @@ def is_set(self) -> bool:
         ...
 
 
+class Timer:
+    def __init__(self, action: Callable) -> None:
+        ...
+
+    async def schedule(self, when: float) -> None:
+        ...
+
+    async def stop(self) -> None:
+        ...
+
+    async def run(self) -> None:
+        ...
+
+
 class WorkerContext(Protocol):
     event_class: Type[Event]
     terminate: Event
@@ -318,6 +332,9 @@ async def spawn_app(
     def spawn(self, func: Callable, *args: Any) -> None:
         ...
 
+    def create_timer(self, action: Callable) -> Timer:
+        ...
+
     async def __aenter__(self) -> TaskGroup:
         ...