use exceptiongroup.catch

pgjones · May 9, 2024 · 56d7dbf · 56d7dbf
1 parent 36cf97a
commit 56d7dbf
Showing 1 changed file with 26 additions and 21 deletions.
diff --git a/src/hypercorn/trio/tcp_server.py b/src/hypercorn/trio/tcp_server.py
@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import sys
 from math import inf
 from typing import Any, Generator, Optional
 
+import exceptiongroup
 import trio
 
 from .task_group import TaskGroup
@@ -51,28 +53,31 @@ async def run(self) -> None:
             socket = self.stream.socket
             ssl = False
 
+        def log_handler(e: Exception):
+            if self.config.log.error_logger is not None:
+                self.config.log.error_logger.exception("Internal hypercorn error")
+
         try:
-            client = parse_socket_addr(socket.family, socket.getpeername())
-            server = parse_socket_addr(socket.family, socket.getsockname())
-
-            async with TaskGroup() as task_group:
-                self._task_group = task_group
-                self.protocol = ProtocolWrapper(
-                    self.app,
-                    self.config,
-                    self.context,
-                    task_group,
-                    ssl,
-                    client,
-                    server,
-                    self.protocol_send,
-                    alpn_protocol,
-                )
-                await self.protocol.initiate()
-                await self._start_idle()
-                await self._read_data()
-        except (OSError, BaseExceptionGroup):
-            await self.config.log.exception("Internal hypercorn error")
+            with exceptiongroup.catch({Exception: log_handler}):
+                client = parse_socket_addr(socket.family, socket.getpeername())
+                server = parse_socket_addr(socket.family, socket.getsockname())
+
+                async with TaskGroup() as task_group:
+                    self._task_group = task_group
+                    self.protocol = ProtocolWrapper(
+                        self.app,
+                        self.config,
+                        self.context,
+                        task_group,
+                        ssl,
+                        client,
+                        server,
+                        self.protocol_send,
+                        alpn_protocol,
+                    )
+                    await self.protocol.initiate()
+                    await self._start_idle()
+                    await self._read_data()
         finally:
             await self._close()