refactor(actor)!: Resolve and cache handlers by message type

src/powerapi/actor/actor.py

@@ -137,7 +137,7 @@ def _process_received_messages(self) -> None:
             return  # Timeout
 
         logging.debug('Received message: %s', msg)
-        self.state.dispatch_message(msg)
+        self.state.handle_message(msg)
 
     def _teardown_actor(self) -> None:
         """

src/powerapi/actor/handler.py

@@ -119,4 +119,4 @@ def handle(self, msg: PoisonPillMessage) -> None:
 
         self.state.alive = False
         while (pending_msg := self.state.actor.socket_interface.receive(timeout=100)) is not None:
-            self.state.dispatch_message(pending_msg)
+            self.state.handle_message(pending_msg)

src/powerapi/actor/state.py

@@ -54,9 +54,11 @@ def __init__(self, actor):
         self.initialized = False
         self.alive = True
 
-        self.handlers: dict[str, Handler] = {}
+        self.handlers: dict[type[Message], Handler] = {}
         self.supervisor = Supervisor()
 
+        self._resolved_handlers: dict[type[Message], Handler] = {}
+
     def initialize(self) -> None:
         """
         Initialize resources after the actor receives a start message.
@@ -69,30 +71,48 @@ def teardown(self) -> None:
 
     def get_corresponding_handler(self, msg: Message) -> Handler:
         """
-        Return the corresponding handler for the given message type.
+        Return the most specific registered handler, caching successful lookups.
         :param msg: The message
         :return: The handler for the given message type
         :raises KeyError: If the message type does not have a corresponding handler
         """
-        return self.handlers[msg.__class__.__name__]
+        message_type = type(msg)
+        try:
+            return self._resolved_handlers[message_type]
+        except KeyError:
+            return self._resolve_handler(message_type)
+
+    def _resolve_handler(self, message_type: type[Message]) -> Handler:
+        """
+        Resolve a handler in method resolution order and cache the successful lookup.
+        :param message_type: The message class to resolve
+        :return: The most specific registered handler
+        :raises KeyError: If no handler is registered for the class or its bases
+        """
+        for base in message_type.__mro__:
+            handler = self.handlers.get(base)
+            if handler is not None:
+                self._resolved_handlers[message_type] = handler
+                return handler
 
-    def add_handler(self, message_type: type[Message], handler: Handler, include_subclasses: bool = True) -> None:
+        raise KeyError(message_type)
+
+    def add_handler(self, message_type: type[Message], handler: Handler) -> None:
         """
-        Add a handler for the given message type.
+        Register a handler for a message class and its subclasses.
+        Exact registrations take precedence over base classes in method resolution order.
+        Registering the same class again replaces its handler.
+        Every registration clears the handler resolution cache.
         :param message_type: The message type
         :param handler: The corresponding handler
-        :param include_subclasses: Whether to include subclasses of the message type
         """
-        self.handlers[message_type.__name__] = handler
-
-        if include_subclasses:
-            for child_type in message_type.__subclasses__():
-                self.handlers[child_type.__name__] = handler
+        self.handlers[message_type] = handler
+        self._resolved_handlers.clear()
 
-    def dispatch_message(self, msg: Message) -> None:
+    def handle_message(self, msg: Message) -> None:
         """
-        Dispatch a message to its registered handler.
-        :param msg: Message to dispatch
+        Handle a message using its registered handler.
+        :param msg: Message to handle
         """
         try:
             handler = self.get_corresponding_handler(msg)

tests/unit/actor/test_actor.py

@@ -68,12 +68,6 @@ def __eq__(self, other):
         return self.value == other.value
 
 
-class DummyMessageSubtype(DummyMessage):
-    """
-    Subtype of a dummy message class.
-    """
-
-
 class FailingStartMessageHandler(Handler[StartMessage]):
     """
     Failing start message handler.
@@ -96,17 +90,12 @@ class LoopbackMessageHandler(Handler[DummyMessage]):
     Set the processed flag to received messages and sends them back to the control channel.
     """
 
-    def handle(self, message: DummyMessage) -> None:
+    def handle(self, msg: DummyMessage) -> None:
         """
         Handle a message.
         """
-        message.processed = True
-        self.state.actor.send_control(message)
-
-
-class KeyErrorMessageHandler(Handler[DummyMessage]):
-    def handle(self, message: DummyMessage) -> None:
-        raise KeyError('handler failure')
+        msg.processed = True
+        self.state.actor.send_control(msg)
 
 
 class DummyActorState(State):
@@ -232,66 +221,6 @@ def test_send_data_message_to_actor(started_loopback_actor):
     assert ack_msg.processed is True
 
 
-def test_retrieve_handler_for_known_message_type():
-    """
-    Test retrieving the corresponding handler for a known message type.
-    """
-    state = DummyActorState(None)
-    dummy_handler = LoopbackMessageHandler(state)
-    state.add_handler(DummyMessage, dummy_handler)
-
-    msg_handler = state.get_corresponding_handler(DummyMessage('test-dummy'))
-    assert msg_handler is dummy_handler
-
-
-def test_retrieve_handler_for_known_message_subtype():
-    """
-    Test retrieving the corresponding handler for a subtype of a known message type.
-    """
-    state = DummyActorState(None)
-    dummy_handler = LoopbackMessageHandler(state)
-    state.add_handler(DummyMessage, dummy_handler, include_subclasses=True)
-
-    msg_handler = state.get_corresponding_handler(DummyMessageSubtype('test-dummy-subtype'))
-    assert msg_handler is dummy_handler
-
-
-def test_retrieve_handler_for_known_message_with_subtype_excluded():
-    """
-    Test that retrieving the corresponding handler for an excluded subtype of a known message type raises an error.
-    """
-    state = DummyActorState(None)
-    dummy_handler = LoopbackMessageHandler(state)
-    state.add_handler(DummyMessage, dummy_handler, include_subclasses=False)
-
-    with pytest.raises(KeyError, match='DummyMessageSubtype'):
-        state.get_corresponding_handler(DummyMessageSubtype('test-dummy-subtype'))
-
-
-def test_retrieve_handler_for_unknown_message_type():
-    """
-    Test that retrieving the corresponding handler for an unknown message type raises an error.
-    """
-    state = DummyActorState(None)
-
-    with pytest.raises(KeyError, match='DummyMessage'):
-        state.get_corresponding_handler(DummyMessage('test-dummy'))
-
-
-def test_handler_key_error_is_not_mistaken_for_an_unknown_message():
-    """
-    Exceptions raised by handlers should not be treated as unknown messages.
-    """
-    actor = Actor('pytest')
-    state = DummyActorState(actor)
-    actor.state = state
-    handler = KeyErrorMessageHandler(state)
-    state.add_handler(DummyMessage, handler)
-
-    with pytest.raises(KeyError, match='handler failure'):
-        state.dispatch_message(DummyMessage('test-dummy'))
-
-
 def test_start_message_handler_initializes_state():
     """
     The start-message handler should initialize the actor state.
@@ -319,7 +248,7 @@ def test_poison_pill_dispatches_pending_messages_before_state_teardown():
 
     handler.handle(PoisonPillMessage())
 
-    state.dispatch_message.assert_called_once_with(pending_message)
+    state.handle_message.assert_called_once_with(pending_message)
     state.teardown.assert_not_called()
     assert state.alive is False