robolib.plugins_manager

  1"""
  2plugins_manager.py A list of plugins and the orchestration of events between them.
  3Import level: 4
  4"""
  5from __future__ import annotations
  6from microecs import EntityId, World
  7from microspec import Protocol
  8from micronetcode import Message
  9
 10from robolib.utils import logger
 11from robolib.plugin import Plugin, Event, Response, ResponseKind
 12from robolib.traits import Serializable
 13from robolib.netcode import RobosimClient, RobosimClientState
 14
 15class PluginsManager(Serializable):
 16    """
 17    Manager on top of all loaded plugins. It routes TCP (protocol) messages to and from each plugin. Also handles the
 18    ordering in which events are applied (last wins). The order of plugins determine priority in case >=2 target the
 19    same behavior, e.g. manual move plugin vs automatic trajectory plugin.
 20    Note: this class doesn't do any channel (queue) transmission. Both relevant public methods (io_handler and
 21    response_handler) return the replies of the plugins (or errors). The queueing is done in the main loop.
 22    Order of calling (bracketed ones happen outside, inlcuding channel response):
 23    [start tick] -> [tcp recv] -> io_handler (message to events + update ecs data from events) ->
 24                 -> [physics system] -> response_handler -> [tcp send] -> [end tick]
 25    """
 26
 27    def __init__(self, registered_plugins: dict[str, type[Plugin]], plugin_names: list[str]):
 28        Serializable.__init__(self)
 29        self.registered_plugins = registered_plugins
 30        self.plugin_names = plugin_names
 31
 32        # Instantiate plugins here!
 33        self.plugins: list[Plugin] = [registered_plugins[name]() for name in plugin_names]
 34        all_endpoints = [e.name for p in self.plugins for e in p.endpoints]
 35        assert len(set(all_endpoints)) == len(all_endpoints), f"Duplicates: {all_endpoints}"
 36        self.all_endpoints = set(all_endpoints)
 37
 38        self._endpoint_name_to_plugin = {e.name: plugin for plugin in self.plugins for e in plugin.endpoints}
 39        # Protocol validation via microspec
 40        self._protocol = Protocol(endpoints={e.name: e for p in self.plugins for e in p.endpoints})
 41        # Internal buffer between io_handler and response_handler
 42        self._events_this_tick: dict[EntityId, list[Event]] = {}
 43
 44    def on_after_physics(self, world: World):
 45        """Calls all the plugins after the physics loop. All events must be responded to after this phase"""
 46        [plugin.on_after_physics(world) for plugin in self.plugins]
 47
 48    def on_before_physics(self, world: World):
 49        """Calls all the plugins before the physics loop"""
 50        [plugin.on_before_physics(world) for plugin in self.plugins]
 51
 52    def io_handler(self, world: World, messages_for_plugins: list[Message]) -> list[tuple[Message, Response]]:
 53        """
 54        Called at the beginning of each render tick. Sets events_this_tick {eid -> [events]} and modifies
 55        ECS data (TODO-196). Returns the replies of the messages rejected before reaching a plugin (unknown endpoint,
 56        failed validation, client not connected). Accepted messages are answered later, by responses_handler. Must be
 57        called exactly once per tick: it clears events_this_tick, so a second call drops the first's events.
 58        """
 59        # group messages by plugins
 60        messages: dict[Plugin, list[Message]] = {}
 61        early_responses: list[tuple[Message, Response]] = []
 62        for msg in messages_for_plugins:
 63            client: RobosimClient = msg.client
 64            cmd = msg.data["cmd"]
 65            err = None
 66
 67            if (plugin := self._endpoint_name_to_plugin.get(cmd)) is None:
 68                err = f"Unknown message: {msg} ({cmd=}) from channel: {client.channel.idx}"
 69            # Validate the payload using microspec.
 70            elif (valid_err := self._protocol.validate_endpoint(endpoint=cmd, data=msg.data)) is not None:
 71                err = valid_err.error
 72            # Plugin messages are assumed to only work for CONNECTED clients. (TODO-197)
 73            elif client.state != RobosimClientState.CONNECTED:
 74                err = f"Client from channel: {client.channel.idx} not connected"
 75
 76            if err is not None:
 77                logger.warning(err)
 78                early_responses.append((msg, Response.err(err)))
 79                continue
 80
 81            # the plugin message is validated, so no syntax sanitization is needed inside the plugin
 82            messages.setdefault(plugin, []).append(msg)
 83
 84        # Then, after grouping, we go through all the messages and call on_message_receive which is used for entity evts
 85        self._events_this_tick.clear()
 86        n_events = 0
 87        for plugin in self.plugins:
 88            # first, we go through all the messages of this plugin (if any) and create ECS events group by entityt id
 89            for message in messages.get(plugin, []):
 90                client: RobosimClient = message.client
 91                plg_resp = plugin.on_message_receive(world, entity_id=client.robot_eid, message=message)
 92                # for physics responses, the even is further relayed to the ECS and processed by physics system.
 93                if plg_resp.kind == ResponseKind.Physics:
 94                    event = Event(source=plugin, target=client.robot_eid, message=message, payload=plg_resp.payload)
 95                else:
 96                    event = Event(source=plugin, target=client.robot_eid, message=message, response=plg_resp)
 97                self._events_this_tick.setdefault(client.robot_eid, []).append(event)
 98                n_events += 1
 99
100            # Then, we call on_tick on this plugin after it aggregated all the data
101            for event in plugin.on_tick(world):
102                self._events_this_tick.setdefault(event.target, []).append(event)
103                n_events += 1
104
105        logger.log_every_s(f"Processing {n_events} across {len(self._events_this_tick)} ents this tick", "TRACE", True)
106
107        # Update ECS data from events
108        for entity_id, entity_events in self._events_this_tick.items():
109            for event in reversed(entity_events):
110                if event.payload is not None: # the last event targeting this entity wins this tick
111                    world.get_entity(entity_id).set_data(**event.payload) # the old e.k[:]=v in one step
112                    event.applied = True
113                    break # only the last event targetting the entity wins this tick (ordering matters!)
114
115        return early_responses
116
117    def responses_handler(self, world: World) -> list[tuple[Message, Response]]:
118        """Called at the end of each render tick to ensure all the received messages have their responses"""
119        responses = []
120        for entity_events in self._events_this_tick.values():
121            for event in entity_events:
122                # events w/o a message are generated via on_tick and they are internal (e.g. wind or default moves)
123                if event.message is None:
124                    continue
125                # respond to this unresponded-yet event (e.g. from on_message_receive)
126                if event.response is None:
127                    event.response = event.source.on_message_response(world, event)
128                    assert event.response.kind in (ResponseKind.Ok, ResponseKind.Error), event.response
129                    assert event.response.payload is not None, event.response
130                responses.append((event.message, event.response))
131        return responses
132
133    def to_dict(self) -> dict:
134        return {name: plugin.to_dict() for name, plugin in zip(self.plugin_names, self.plugins)}
135
136    # pylint: disable=arguments-differ
137    @staticmethod
138    def from_dict(registered_plugins: dict[str, type[Plugin]], state: dict) -> PluginsManager:
139        plugin_names = []
140        for plugin_name in state.keys():
141            if plugin_name not in registered_plugins:
142                logger.warning(f"Plugin: '{plugin_name}' not in {list(registered_plugins)}. Skipping.")
143                continue
144            plugin_names.append(plugin_name)
145
146        res = PluginsManager(registered_plugins, plugin_names)
147        for plugin_name, plugin in zip(res.plugin_names, res.plugins):
148            plugin.load_state_dict(state[plugin_name])
149        return res
150
151    def __repr__(self):
152        plugins = "".join(f"\n  - {type(s).__name__}. Endpoints: {s.endpoints}" for s in self.plugins)
153        return f"[PluginsManager]\n- Endpoints: {len(self.all_endpoints)}\n- Plugins ({len(self.plugins)}):{plugins}"