💅 Propagate error causes via asyncio protocols

This is supposed to unify setting exceptions on the future objects,
allowing to also attach their causes whenever available.

It's also supposed to help with tracking down what's happening
with #4581.

Author: webknjaz

aiohttp/_http_parser.pyx

@@ -19,7 +19,7 @@ from multidict import CIMultiDict as _CIMultiDict, CIMultiDictProxy as _CIMultiD
 from yarl import URL as _URL
 
 from aiohttp import hdrs
-from aiohttp.helpers import DEBUG
+from aiohttp.helpers import DEBUG, set_exception
 
 from .http_exceptions import (
     BadHttpMessage,
@@ -763,11 +763,13 @@ cdef int cb_on_body(cparser.llhttp_t* parser,
     cdef bytes body = at[:length]
     try:
         pyparser._payload.feed_data(body, length)
-    except BaseException as exc:
+    except BaseException as underlying_exc:
+        reraised_exc = underlying_exc
         if pyparser._payload_exception is not None:
-            pyparser._payload.set_exception(pyparser._payload_exception(str(exc)))
-        else:
-            pyparser._payload.set_exception(exc)
+            reraised_exc = pyparser._payload_exception(str(underlying_exc))
+
+        set_exception(pyparser._payload, reraised_exc, underlying_exc)
+
         pyparser._payload_error = 1
         return -1
     else:

aiohttp/base_protocol.py

@@ -1,6 +1,7 @@
 import asyncio
 from typing import Optional, cast
 
+from .helpers import set_exception
 from .tcp_helpers import tcp_nodelay
 
 
@@ -76,7 +77,11 @@ def connection_lost(self, exc: Optional[BaseException]) -> None:
         if exc is None:
             waiter.set_result(None)
         else:
-            waiter.set_exception(exc)
+            set_exception(
+                waiter,
+                ConnectionError("Connection lost"),
+                exc,
+            )
 
     async def _drain_helper(self) -> None:
         if not self.connected:

aiohttp/client_proto.py

@@ -4,13 +4,15 @@
 
 from .base_protocol import BaseProtocol
 from .client_exceptions import (
+    ClientConnectionError,
     ClientOSError,
     ClientPayloadError,
     ServerDisconnectedError,
     SocketTimeoutError,
 )
 from .helpers import (
     BaseTimerContext,
+    ErrorableMixin,
     set_exception,
     set_result,
     status_code_must_be_empty_body,
@@ -80,41 +82,71 @@ def is_connected(self) -> bool:
     def connection_lost(self, exc: Optional[BaseException]) -> None:
         self._drop_timeout()
 
-        if exc is not None:
-            set_exception(self.closed, exc)
-        else:
+        original_connection_error = exc
+        reraised_exc = original_connection_error
+
+        connection_closed_cleanly = original_connection_error is None
+        """Whether the connection got a clean EOF."""
+
+        if connection_closed_cleanly:
             set_result(self.closed, None)
+        else:
+            set_exception(
+                self.closed,
+                ClientConnectionError(
+                    f"Connection lost: {original_connection_error !s}",
+                ),
+                original_connection_error,
+            )
 
         if self._payload_parser is not None:
-            with suppress(Exception):
+            with suppress(Exception):  # FIXME: log this somehow?
                 self._payload_parser.feed_eof()
 
         uncompleted = None
         if self._parser is not None:
             try:
                 uncompleted = self._parser.feed_eof()
-            except Exception as e:
+            except Exception as underlying_exc:
                 if self._payload is not None:
-                    exc = ClientPayloadError("Response payload is not completed")
-                    exc.__cause__ = e
-                    self._payload.set_exception(exc)
+                    client_payload_exc_msg = (
+                        f"Response payload is not completed: {underlying_exc !r}"
+                    )
+                    if not connection_closed_cleanly:
+                        client_payload_exc_msg = (
+                            f"{client_payload_exc_msg !s}. "
+                            f"{original_connection_error !r}"
+                        )
+                    set_exception(
+                        self._payload,
+                        ClientPayloadError(client_payload_exc_msg),
+                        underlying_exc,
+                    )
 
         if not self.is_eof():
-            if isinstance(exc, OSError):
-                exc = ClientOSError(*exc.args)
-            if exc is None:
-                exc = ServerDisconnectedError(uncompleted)
+            if isinstance(original_connection_error, OSError):
+                reraised_exc = ClientOSError(*original_connection_error.args)
+            if connection_closed_cleanly:
+                reraised_exc = ServerDisconnectedError(uncompleted)
             # assigns self._should_close to True as side effect,
             # we do it anyway below
-            self.set_exception(exc)
+            set_exc_kwargs = (
+                {}
+                if connection_closed_cleanly
+                else {
+                    "exc_cause": original_connection_error,
+                }
+            )
+            assert reraised_exc is not None
+            self.set_exception(reraised_exc, **set_exc_kwargs)
 
         self._should_close = True
         self._parser = None
         self._payload = None
         self._payload_parser = None
         self._reading_paused = False
 
-        super().connection_lost(exc)
+        super().connection_lost(reraised_exc)
 
     def eof_received(self) -> None:
         # should call parser.feed_eof() most likely
@@ -128,10 +160,14 @@ def resume_reading(self) -> None:
         super().resume_reading()
         self._reschedule_timeout()
 
-    def set_exception(self, exc: BaseException) -> None:
+    def set_exception(
+        self,
+        exc: BaseException,
+        exc_cause: BaseException = ErrorableMixin._EXC_SENTINEL,
+    ) -> None:
         self._should_close = True
         self._drop_timeout()
-        super().set_exception(exc)
+        super().set_exception(exc, exc_cause)
 
     def set_parser(self, parser: Any, payload: Any) -> None:
         # TODO: actual types are:
@@ -208,7 +244,7 @@ def _on_read_timeout(self) -> None:
         exc = SocketTimeoutError("Timeout on reading data from socket")
         self.set_exception(exc)
         if self._payload is not None:
-            self._payload.set_exception(exc)
+            set_exception(self._payload, exc)
 
     def data_received(self, data: bytes) -> None:
         self._reschedule_timeout()
@@ -234,14 +270,19 @@ def data_received(self, data: bytes) -> None:
                 # parse http messages
                 try:
                     messages, upgraded, tail = self._parser.feed_data(data)
-                except BaseException as exc:
+                except BaseException as underlying_exc:
                     if self.transport is not None:
                         # connection.release() could be called BEFORE
                         # data_received(), the transport is already
                         # closed in this case
                         self.transport.close()
                     # should_close is True after the call
-                    self.set_exception(exc)
+                    self.set_exception(
+                        ClientPayloadError(
+                            f"Unable to parse response payload: {underlying_exc !r}"
+                        ),
+                        underlying_exc,
+                    )
                     return
 
                 self._upgraded = upgraded

aiohttp/client_reqrep.py

@@ -53,6 +53,7 @@
     noop,
     parse_mimetype,
     reify,
+    set_exception,
     set_result,
 )
 from .http import (
@@ -566,20 +567,33 @@ async def write_bytes(
 
                 for chunk in self.body:
                     await writer.write(chunk)  # type: ignore[arg-type]
-        except OSError as exc:
-            if exc.errno is None and isinstance(exc, asyncio.TimeoutError):
-                protocol.set_exception(exc)
-            else:
-                new_exc = ClientOSError(
-                    exc.errno, "Can not write request body for %s" % self.url
+        except OSError as underlying_exc:
+            reraised_exc = underlying_exc
+
+            exc_is_not_timeout = underlying_exc.errno is not None or not isinstance(
+                underlying_exc, asyncio.TimeoutError
+            )
+            if exc_is_not_timeout:
+                reraised_exc = ClientOSError(
+                    underlying_exc.errno,
+                    f"Can not write request body for {self.url !s}",
                 )
-                new_exc.__context__ = exc
-                new_exc.__cause__ = exc
-                protocol.set_exception(new_exc)
+
+                # FIXME: Is setting `__context__` harmful?
+                # FIXME: Especially, given that `__cause__` is also set...
+                reraised_exc.__context__ = underlying_exc
+
+            set_exception(protocol, reraised_exc, underlying_exc)
         except asyncio.CancelledError:
             await writer.write_eof()
-        except Exception as exc:
-            protocol.set_exception(exc)
+        except Exception as underlying_exc:
+            set_exception(
+                protocol,
+                ClientConnectionError(
+                    f"Failed to send bytes into the underlying connection {conn !s}",
+                ),
+                underlying_exc,
+            )
         else:
             await writer.write_eof()
             protocol.start_timeout()
@@ -1019,7 +1033,7 @@ def _notify_content(self) -> None:
         content = self.content
         # content can be None here, but the types are cheated elsewhere.
         if content and content.exception() is None:  # type: ignore[truthy-bool]
-            content.set_exception(ClientConnectionError("Connection closed"))
+            set_exception(content, ClientConnectionError("Connection closed"))
         self._released = True
 
     async def wait_for_close(self) -> None:

aiohttp/helpers.py

@@ -797,9 +797,38 @@ def set_result(fut: "asyncio.Future[_T]", result: _T) -> None:
         fut.set_result(result)
 
 
-def set_exception(fut: "asyncio.Future[_T]", exc: BaseException) -> None:
-    if not fut.done():
-        fut.set_exception(exc)
+class ErrorableMixin:
+    _EXC_SENTINEL = BaseException()
+
+    def set_exception(
+        self,
+        exc: BaseException,
+        exc_cause: BaseException = _EXC_SENTINEL,
+    ):
+        raise NotImplementedError
+
+
+def set_exception(
+    fut: "asyncio.Future[_T] | ErrorableMixin",
+    exc: BaseException,
+    exc_cause: BaseException = ErrorableMixin._EXC_SENTINEL,
+) -> None:
+    """Set future exception.
+
+    If the future is marked as complete, this function is a no-op.
+
+    :param exc_cause: An exception that is a direct cause of ``exc``.
+                      Only set if provided.
+    """
+    if asyncio.isfuture(fut) and fut.done():
+        return
+
+    exc_is_sentinel = exc_cause is ErrorableMixin._EXC_SENTINEL
+    exc_causes_itself = exc is exc_cause
+    if not exc_is_sentinel and not exc_causes_itself:
+        exc.__cause__ = exc_cause
+
+    fut.set_exception(exc)
 
 
 @functools.total_ordering

aiohttp/http_parser.py

@@ -31,7 +31,9 @@
     DEBUG,
     NO_EXTENSIONS,
     BaseTimerContext,
+    ErrorableMixin,
     method_must_be_empty_body,
+    set_exception,
     status_code_must_be_empty_body,
 )
 from .http_exceptions import (
@@ -439,13 +441,16 @@ def get_content_length() -> Optional[int]:
                 assert self._payload_parser is not None
                 try:
                     eof, data = self._payload_parser.feed_data(data[start_pos:], SEP)
-                except BaseException as exc:
+                except BaseException as underlying_exc:
+                    reraised_exc = underlying_exc
                     if self.payload_exception is not None:
-                        self._payload_parser.payload.set_exception(
-                            self.payload_exception(str(exc))
-                        )
-                    else:
-                        self._payload_parser.payload.set_exception(exc)
+                        reraised_exc = self.payload_exception(str(underlying_exc))
+
+                    set_exception(
+                        self._payload_parser.payload,
+                        reraised_exc,
+                        underlying_exc,
+                    )
 
                     eof = True
                     data = b""
@@ -826,7 +831,7 @@ def feed_data(
                             exc = TransferEncodingError(
                                 chunk[:pos].decode("ascii", "surrogateescape")
                             )
-                            self.payload.set_exception(exc)
+                            set_exception(self.payload, exc)
                             raise exc
                         size = int(bytes(size_b), 16)
 
@@ -929,8 +934,12 @@ def __init__(self, out: StreamReader, encoding: Optional[str]) -> None:
         else:
             self.decompressor = ZLibDecompressor(encoding=encoding)
 
-    def set_exception(self, exc: BaseException) -> None:
-        self.out.set_exception(exc)
+    def set_exception(
+        self,
+        exc: BaseException,
+        exc_cause: BaseException = ErrorableMixin._EXC_SENTINEL,
+    ) -> None:
+        set_exception(self.out, exc, exc_cause)
 
     def feed_data(self, chunk: bytes, size: int) -> None:
         if not size:

aiohttp/http_websocket.py

@@ -25,7 +25,7 @@
 
 from .base_protocol import BaseProtocol
 from .compression_utils import ZLibCompressor, ZLibDecompressor
-from .helpers import NO_EXTENSIONS
+from .helpers import NO_EXTENSIONS, set_exception
 from .streams import DataQueue
 
 __all__ = (
@@ -305,7 +305,7 @@ def feed_data(self, data: bytes) -> Tuple[bool, bytes]:
             return self._feed_data(data)
         except Exception as exc:
             self._exc = exc
-            self.queue.set_exception(exc)
+            set_exception(self.queue, exc)
             return True, b""
 
     def _feed_data(self, data: bytes) -> Tuple[bool, bytes]: