Move ChannelRegistry implementation to _internal/_channels.py

This is a stand-alone class and users shouldn't be tied to the `actor`
module.

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

Author: shsms

benchmarks/timeseries/benchmark_datasourcing.py

@@ -20,11 +20,8 @@
 from frequenz.client.microgrid import ComponentMetricId
 
 from frequenz.sdk import microgrid
-from frequenz.sdk.actor import (
-    ChannelRegistry,
-    ComponentMetricRequest,
-    DataSourcingActor,
-)
+from frequenz.sdk._internal._channels import ChannelRegistry
+from frequenz.sdk.actor import ComponentMetricRequest, DataSourcingActor
 
 try:
     from tests.timeseries.mock_microgrid import MockMicrogrid

src/frequenz/sdk/_internal/_channels.py

@@ -4,12 +4,18 @@
 """General purpose classes for use with channels."""
 
 import abc
+import dataclasses
+import logging
+import traceback
 import typing
 
-from frequenz.channels import Receiver
+from frequenz.channels import Broadcast, Receiver
+
+_logger = logging.getLogger(__name__)
 
 T_co = typing.TypeVar("T_co", covariant=True)
 U_co = typing.TypeVar("U_co", covariant=True)
+T = typing.TypeVar("T")
 
 
 class ReceiverFetcher(typing.Generic[T_co], typing.Protocol):
@@ -55,3 +61,141 @@ def new_receiver(self, *, limit: int = 50) -> Receiver[U_co]:
             A receiver instance.
         """
         return self._mapping_function(self._fetcher.new_receiver(limit=limit))
+
+
+class ChannelRegistry:
+    """Dynamically creates, own and provide access to broadcast channels.
+
+    It can be used by actors to dynamically establish a communication channel
+    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:
+        """Initialize this registry.
+
+        Args:
+            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, _Entry] = {}
+
+    @property
+    def name(self) -> str:
+        """The name of this registry."""
+        return self._name
+
+    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.
+
+        Returns:
+            The message type of the channel.
+
+        Raises:
+            KeyError: If the channel does not exist.
+        """
+        entry = self._channels.get(key)
+        if entry is None:
+            raise KeyError(f"No channel for key {key!r} exists.")
+        return entry.message_type
+
+    def __contains__(self, key: str) -> bool:
+        """Check whether the channel for the given `key` exists."""
+        return key in self._channels
+
+    def get_or_create(self, message_type: type[T], key: str) -> Broadcast[T]:
+        """Get or create a channel for 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.
+
+        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:
+            message_type: The type of the message that is sent through the channel.
+            key: The key to identify the channel.
+
+        Returns:
+            The channel for the given key.
+
+        Raises:
+            ValueError: If the channel exists and the message type does not match.
+        """
+        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(name=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 typing.cast(Broadcast[T], entry.channel)
+
+    async def close_and_remove(self, key: str) -> None:
+        """Remove the channel for the given key.
+
+        Args:
+            key: The key to identify the channel.
+
+        Raises:
+            KeyError: If the channel does not exist.
+        """
+        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/__init__.py

@@ -599,7 +599,6 @@ async def main() -> None:  # (6)!
 from ..timeseries._resampling import ResamplerConfig
 from ._actor import Actor
 from ._background_service import BackgroundService
-from ._channel_registry import ChannelRegistry
 from ._config_managing import ConfigManagingActor
 from ._data_sourcing import ComponentMetricId, ComponentMetricRequest, DataSourcingActor
 from ._resampling import ComponentMetricsResamplingActor
@@ -608,7 +607,6 @@ async def main() -> None:  # (6)!
 __all__ = [
     "Actor",
     "BackgroundService",
-    "ChannelRegistry",
     "ComponentMetricId",
     "ComponentMetricRequest",
     "ComponentMetricsResamplingActor",

src/frequenz/sdk/actor/_channel_registry.py

@@ -5,8 +5,8 @@
 
 from frequenz.channels import Receiver
 
+from ..._internal._channels import ChannelRegistry
 from .._actor import Actor
-from .._channel_registry import ChannelRegistry
 from ._component_metric_request import ComponentMetricRequest
 from .microgrid_api_source import MicrogridApiSource
 

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

@@ -18,10 +18,10 @@
     MeterData,
 )
 
+from ..._internal._channels import ChannelRegistry
 from ...microgrid import connection_manager
 from ...timeseries import Sample
 from ...timeseries._quantities import Quantity
-from .._channel_registry import ChannelRegistry
 from ._component_metric_request import ComponentMetricRequest
 
 _logger = logging.getLogger(__name__)

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

@@ -16,10 +16,10 @@
 from frequenz.client.microgrid import ComponentCategory, ComponentType, InverterType
 from typing_extensions import override
 
+from ..._internal._channels import ChannelRegistry
 from ...timeseries import Power
 from ...timeseries._base_types import Bounds, SystemBounds
 from .._actor import Actor
-from .._channel_registry import ChannelRegistry
 from ._base_classes import Algorithm, BaseAlgorithm, Proposal, ReportRequest, _Report
 from ._matryoshka import Matryoshka
 

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

@@ -11,11 +11,11 @@
 from frequenz.channels import Receiver, Sender
 
 from .._internal._asyncio import cancel_and_await
+from .._internal._channels import ChannelRegistry
 from ..timeseries import Sample
 from ..timeseries._quantities import Quantity
 from ..timeseries._resampling import Resampler, ResamplerConfig, ResamplingError
 from ._actor import Actor
-from ._channel_registry import ChannelRegistry
 from ._data_sourcing import ComponentMetricRequest
 
 _logger = logging.getLogger(__name__)