Make the registry type-aware (#830)

The channel registry is using `Any` as the message type, which is not
type safe, as it completely *disables* type checking for the messages.

This commit makes the channel registry type-aware, so channels are
stored with their message type and the registry checks that the same
channel is not used for different message types.

This also makes the registry just a plain container for channels, the
wrapper methods to create new senders and receivers, and to configure
the `resend_latest` flag are removed. The `ReceiverFetcher` abstraction
is also not needed if we just return the channel directly, as the
channel itself is a `ReceiverFetcher`.

Also the method to close and remove a channel is made public and the
name more explicit, as it is used in normal code paths.

The new registry only provide 2 main methods:

* `get_or_create()`: Get or create a channel for the given key, doing
the type checking to make sure the requested message type matches the
existing channel message type if it already exists.

* `close_and_remove()`: Close and remove the channel for the given key.

This change uncovered 5 issues:

* `MockMigrogrid/Resampler: Fail on unhandled exceptions` (in this PR,
but more generally #826)
* `Fix checks in tests` (in this PR)
* `Fix missing conversion to `QuantityT` (workaroud in this PR, but
should be better solved by #821)
* #807
* #823

Fixes #806.

Author: llucax

RELEASE_NOTES.md

@@ -50,10 +50,28 @@
 
 - The `microgrid.frequency()` method no longer supports passing the `component` parameter. Instead the best component is automatically selected.
 
+- The `actor.ChannelRegistry` was rewritten to be type-aware and just a container of channels. You now need to provide the type of message that will be contained by the channel and use the `get_or_create()` method to get a channel and the `stop_and_remove()` method to stop and remove a channel. Once you get a channel you can create new senders and receivers, or set channel options, as usual. Please read the docs for a full description, but in general this:
+
+    ```python
+    r = registry.new_receiver(name)
+    s = registry.new_sender(name)
+    ```
+
+    Should be replaced by:
+
+    ```python
+    r = registry.get_or_create(T, name).new_receiver()
+    s = registry.get_or_create(T, name).new_sender()
+    ```
+
+- The `ReceiverFetcher` interface was slightly changed to make `maxsize` a keyword-only argument. This is to make it compatible with the `Broadcast` channel, so it can be considered a `ReceiverFetcher`.
+
 ## New Features
 
 - A new method `microgrid.voltage()` was added to allow easy access to the phase-to-neutral 3-phase voltage of the microgrid.
 
+- The `actor.ChannelRegistry` is now type-aware.
+
 ## Bug Fixes
 
 - 0W power requests are now not adjusted to exclusion bounds by the `PowerManager` and `PowerDistributor`, and are sent over to the microgrid API directly.

benchmarks/timeseries/benchmark_datasourcing.py

@@ -108,7 +108,9 @@ async def consume(channel: Receiver[Any]) -> None:
                 "current_phase_requests", evc_id, component_metric_id, None
             )
 
-            recv_channel = channel_registry.new_receiver(request.get_channel_name())
+            recv_channel = channel_registry.get_or_create(
+                ComponentMetricRequest, request.get_channel_name()
+            ).new_receiver()
 
             await request_sender.send(request)
             consume_tasks.append(asyncio.create_task(consume(recv_channel)))

src/frequenz/sdk/_internal/_channels.py

@@ -16,7 +16,7 @@ class ReceiverFetcher(typing.Generic[T], typing.Protocol):
     """An interface that just exposes a `new_receiver` method."""
 
     @abc.abstractmethod
-    def new_receiver(self, maxsize: int = 50) -> Receiver[T]:
+    def new_receiver(self, *, maxsize: int = 50) -> Receiver[T]:
         """Get a receiver from the channel.
 
         Args:

src/frequenz/sdk/actor/_channel_registry.py

@@ -3,132 +3,148 @@
 
 """A class that would dynamically create, own and provide access to channels."""
 
-from __future__ import annotations
+import dataclasses
+import logging
+import traceback
+from typing import TypeVar, cast
 
-import typing
+from frequenz.channels import Broadcast
 
-from frequenz.channels import Broadcast, Receiver, Sender
-
-from .._internal._channels import ReceiverFetcher
+_T = TypeVar("_T")
+_logger = logging.getLogger(__name__)
 
 
 class ChannelRegistry:
-    """Dynamically creates, own and provide access to channels.
+    """Dynamically creates, own and provide access to broadcast channels.
 
     It can be used by actors to dynamically establish a communication channel
-    between each other.  Channels are identified by string names.
+    between each other.
+
+    The registry is responsible for creating channels when they are first requested via
+    the [`get_or_create()`][frequenz.sdk.actor.ChannelRegistry.get_or_create] method.
+
+    The registry also stores type information to make sure that the same channel is not
+    used for different message types.
+
+    Since the registry owns the channels, it is also responsible for closing them when
+    they are no longer needed. There is no way to remove a channel without closing it.
+
+    Note:
+        This registry stores [`Broadcast`][frequenz.channels.Broadcast] channels.
     """
 
     def __init__(self, *, name: str) -> None:
-        """Create a `ChannelRegistry` instance.
+        """Initialize this registry.
 
         Args:
-            name: A unique name for the registry.
+            name: A name to identify the registry in the logs. This name is also used as
+                a prefix for the channel names.
         """
         self._name = name
-        self._channels: dict[str, Broadcast[typing.Any]] = {}
+        self._channels: dict[str, _Entry] = {}
 
-    def set_resend_latest(self, key: str, resend_latest: bool) -> None:
-        """Set the `resend_latest` flag for a given channel.
+    @property
+    def name(self) -> str:
+        """The name of this registry."""
+        return self._name
 
-        This flag controls whether the latest value of the channel should be resent to
-        new receivers, in slow streams.
-
-        `resend_latest` is `False` by default.  It is safe to be set in data/reporting
-        channels, but is not recommended for use in channels that stream control
-        instructions.
+    def message_type(self, key: str) -> type:
+        """Get the message type of the channel for the given key.
 
         Args:
             key: The key to identify the channel.
-            resend_latest: Whether to resend the latest value to new receivers, for the
-                given channel.
-        """
-        if key not in self._channels:
-            self._channels[key] = Broadcast(f"{self._name}-{key}")
-        # This attribute is protected in the current version of the channels library,
-        # but that will change in the future.
-        self._channels[key].resend_latest = resend_latest
-
-    def new_sender(self, key: str) -> Sender[typing.Any]:
-        """Get a sender to a dynamically created channel with the given key.
-
-        Args:
-            key: A key to identify the channel.
 
         Returns:
-            A sender to a dynamically created channel with the given key.
-        """
-        if key not in self._channels:
-            self._channels[key] = Broadcast(f"{self._name}-{key}")
-        return self._channels[key].new_sender()
-
-    def new_receiver(self, key: str, maxsize: int = 50) -> Receiver[typing.Any]:
-        """Get a receiver to a dynamically created channel with the given key.
+            The message type of the channel.
 
-        Args:
-            key: A key to identify the channel.
-            maxsize: The maximum size of the receiver.
-
-        Returns:
-            A receiver for a dynamically created channel with the given key.
+        Raises:
+            KeyError: If the channel does not exist.
         """
-        if key not in self._channels:
-            self._channels[key] = Broadcast(f"{self._name}-{key}")
-        return self._channels[key].new_receiver(maxsize=maxsize)
+        entry = self._channels.get(key)
+        if entry is None:
+            raise KeyError(f"No channel for key {key!r} exists.")
+        return entry.message_type
 
-    def new_receiver_fetcher(self, key: str) -> ReceiverFetcher[typing.Any]:
-        """Get a receiver fetcher to a dynamically created channel with the given key.
+    def __contains__(self, key: str) -> bool:
+        """Check whether the channel for the given `key` exists."""
+        return key in self._channels
 
-        Args:
-            key: A key to identify the channel.
-
-        Returns:
-            A receiver fetcher for a dynamically created channel with the given key.
-        """
-        if key not in self._channels:
-            self._channels[key] = Broadcast(f"{self._name}-{key}")
-        return _RegistryReceiverFetcher(self, key)
+    def get_or_create(self, message_type: type[_T], key: str) -> Broadcast[_T]:
+        """Get or create a channel for the given key.
 
-    async def _close_channel(self, key: str) -> None:
-        """Close a channel with the given key.
+        If a channel for the given key already exists, the message type of the existing
+        channel is checked against the requested message type. If they do not match,
+        a `ValueError` is raised.
 
-        This method is private and should only be used in special cases.
+        Note:
+            The types have to match exactly, it doesn't do a subtype check due to
+            technical limitations. In the future subtype checks might be supported.
 
         Args:
-            key: A key to identify the channel.
-        """
-        if key in self._channels:
-            if channel := self._channels.pop(key, None):
-                await channel.close()
-
-
-T = typing.TypeVar("T")
-
-
-class _RegistryReceiverFetcher(typing.Generic[T]):
-    """A receiver fetcher that is bound to a channel registry and a key."""
+            message_type: The type of the message that is sent through the channel.
+            key: The key to identify the channel.
 
-    def __init__(
-        self,
-        registry: ChannelRegistry,
-        key: str,
-    ) -> None:
-        """Create a new instance of a receiver fetcher.
+        Returns:
+            The channel for the given key.
 
-        Args:
-            registry: The channel registry.
-            key: A key to identify the channel.
+        Raises:
+            ValueError: If the channel exists and the message type does not match.
         """
-        self._registry = registry
-        self._key = key
-
-    def new_receiver(self, maxsize: int = 50) -> Receiver[T]:
-        """Get a receiver from the channel.
+        if key not in self._channels:
+            if _logger.isEnabledFor(logging.DEBUG):
+                _logger.debug(
+                    "Creating a new channel for key %r with type %s at:\n%s",
+                    key,
+                    message_type,
+                    "".join(traceback.format_stack(limit=10)[:9]),
+                )
+            self._channels[key] = _Entry(message_type, Broadcast(f"{self._name}-{key}"))
+
+        entry = self._channels[key]
+        if entry.message_type is not message_type:
+            exception = ValueError(
+                f"Type mismatch, a channel for key {key!r} exists and the requested "
+                f"message type {message_type} is not the same as the existing "
+                f"message type {entry.message_type}."
+            )
+            if _logger.isEnabledFor(logging.DEBUG):
+                _logger.debug(
+                    "%s at:\n%s",
+                    str(exception),
+                    # We skip the last frame because it's this method, and limit the
+                    # stack to 9 frames to avoid adding too much noise.
+                    "".join(traceback.format_stack(limit=10)[:9]),
+                )
+            raise exception
+
+        return cast(Broadcast[_T], entry.channel)
+
+    async def close_and_remove(self, key: str) -> None:
+        """Remove the channel for the given key.
 
         Args:
-            maxsize: The maximum size of the receiver.
+            key: The key to identify the channel.
 
-        Returns:
-            A receiver instance.
+        Raises:
+            KeyError: If the channel does not exist.
         """
-        return self._registry.new_receiver(self._key, maxsize)
+        entry = self._channels.pop(key, None)
+        if entry is None:
+            raise KeyError(f"No channel for key {key!r} exists.")
+        await entry.channel.close()
+
+
+@dataclasses.dataclass(frozen=True)
+class _Entry:
+    """An entry in a channel registry."""
+
+    message_type: type
+    """The type of the message that is sent through the channel in this entry."""
+
+    # We use object instead of Any to minimize the chances of hindering type checking.
+    # If for some reason the channel is not casted to the proper underlaying type, when
+    # using object at least accessing any member that's not part of the object base
+    # class will yield a type error, while if we used Any, it would not and the issue
+    # would be much harder to find.
+    channel: Broadcast[object]
+    """The channel in this entry."""

src/frequenz/sdk/actor/_data_sourcing/microgrid_api_source.py

@@ -325,7 +325,9 @@ def _get_metric_senders(
             (
                 self._get_data_extraction_method(category, metric),
                 [
-                    self._registry.new_sender(request.get_channel_name())
+                    self._registry.get_or_create(
+                        Sample[Quantity], request.get_channel_name()
+                    ).new_sender()
                     for request in req_list
                 ],
             )
@@ -379,8 +381,7 @@ def process_msg(data: Any) -> None:
 
         await asyncio.gather(
             *[
-                # pylint: disable=protected-access
-                self._registry._close_channel(r.get_channel_name())
+                self._registry.close_and_remove(r.get_channel_name())
                 for requests in self._req_streaming_metrics[comp_id].values()
                 for r in requests
             ]

src/frequenz/sdk/actor/_power_managing/_power_managing_actor.py

@@ -211,14 +211,16 @@ async def _run(self) -> None:
 
                 if component_ids not in self._subscriptions:
                     self._subscriptions[component_ids] = {
-                        priority: self._channel_registry.new_sender(
-                            sub.get_channel_name()
-                        )
+                        priority: self._channel_registry.get_or_create(
+                            _Report, sub.get_channel_name()
+                        ).new_sender()
                     }
                 elif priority not in self._subscriptions[component_ids]:
                     self._subscriptions[component_ids][
                         priority
-                    ] = self._channel_registry.new_sender(sub.get_channel_name())
+                    ] = self._channel_registry.get_or_create(
+                        _Report, sub.get_channel_name()
+                    ).new_sender()
 
                 if sub.component_ids not in self._bound_tracker_tasks:
                     self._add_bounds_tracker(sub.component_ids)

src/frequenz/sdk/actor/_resampling.py

@@ -77,11 +77,15 @@ async def _subscribe(self, request: ComponentMetricRequest) -> None:
         )
         data_source_channel_name = data_source_request.get_channel_name()
         await self._data_sourcing_request_sender.send(data_source_request)
-        receiver = self._channel_registry.new_receiver(data_source_channel_name)
+        receiver = self._channel_registry.get_or_create(
+            Sample[Quantity], data_source_channel_name
+        ).new_receiver()
 
         # This is a temporary hack until the Sender implementation uses
         # exceptions to report errors.
-        sender = self._channel_registry.new_sender(request_channel_name)
+        sender = self._channel_registry.get_or_create(
+            Sample[Quantity], request_channel_name
+        ).new_sender()
 
         async def sink_adapter(sample: Sample[Quantity]) -> None:
             await sender.send(sample)

src/frequenz/sdk/timeseries/_base_types.py

@@ -6,11 +6,10 @@
 import dataclasses
 import enum
 import functools
-import typing
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import Generic, Self, overload
+from typing import Generic, Self, TypeVar, overload
 
 from ._quantities import Power, QuantityT
 
@@ -140,7 +139,7 @@ def map(
         )
 
 
-_T = typing.TypeVar("_T")
+_T = TypeVar("_T")
 
 
 @dataclass(frozen=True)

src/frequenz/sdk/timeseries/_grid_frequency.py

@@ -15,7 +15,7 @@
 from ..microgrid import connection_manager
 from ..microgrid.component import Component, ComponentCategory, ComponentMetricId
 from ..timeseries._base_types import Sample
-from ..timeseries._quantities import Frequency
+from ..timeseries._quantities import Frequency, Quantity
 
 if TYPE_CHECKING:
     # Imported here to avoid a circular import.
@@ -95,9 +95,9 @@ def new_receiver(self) -> Receiver[Sample[Frequency]]:
         Returns:
             A receiver that will receive grid frequency samples.
         """
-        receiver = self._channel_registry.new_receiver(
-            self._component_metric_request.get_channel_name()
-        )
+        receiver = self._channel_registry.get_or_create(
+            Sample[Quantity], self._component_metric_request.get_channel_name()
+        ).new_receiver()
 
         if not self._task:
             self._task = asyncio.create_task(self._send_request())