Clean up retry-tasks on shutdown

Signed-off-by: Mathias L. Baumann <[email protected]>

Author: Marenz

src/frequenz/dispatch/_actor_dispatcher.py

@@ -155,6 +155,12 @@ def new_receiver(self) -> Receiver[Dispatch]:
             """
             return self._channel.new_receiver()
 
+        def cleanup(self) -> None:
+            """Clean up the retry manager."""
+            for task in self._tasks:
+                task.cancel()
+            self._tasks.clear()
+
         def retry(self, dispatch: Dispatch) -> None:
             """Retry a dispatch.
 
@@ -290,13 +296,16 @@ async def _run(self) -> None:
             async for dispatch in self._dispatch_rx:
                 await self._handle_dispatch(dispatch)
         else:
-            retry_recv = self._retrier.new_receiver()
-
-            async for selected in select(retry_recv, self._dispatch_rx):
-                if retry_recv.triggered(selected):
-                    self._retrier.retry(selected.message)
-                elif self._dispatch_rx.triggered(selected):
-                    await self._handle_dispatch(selected.message)
+            try:
+                retry_recv = self._retrier.new_receiver()
+
+                async for selected in select(retry_recv, self._dispatch_rx):
+                    if retry_recv.triggered(selected):
+                        self._retrier.retry(selected.message)
+                    elif self._dispatch_rx.triggered(selected):
+                        await self._handle_dispatch(selected.message)
+            finally:
+                self._retrier.cleanup()
 
     async def _handle_dispatch(self, dispatch: Dispatch) -> None:
         """Handle a dispatch.