Disconnect dispatcher at cleanup

Author: Marenz

RELEASE_NOTES.md

@@ -14,4 +14,4 @@
 
 ## Bug Fixes
 
-<!-- Here goes notable bug fixes that are worth a special mention or explanation -->
+- Fixed an issue in the `Dispatcher` class where the client connection was not properly disconnected during cleanup, potentially causing unclosed socket errors.

src/frequenz/dispatch/_dispatcher.py

@@ -9,6 +9,7 @@
 import logging
 from asyncio import Event
 from datetime import timedelta
+from types import TracebackType
 from typing import Awaitable, Callable, Self
 
 from frequenz.channels import Receiver
@@ -221,6 +222,7 @@ def __init__(
         self._actor_dispatchers: dict[str, ActorDispatcher] = {}
         self._empty_event = Event()
         self._empty_event.set()
+        self._disconnecting_futurgite: asyncio.Future[None] | None = None
 
     @override
     def start(self) -> None:
@@ -235,19 +237,23 @@ def is_running(self) -> bool:
 
     @override
     async def wait(self) -> None:
-        """Wait until all actor dispatches are stopped."""
-        await asyncio.gather(self._bg_service.wait(), self._empty_event.wait())
+        """Wait until all actor dispatches are stopped and client is disconnected."""
+        if self._disconnecting_future is not None:
+            await self._disconnecting_future
 
+        await asyncio.gather(self._bg_service.wait(), self._empty_event.wait())
         self._actor_dispatchers.clear()
 
-    @override
     def cancel(self, msg: str | None = None) -> None:
-        """Stop the local dispatch service."""
+        """Stop the local dispatch service and initiate client disconnection."""
         self._bg_service.cancel(msg)
 
         for instance in self._actor_dispatchers.values():
             instance.cancel()
 
+        # Initiate client disconnection asynchronously
+        self._disconnecting_future = asyncio.ensure_future(self._client.disconnect())
+
     async def wait_for_initialization(self) -> None:
         """Wait until the background service is initialized."""
         await self._bg_service.wait_for_initialization()
@@ -358,9 +364,29 @@ async def __aenter__(self) -> Self:
             This background service.
         """
         await super().__aenter__()
+        await self._client.__aenter__()
         await self.wait_for_initialization()
         return self
 
+    @override
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        """Exit an async context.
+
+        Stop this background service and disconnect the client.
+
+        Args:
+            exc_type: The exception type.
+            exc_value: The exception value.
+            traceback: The traceback object.
+        """
+        await self._client.__aexit__(exc_type, exc_value, traceback)
+        await self._client.disconnect()
+
     def new_lifecycle_events_receiver(
         self, dispatch_type: str
     ) -> Receiver[DispatchEvent]: