|
1 | 1 | import asyncio
|
2 | 2 | from contextlib import suppress
|
3 |
| -from typing import Any, Optional, Tuple |
| 3 | +from typing import Any, Optional, Tuple, Union |
4 | 4 |
|
5 | 5 | from .base_protocol import BaseProtocol
|
6 | 6 | from .client_exceptions import (
|
@@ -45,7 +45,27 @@ def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
|
45 | 45 | self._read_timeout_handle: Optional[asyncio.TimerHandle] = None
|
46 | 46 |
|
47 | 47 | self._timeout_ceil_threshold: Optional[float] = 5
|
48 |
| - self.closed: asyncio.Future[None] = self._loop.create_future() |
| 48 | + |
| 49 | + self._closed: Union[None, asyncio.Future[None]] = None |
| 50 | + self._connection_lost_called = False |
| 51 | + |
| 52 | + @property |
| 53 | + def closed(self) -> Union[None, asyncio.Future[None]]: |
| 54 | + """Future that is set when the connection is closed. |
| 55 | +
|
| 56 | + This property returns a Future that will be completed when the connection |
| 57 | + is closed. The Future is created lazily on first access to avoid creating |
| 58 | + futures that will never be awaited. |
| 59 | +
|
| 60 | + Returns: |
| 61 | + - A Future[None] if the connection is still open or was closed after |
| 62 | + this property was accessed |
| 63 | + - None if connection_lost() was already called before this property |
| 64 | + was ever accessed (indicating no one is waiting for the closure) |
| 65 | + """ |
| 66 | + if self._closed is None and not self._connection_lost_called: |
| 67 | + self._closed = self._loop.create_future() |
| 68 | + return self._closed |
49 | 69 |
|
50 | 70 | @property
|
51 | 71 | def upgraded(self) -> bool:
|
@@ -79,30 +99,31 @@ def is_connected(self) -> bool:
|
79 | 99 | return self.transport is not None and not self.transport.is_closing()
|
80 | 100 |
|
81 | 101 | def connection_lost(self, exc: Optional[BaseException]) -> None:
|
| 102 | + self._connection_lost_called = True |
82 | 103 | self._drop_timeout()
|
83 | 104 |
|
84 | 105 | original_connection_error = exc
|
85 | 106 | reraised_exc = original_connection_error
|
86 | 107 |
|
87 | 108 | connection_closed_cleanly = original_connection_error is None
|
88 | 109 |
|
89 |
| - if connection_closed_cleanly: |
90 |
| - set_result(self.closed, None) |
91 |
| - else: |
92 |
| - assert original_connection_error is not None |
93 |
| - set_exception( |
94 |
| - self.closed, |
95 |
| - ClientConnectionError( |
96 |
| - f"Connection lost: {original_connection_error !s}", |
97 |
| - ), |
98 |
| - original_connection_error, |
99 |
| - ) |
100 |
| - # Mark the exception as retrieved to prevent |
101 |
| - # "Future exception was never retrieved" warnings |
102 |
| - # The exception is always passed on through |
103 |
| - # other means, so this is safe |
104 |
| - with suppress(Exception): |
105 |
| - self.closed.exception() |
| 110 | + if self._closed is not None: |
| 111 | + # If someone is waiting for the closed future, |
| 112 | + # we should set it to None or an exception. If |
| 113 | + # self._closed is None, it means that |
| 114 | + # connection_lost() was called already |
| 115 | + # or nobody is waiting for it. |
| 116 | + if connection_closed_cleanly: |
| 117 | + set_result(self._closed, None) |
| 118 | + else: |
| 119 | + assert original_connection_error is not None |
| 120 | + set_exception( |
| 121 | + self._closed, |
| 122 | + ClientConnectionError( |
| 123 | + f"Connection lost: {original_connection_error !s}", |
| 124 | + ), |
| 125 | + original_connection_error, |
| 126 | + ) |
106 | 127 |
|
107 | 128 | if self._payload_parser is not None:
|
108 | 129 | with suppress(Exception): # FIXME: log this somehow?
|
|
0 commit comments