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, 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 (first 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    Order of calling:
 21    [start tick] -> [tcp recv] -> io_handler (message to events + update ecs data from events) ->
 22                 -> [physics system] -> response_handler -> [tcp send] -> [end tick]
 23    """
 24
 25    def __init__(self, registered_plugins: dict[str, type[Plugin]], plugin_names: list[str]):
 26        Serializable.__init__(self)
 27        self.registered_plugins = registered_plugins
 28        self.plugin_names = plugin_names
 29
 30        # instantiate plugins here!
 31        self.plugins: list[Plugin] = [registered_plugins[name]() for name in plugin_names]
 32        all_endpoints = [e.name for p in self.plugins for e in p.endpoints]
 33        assert len(set(all_endpoints)) == len(all_endpoints), f"Duplicates: {all_endpoints}"
 34        self.all_endpoints = set(all_endpoints)
 35
 36        self._endpoint_name_to_plugin = {e.name: plugin for plugin in self.plugins for e in plugin.endpoints}
 37        self._protocol = Protocol(endpoints={e.name: e for p in self.plugins for e in p.endpoints}) # formal validation
 38        self._events_this_tick: dict[EntityId, list[Event]] = {}
 39
 40    def on_after_physics(self, world: World):
 41        """Calls all the plugins after the physics loop. All events must be responded to after this phase"""
 42        [plugin.on_after_physics(world) for plugin in self.plugins]
 43
 44    def on_before_physics(self, world: World):
 45        """Calls all the plugins before the physics loop"""
 46        [plugin.on_before_physics(world) for plugin in self.plugins]
 47
 48    def io_handler(self, world: World, messages_for_plugins: list[Message]):
 49        """Called at the beginning of each render tick. Sets events_this_tick {eid -> [events]} and modifies ECS data"""
 50        # group messages by plugins
 51        messages: dict[Plugin, list[Message]] = {}
 52        for msg in messages_for_plugins:
 53            assert "cmd" in msg.data, msg
 54            client: RobosimClient = msg.client
 55            channel = client.channel
 56            cmd = msg.data["cmd"]
 57
 58            if (plugin := self._endpoint_name_to_plugin.get(cmd)) is None:
 59                logger.error(err := f"Unknown message: {msg} ({cmd=}) from channel: {channel.idx}")
 60                channel.main2tcp.put({"error": err})
 61                continue
 62
 63            # validate the payload using microspec formal validation.
 64            if (err := self._protocol.validate_endpoint(endpoint=cmd, data=msg.data)) is not None:
 65                logger.error(err)
 66                channel.main2tcp.put({"error": err.error})
 67                continue
 68
 69            # TODO: some messages are ok to be sent by not connected clients (e.g. mission_get_state)
 70            if client.state != RobosimClientState.CONNECTED:
 71                logger.error(err := f"Client from channel: {channel.idx} not connected")
 72                channel.main2tcp.put({"error": err})
 73                continue
 74
 75            # the plugin message is validated, so no syntax sanitization is needed inside the plugin
 76            messages.setdefault(plugin, []).append(msg)
 77
 78        # Then, after grouping, we go through all the messages and call on_message_receive which is used for entity evts
 79        self._events_this_tick.clear()
 80        n_events = 0
 81        for plugin in self.plugins:
 82            # first, we go through all the messages of this plugin (if any) and create ECS events group by entityt id
 83            for message in messages.get(plugin, []):
 84                client: RobosimClient = message.client
 85                plg_response = plugin.on_message_receive(world, entity_id=client.robot_eid, message=message)
 86                # for physics responses, the even is further relayed to the ECS and processed by physics system.
 87                if plg_response.kind == ResponseKind.Physics:
 88                    event_data = {"payload": plg_response.payload}
 89                else:
 90                    event_data = {"response": plg_response.payload}
 91                event = Event(source=plugin, target=client.robot_eid, message=message, **event_data)
 92                self._events_this_tick.setdefault(client.robot_eid, []).append(event)
 93                n_events += 1
 94
 95            # Then, we call on_tick on this plugin after it aggregated all the data
 96            for event in plugin.on_tick(world):
 97                self._events_this_tick.setdefault(event.target, []).append(event)
 98                n_events += 1
 99
100        logger.log_every_s(f"Processing {n_events} across {len(self._events_this_tick)} ents this tick", "TRACE", True)
101
102        # Update ECS data from events
103        for entity_id, entity_events in self._events_this_tick.items():
104            for event in reversed(entity_events):
105                if event.payload is not None: # the first event targeting this entity wins this tick
106                    world.get_entity(entity_id).set_data(**event.payload) # the old e.k[:]=v in one step
107                    event.applied = True
108                    break # only the 1st event targetting the entity wins this tick (ordering matters!)
109        world.update()
110
111    def responses_handler(self, world: World):
112        """Called at the end of each render tick to ensure all the received messages have their responses"""
113        for entity_events in self._events_this_tick.values():
114            for event in entity_events:
115                # events w/o a message are generated via on_tick and they are internal (e.g. wind or default moves)
116                if event.message is None:
117                    continue
118                # respond to this unresponded-yet event (e.g. from on_message_receive)
119                if event.response is None:
120                    plg_response = event.source.on_message_response(world, event)
121                    assert plg_response.kind in (ResponseKind.Ok, ResponseKind.Error), plg_response
122                    assert plg_response.payload is not None, plg_response
123                    event.response = plg_response.payload
124                event.message.client.channel.main2tcp.put(event.response)
125
126    def to_dict(self) -> dict:
127        return {name: plugin.to_dict() for name, plugin in zip(self.plugin_names, self.plugins)}
128
129    # pylint: disable=arguments-differ
130    @staticmethod
131    def from_dict(registered_plugins: dict[str, type[Plugin]], state: dict) -> PluginsManager:
132        plugin_names = []
133        for plugin_name in state.keys():
134            if plugin_name not in registered_plugins:
135                logger.warning(f"Plugin: '{plugin_name}' not in {list(registered_plugins)}. Skipping.")
136                continue
137            plugin_names.append(plugin_name)
138
139        res = PluginsManager(registered_plugins, plugin_names)
140        for plugin_name, plugin in zip(res.plugin_names, res.plugins):
141            plugin.load_state_dict(state[plugin_name])
142        return res
143
144    def __repr__(self):
145        plugins = "".join(f"\n  - {type(s).__name__}. Endpoints: {s.endpoints}" for s in self.plugins)
146        return f"[PluginsManager]\n- Endpoints: {len(self.all_endpoints)}\n- Plugins ({len(self.plugins)}):{plugins}"