fine-tuning steps to properly cancel or exit (CLOSING > DRAINING) and finally closing a connection. Still not sure why __aexit__ is not called for a client connection used as an async context manager but adding the finally-block to QuicClient.connect is the last resort. Closing from server side needs to be investigated when implementing the test case where endpoint switches to no-DATAGRAM after establishing the connection and then the peer sends a DATAGRAM, which leads to PROTOCOL_VIOLATION and closing the connection from one side (namely server). See details in comment to Issue #16

Author: lilalinda

examples/echo_quic_client.py

@@ -66,7 +66,7 @@ async def parent(num: int = 0):
             data = await client_conn.receive()
             print(f"receiver: got data {data!r}")
 
-    print(f'Stopping client {num}')
+    print(f'Stopping client {num} in state={client_conn.state}')
 
 
 async def two_clients():

quicly/connection.py

@@ -174,15 +174,14 @@ def close(self) -> None:
         if self._closed:
             return
         self._closed = True
-        # TODO: cancel timers?
+        # close and clean up resources:
         self.on_tx_recv.close()
         self.sending_ch.close()
-        # TODO: how to communicate to Endpoint?
-        # if self.endpoint.connections.get(self.remote_address) is self:
-        #     del self.endpoint.connections[self.remote_address]
-        # Will wake any tasks waiting on self.q.r.receive() with a ClosedResourceError:
         self._stream_q.r.close()
         self._datagram_q.r.close()
+        # TODO: how to communicate to Endpoint so that they can clean up?
+        # if self.endpoint.connections.get(self.remote_address) is self:
+        #     del self.endpoint.connections[self.remote_address]
 
     async def aclose(self) -> None:
         """
@@ -209,7 +208,7 @@ async def __aexit__(
 
     async def __anext__(self):
         # using this object to iterate asynchronously uses ReceiveChannel semantics (DATAGRAM frames).
-        # `iter_stream_chunks()` is the async iterator for ReceiveStream semantics (STREAM frames).
+        # `aiter_stream_chunks()` is the async iterator for ReceiveStream semantics (STREAM frames).
         # TODO: we could also switch to ReceiveStream if DATAGRAM are not supported by local transport parameters?
         try:
             return await self.receive()  # ← Channel semantics instead of Stream!
@@ -219,7 +218,7 @@ async def __anext__(self):
     # TODO: this is only needed while Connection inherits both, ReceiveChannel and ReceiveStream:
     # TODO: use @as_safe_channel decorator?
     @trio.as_safe_channel
-    async def iter_stream_chunks(self, max_bytes: int | None = None):
+    async def aiter_stream_chunks(self, max_bytes: int | None = None):
         # since we direct `__anext__` above to use the channel receive semantics, if one wanted to use STREAM method
         # receive_some() in an async for loop, use this generator function
         while True:
@@ -445,7 +444,7 @@ async def do_handshake(self, hello_payload: bytes, remote_address: NetworkAddres
     async def on_tx(self, qpkt: QuicPacket) -> None:
         if self.is_closed:
             self._qlog.debug(f"Cannot send QUIC packet during connection closed - dropping.", packet=qpkt)
-            return
+            # return
 
         possible_ack_frame = self._pn_space.to_ack_frame(QuicFrameType.ACK, now=trio.current_time())
         if possible_ack_frame is not None:

quicly/endpoint.py

@@ -178,15 +178,16 @@ def __del__(self) -> None:
             )
 
     def close(self) -> None:
-        """Close this socket, and all associated QUIC connections.
-        This object can also be used as a context manager.
+        """
+        Close this socket, and all associated QUIC connections.
+        This object can also be used as a context manager as it implements __enter__ and __exit__.
         """
         self._closed = True
-        self.socket.close()
-        for stream in list(self._connections.values()):
-            stream.close()
+        for connection in list(self._connections.values()):
+            connection.close()
         self._new_connections_q.r.close()  # alerts anyone waiting on receive(), e.g., the server or client handshake
         self._send_q.r.close()
+        self.socket.close()
         self.dump_qlog()
 
     def __enter__(self) -> Self:
@@ -217,13 +218,13 @@ async def _datagram_received(self, udp_payload: bytes, remote_address: Any) -> N
             # connection for this DCID not yet established:
             await self._new_connections_q.s.send((udp_payload, remote_address, destination_cid))
         else:
-            self._qlog.debug(f"UDP datagram from known CID={hexdump(destination_cid)}", size=len(udp_payload))
+            # self._qlog.debug(f"UDP datagram from known CID={hexdump(destination_cid)}", size=len(udp_payload))
             await destination.on_rx(list(decode_udp_packet(udp_payload, destination_cid)), remote_address)
 
     def _config_logger(self, level: str | int):
         for h in self._qlog.handlers:
             if h.name == 'structlog-console':
-                h.setLevel("DEBUG")  # TODO: level)
+                h.setLevel("DEBUG")  # TODO: level)  Linda: replace!
 
     def dump_qlog(self):
         if isinstance(self._mem_qlog, QlogMemoryCollector):
@@ -301,7 +302,7 @@ async def handler(quic_connection):
             async def handle_connection(new_connection: SimpleQuicConnection) -> None:
                 async with new_connection:
                     await handler(new_connection, *args)
-                # TODO: if we end up here, we should remove the connection from the endpoint's list!
+                # TODO: if we end up here, we should remove the connection from the endpoint's list!?
 
             async with trio.open_nursery() as nursery:
                 self.start_endpoint(nursery)
@@ -435,6 +436,12 @@ async def watch_state():
                         yield connection  # give caller control; loops continue running in the nursery
                 finally:
                     # on context exit: cancel everything cleanly
+                    await connection.enter_closing(
+                        QuicFrame(
+                            QuicFrameType.APPLICATION_CLOSE,
+                            content=ConnectionCloseFrame(
+                                QuicErrorCode.NO_ERROR,
+                                reason=b'finally block'))
+                        if connection.state == ConnectionState.ESTABLISHED else None)
                     nursery.cancel_scope.cancel()
-                    await connection.aclose()
                     self.dump_qlog()

tests/test_endpoints.py

@@ -86,7 +86,7 @@ async def echo_handler(server_channel: SimpleQuicConnection) -> None:
                             await trio.sleep(delay)
                             await server_channel.send(payload)
                     else:
-                        async with server_channel.iter_stream_chunks() as recv_chan:
+                        async with server_channel.aiter_stream_chunks() as recv_chan:
                             async for payload in recv_chan:
                                 await trio.sleep(delay)
                                 await server_channel.send_all(payload)
@@ -102,8 +102,8 @@ async def echo_handler(server_channel: SimpleQuicConnection) -> None:
             if autocancel:
                 nursery.cancel_scope.cancel()
 
-# TODO: @parametrize_ipv6, also remove = False default below
-async def test_smoke_datagram(ipv6: bool = False) -> None:
+@parametrize_ipv6
+async def test_smoke_datagram(ipv6: bool) -> None:
     transport_parameters = {"max_datagram_frame_size" : 1200}  # add DATAGRAM support
     async with quic_echo_server(True, ipv6=ipv6, delay=0,
                                 transport_parameters=transport_parameters) as (_server_endpoint, address):
@@ -121,9 +121,11 @@ async def test_smoke_datagram(ipv6: bool = False) -> None:
                 await client_channel.send(b"goodbye")
                 answer = await client_channel.receive()
                 assert answer == b"goodbye"
+            assert client_channel.state == ConnectionState.DRAINING
 
 # TODO: test_fast_start (sending bytes with INITIAL...)
 
+# TODO: @parametrize_ipv6, also remove = False default below
 async def test_one_sided_datagram1(ipv6: bool = False) -> None:
     transport_parameters = {"max_datagram_frame_size" : 2400}  # add DATAGRAM support
     async with quic_echo_server(True, ipv6=ipv6, delay=0,
@@ -135,7 +137,6 @@ async def test_one_sided_datagram1(ipv6: bool = False) -> None:
             async with client.connect((get_localhost(ipv6, use_wildcard=False),) + address[1:],
                                       client_config) as connection:
                 assert connection.state == ConnectionState.ESTABLISHED
-                await trio.sleep(0.1)  # let handshake finalize for server
                 # server changes to not supporting DATAGRAM:
                 server = cast(QuicServer, _server_endpoint)
                 assert connection.host_cid in server._connections.keys()
@@ -146,6 +147,7 @@ async def test_one_sided_datagram1(ipv6: bool = False) -> None:
                 # TODO: PROTOCOL_VIOLATION!
                 await trio.sleep_forever() # let server close connection with PROTOCOL_VIOLATION
 
+# TODO: @parametrize_ipv6, also remove = False default below
 async def test_one_sided_datagram2(ipv6: bool = False) -> None:
     # server does not support DATAGRAM:
     async with quic_echo_server(True, ipv6=ipv6, delay=0) as (_server_endpoint, address):
@@ -158,8 +160,8 @@ async def test_one_sided_datagram2(ipv6: bool = False) -> None:
                 await connection.send(b"hello")
                 await trio.sleep(1) # let server close connection with PROTOCOL_VIOLATION
 
-# TODO: @parametrize_ipv6, also remove = False default below
-async def test_handshake_datagram(ipv6: bool = False) -> None:
+@parametrize_ipv6
+async def test_handshake_datagram(ipv6: bool) -> None:
     transport_parameters = {"max_datagram_frame_size" : 2400}  # add DATAGRAM support
     async with quic_echo_server(True, ipv6=ipv6, delay=0,
                                 transport_parameters=transport_parameters) as (_server_endpoint, address):
@@ -175,16 +177,7 @@ async def test_handshake_datagram(ipv6: bool = False) -> None:
                 assert connection.host_cid in server._connections.keys()
                 server_connection = server._connections[connection.host_cid]
                 assert server_connection.state == ConnectionState.ESTABLISHED
-            # Linda: why not called upon leaving async context manager???
-            # await connection.aclose()
-
-            # await trio.sleep(0)  # let closing() commence?
-            print(f"1st client state: {connection.state}")
-            await trio.sleep(1)  # let closing() commence?
-        print(f"2nd client state: {connection.state}")
-        await trio.sleep(0)  # let closing() commence?
-
-            # TODO: check cleanly shutdown?
+            assert connection.state == ConnectionState.DRAINING
 
 @parametrize_ipv6
 async def test_handshake(ipv6: bool) -> None: