Ensure deleted broadcast receivers get cleaned up

... by storing just a weakref to a receiver in the channel object.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/channels/broadcast.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import logging
+import weakref
 from asyncio import Condition
 from collections import deque
 from typing import Deque, Dict, Generic, Optional
@@ -81,7 +82,7 @@ def __init__(self, name: str, resend_latest: bool = False) -> None:
         self._resend_latest = resend_latest
 
         self.recv_cv: Condition = Condition()
-        self.receivers: Dict[UUID, Receiver[T]] = {}
+        self.receivers: Dict[UUID, weakref.ReferenceType[Receiver[T]]] = {}
         self.closed: bool = False
         self._latest: Optional[T] = None
 
@@ -101,17 +102,6 @@ async def close(self) -> None:
         async with self.recv_cv:
             self.recv_cv.notify_all()
 
-    def _drop_receiver(self, uuid: UUID) -> None:
-        """Drop a specific receiver from the list of broadcast receivers.
-
-        Called from the destructors of receivers.
-
-        Args:
-            uuid: a uuid identifying the receiver to be dropped.
-        """
-        if uuid in self.receivers:
-            del self.receivers[uuid]
-
     def get_sender(self) -> Sender[T]:
         """Create a new broadcast sender.
 
@@ -140,7 +130,7 @@ def get_receiver(
         if name is None:
             name = str(uuid)
         recv: Receiver[T] = Receiver(uuid, name, maxsize, self)
-        self.receivers[uuid] = recv
+        self.receivers[uuid] = weakref.ref(recv)
         if self._resend_latest and self._latest is not None:
             recv.enqueue(self._latest)
         return recv
@@ -188,8 +178,15 @@ async def send(self, msg: T) -> bool:
             return False
         # pylint: disable=protected-access
         self._chan._latest = msg
-        for recv in self._chan.receivers.values():
+        stale_refs = []
+        for name, recv_ref in self._chan.receivers.items():
+            recv = recv_ref()
+            if recv is None:
+                stale_refs.append(name)
+                continue
             recv.enqueue(msg)
+        for name in stale_refs:
+            del self._chan.receivers[name]
         async with self._chan.recv_cv:
             self._chan.recv_cv.notify_all()
         return True
@@ -225,11 +222,6 @@ def __init__(self, uuid: UUID, name: str, maxsize: int, chan: Broadcast[T]) -> N
 
         self._active = True
 
-    def __del__(self) -> None:
-        """Drop this receiver from the list of Broadcast receivers."""
-        if self._active:
-            self._chan._drop_receiver(self._uuid)
-
     def enqueue(self, msg: T) -> None:
         """Put a message into this receiver's queue.
 
@@ -295,7 +287,6 @@ def into_peekable(self) -> Peekable[T]:
         Returns:
             A `Peekable` instance.
         """
-        self._chan._drop_receiver(self._uuid)  # pylint: disable=protected-access
         self._active = False
         return Peekable(self._chan)
 

tests/test_broadcast.py

@@ -217,3 +217,25 @@ async def test_broadcast_map() -> None:
 
     assert (await receiver.receive()) is False
     assert (await receiver.receive()) is True
+
+
+async def test_broadcast_receiver_drop() -> None:
+    """Ensure deleted receivers get cleaned up."""
+    chan = Broadcast[int]("input-chan")
+    sender = chan.get_sender()
+
+    receiver1 = chan.get_receiver()
+    receiver2 = chan.get_receiver()
+
+    await sender.send(10)
+
+    assert 10 == await receiver1.receive()
+    assert 10 == await receiver2.receive()
+
+    assert len(chan.receivers) == 2
+
+    del receiver2
+
+    await sender.send(20)
+
+    assert len(chan.receivers) == 1