diff --git a/dimos/core/baked_host.py b/dimos/core/baked_host.py index e058d0ce7e..e4ef793040 100644 --- a/dimos/core/baked_host.py +++ b/dimos/core/baked_host.py @@ -24,7 +24,7 @@ from pydantic import Field, create_model -from dimos.core.native_module import NativeModule, NativeModuleConfig +from dimos.core.native_module import NativeModule, NativeModuleConfig, TopicsMap, TopicValue from dimos.core.stream import IO, In, Out @@ -83,23 +83,23 @@ class BakedHost(NativeModule): _members: Mapping[str, type[NativeModule]] = {} _remaps: Mapping[tuple[str, str], str] = {} - def _member_topics(self, instance: str, topics: Mapping[str, str]) -> dict[str, str]: + def _member_topics(self, instance: str, topics: TopicsMap) -> dict[str, TopicValue]: member = self._members[instance] - resolved = {} + resolved: dict[str, TopicValue] = {} for port in _member_ports(member): name = self._remaps.get((instance, port), port) if name in topics: resolved[port] = topics[name] return resolved - def _argv(self, topics: dict[str, str]) -> list[str]: + def _argv(self, topics: TopicsMap) -> list[str]: """A baked host takes its whole wiring on the launch line, so no flags. The binary rejects an unknown argument rather than ignoring it. """ return [self.config.executable, *self.config.extra_args] - def _stdin_blob(self, topics: dict[str, str]) -> bytes: + def _stdin_blob(self, topics: TopicsMap) -> bytes: sections: dict[str, object] = {} for instance in self._members: member_config = getattr(self.config, f"{instance}_config") diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index b02453e7f5..945819f5db 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -27,7 +27,7 @@ if TYPE_CHECKING: from dimos.protocol.service.system_configurator.base import SystemConfigurator -from dimos.core.module import ModuleBase, is_module_type +from dimos.core.module import ModuleBase, TopicFunnel, is_module_type from dimos.core.stream import IO, In, Out, Transport from dimos.spec.utils import Spec, is_spec from dimos.utils.logging_config import setup_logger @@ -231,13 +231,28 @@ def global_config(self, **kwargs: Any) -> "Blueprint": def remappings( self, remappings: Sequence[ - tuple[type[ModuleBase] | str, str, str | type[ModuleBase] | type[Spec]] + tuple[ + type[ModuleBase] | str, + str, + str + | Sequence[str] + | Mapping[str, Mapping[str, Any]] + | type[ModuleBase] + | type[Spec], + ] ], ) -> "Blueprint": remappings_dict = dict(self.remapping_map) + atoms = list(self.blueprints) for module, old, new in remappings: - remappings_dict[(self._instance_key(module), old)] = new - return replace(self, remapping_map=MappingProxyType(remappings_dict)) + instance_key = self._instance_key(module) + if isinstance(new, (str, type)): + remappings_dict[instance_key, old] = new + else: + atoms = _fan_in(atoms, instance_key, old, TopicFunnel(new)) + return replace( + self, blueprints=tuple(atoms), remapping_map=MappingProxyType(remappings_dict) + ) def _instance_key(self, module: type[ModuleBase] | str) -> str: if isinstance(module, str): @@ -381,6 +396,54 @@ def autoconnect(*blueprints: Blueprint) -> Blueprint: ) +def _fan_in( + atoms: list[BlueprintAtom], instance_key: str, port: str, funnel: TopicFunnel +) -> list[BlueprintAtom]: + """Point one declared input at several streams instead of one. + + The port stops being a stream of its own; each entry takes its place as an + `In` of the same type, so autoconnect, `.namespace()` and transport pins + treat the entries as ordinary streams. The module keeps the port as the + single place messages from all of them arrive. + """ + updated = [] + matched = False + for atom in atoms: + declared = next( + ( + stream + for stream in atom.streams + if atom.name == instance_key + and stream.name == port + and stream.direction in ("in", "inout") + ), + None, + ) + if declared is None: + updated.append(atom) + continue + matched = True + entries = tuple( + StreamRef(name=name, type=declared.type, direction="in") for name in funnel.names + ) + updated.append( + replace( + atom, + kwargs={ + **atom.kwargs, + "topic_funnels": {**atom.kwargs.get("topic_funnels", {}), port: funnel}, + }, + streams=tuple(s for s in atom.streams if s is not declared) + entries, + ) + ) + if not matched: + raise ValueError( + f"cannot fan {port!r} in: {instance_key} declares no such In/IO stream " + "in this blueprint (already fanned in, or a typo?)" + ) + return updated + + def _eliminate_duplicates(blueprints: list[BlueprintAtom]) -> list[BlueprintAtom]: # The duplicates are eliminated in reverse so that newer blueprints override older ones. seen = set() diff --git a/dimos/core/module.py b/dimos/core/module.py index f5aab99421..88d8a1666e 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio -from collections.abc import AsyncGenerator, Callable -from dataclasses import dataclass +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from dataclasses import dataclass, field from functools import partial import inspect import json @@ -41,7 +41,7 @@ from dimos.core.resource import CompositeResource from dimos.core.rpc_client import RpcCall from dimos.core.stream import IO, In, Out, RemoteOut, Transport -from dimos.core.transport_factory import rpc_backend +from dimos.core.transport_factory import make_transport, rpc_backend from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, RPCSpec from dimos.protocol.service.spec import BaseConfig, Configurable from dimos.protocol.tf.tf import TF @@ -103,6 +103,76 @@ def get_loop() -> tuple[asyncio.AbstractEventLoop, threading.Thread | None]: Deployment = Literal["python", "docker"] +@dataclass(init=False) +class TopicFunnel: + """Several same-typed streams that one declared port fans into one handler. + + Built by `Blueprint.remappings()`, not by blueprint authors:: + + .remappings([ + (Fuser, "scan", ["front_lidar", "rear_lidar"]), + (Fuser, "image", {"depth": {"meters_per_unit": 0.001}, "color": {}}), + ]) + """ + + names: list[str] + info: dict[str, dict[str, Any]] + + def __init__(self, entries: "Sequence[str] | Mapping[str, Mapping[str, Any]]") -> None: + self.names = list(entries) + self.info = ( + {name: dict(value) for name, value in entries.items() if value} + if isinstance(entries, Mapping) + else {} + ) + leading_slash = [name for name in self.names if name.startswith("/")] + if leading_slash: + raise ValueError( + f"topic funnel names must be stream names, not topics: {leading_slash} " + "start with '/' (use `left_cam`, not `/left_cam`)" + ) + + +@dataclass(frozen=True) +class TopicMetadata: + """Which stream a message arrived on, for handlers that take `(msg, meta)`. + + `index` is the position in the funnel's `names` (0 for a plain input); + `name` is the stream name as the module declared it, pre-remapping; + `info` is whatever the funnel's `TopicFunnel.info` attached to that name. + """ + + index: int + name: str + info: dict[str, Any] = field(default_factory=dict) + + +def _handler_wants_metadata(handler: Callable[..., Any], label: str) -> bool: + """True if the bound handler takes `(msg, meta)` rather than just `(msg)`.""" + count = len(inspect.signature(handler).parameters) + if count not in (1, 2): + raise TypeError(f"{label} must take (msg) or (msg, meta), not {count} parameters") + return count == 2 + + +def _dispatch_to( + handler: Callable[..., Any], metas: "list[TopicMetadata] | None" +) -> Callable[[int, Any], Any]: + """Adapt a handler to the `(index, msg)` the dispatcher calls with. + + `metas` is None when the handler takes just the message; otherwise it holds + one `TopicMetadata` per stream feeding the port, indexed the same way. + """ + + async def invoke(index: int, msg: Any) -> None: + if metas is None: + await handler(msg) + else: + await handler(msg, metas[index]) + + return invoke + + class ModuleConfig(BaseConfig): rpc_transport: type[RPCSpec] = Field(default_factory=rpc_backend) default_rpc_timeout: float = DEFAULT_RPC_TIMEOUT @@ -113,6 +183,9 @@ class ModuleConfig(BaseConfig): # once (see BlueprintAtom.instance_name). Changes the RPC topic prefix # from the class name to this name. instance_name: str | None = None + # Declared In/IO port -> the streams that port's single handler receives. + # Set by `Blueprint.remappings()`, not by blueprint authors. + topic_funnels: dict[str, TopicFunnel] = Field(default_factory=dict) g: GlobalConfig = global_config @@ -146,9 +219,11 @@ class ModuleBase(Configurable, CompositeResource): _main_gen: AsyncGenerator[None, None] | None = None _tools: dict[str, Any] _tools_lock: threading.Lock + _funnel_streams: dict[str, In[Any]] def __init__(self, config_args: dict[str, Any]) -> None: super().__init__(**config_args) + self._funnel_streams = self._make_funnel_streams() self._module_closed_lock = threading.Lock() self._tools = {} self._tools_lock = threading.Lock() @@ -638,12 +713,18 @@ def _stop_main(self) -> None: def _auto_bind_handlers(self) -> None: """ For each declared `x: In[T]` or `x: IO[T]`, if `async def handle_x` exists, - subscribe it via process_observable so it runs on self._loop. + subscribe it to that port's stream so it runs on self._loop. A port that + `.remappings()` fanned in is subscribed to every stream of its funnel + instead, all reaching the one handler. + + A native module defines no such method — its subprocess subscribes + instead — so this is a no-op there. """ # Validate every handler before subscribing any of them. - bindings: list[tuple[Any, Callable[[Any], Any]]] = [] - for input_name, in_stream in {**self.inputs, **self.ios}.items(): - handler = getattr(self, f"handle_{input_name}", None) + bindings: list[tuple[list[In[Any] | IO[Any]], Callable[[int, Any], Any]]] = [] + ports: dict[str, In[Any] | IO[Any]] = {**self.inputs, **self.ios} + for port, declared in ports.items(): + handler = getattr(self, f"handle_{port}", None) if handler is None: continue # Async @rpc wraps the coroutine fn in a sync dispatcher. Unwrap it @@ -653,18 +734,64 @@ def _auto_bind_handlers(self) -> None: handler = handler.aio.__get__(self, type(self)) if not inspect.iscoroutinefunction(handler): raise TypeError( - f"{type(self).__name__}.handle_{input_name} must be `async def` " + f"{type(self).__name__}.handle_{port} must be `async def` " "(use a manual self..subscribe(...) for sync handlers)" ) - bindings.append((in_stream, handler)) - - for in_stream, handler in bindings: - # process_observable runs each handler through a per-subscription - # dispatcher task on self._loop that serializes invocations and - # keeps only the latest unprocessed message. We subscribe to - # pure_observable() because the dispatcher already provides - # backpressure. - self.process_observable(in_stream.pure_observable(), handler) + funnel = self.config.topic_funnels.get(port) + names = funnel.names if funnel else [port] + streams: list[In[Any] | IO[Any]] = ( + [self._funnel_streams[name] for name in names] if funnel else [declared] + ) + metas = None + if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{port}"): + metas = [ + TopicMetadata( + index=index, name=name, info=funnel.info.get(name, {}) if funnel else {} + ) + for index, name in enumerate(names) + ] + bindings.append((streams, _dispatch_to(handler, metas))) + + for streams, invoke in bindings: + # One dispatcher task per port, on self._loop: it serializes the + # handler and holds only the latest unprocessed message per stream. + on_msg, dispatcher = self._make_keyed_dispatch(invoke) + self.register_disposable(dispatcher) + for index, stream in enumerate(streams): + self.register_disposable(Disposable(stream.subscribe(partial(on_msg, index)))) + + def _make_funnel_streams(self) -> "dict[str, In[Any]]": + """One `In` per topic-funnel entry, keyed by stream name. + + These are wired like declared ports (`set_transport` finds them, the + blueprint machinery remaps/namespaces/pins them), but they are not + attributes, so `inputs` and the native launch line never see them as + ports of their own — the funnel's port stands in for all of them. + + Each starts on the transport its name alone implies, which a coordinator + overwrites when it wires the blueprint and standalone use keeps. + """ + streams: dict[str, In[Any]] = {} + for port, funnel in self.config.topic_funnels.items(): + declared = self.inputs.get(port) or self.ios.get(port) + if declared is None: + raise ValueError( + f"topic funnel port {port!r} is not an In or IO stream of {type(self).__name__}" + ) + for name in funnel.names: + if name != port and name in {**self.inputs, **self.outputs, **self.ios}: + raise ValueError( + f"topic funnel {port!r} entry {name!r} collides with a " + f"declared stream of {type(self).__name__}" + ) + if name in streams: + raise ValueError( + f"topic funnel {port!r} entry {name!r} appears in more than one funnel" + ) + stream: In[Any] = In(declared.type, name, self) + stream.transport = make_transport(name, declared.type) + streams[name] = stream + return streams def _make_async_dispatch( self, async_handler: Callable[[Any], Any] @@ -680,45 +807,60 @@ def _make_async_dispatch( message is kept (LATEST policy). - The returned Disposable cancels the dispatcher task. """ + + on_msg, disposable = self._make_keyed_dispatch(_dispatch_to(async_handler, None)) + return partial(on_msg, 0), disposable + + def _make_keyed_dispatch( + self, async_handler: Callable[[int, Any], Any] + ) -> tuple[Callable[[int, Any], None], "DisposableBase"]: + """`_make_async_dispatch` generalized to a mailbox keyed by sender. + + The mailbox keeps the latest unprocessed message *per key* and drains them + oldest-waiting first, so one chatty topic in a group cannot starve its + siblings the way a single shared slot would. One key is the degenerate case + and behaves exactly like a single LATEST slot. + """ loop = self._loop if loop is None or not loop.is_running(): raise RuntimeError(f"{type(self).__name__}._loop is not running") - async def _bootstrap() -> tuple[asyncio.Event, dict[str, Any], asyncio.Task[None]]: + async def _bootstrap() -> tuple[asyncio.Event, dict[int, Any], asyncio.Task[None]]: event = asyncio.Event() - slot: dict[str, Any] = {"value": None, "has_value": False} + # Insertion-ordered, and re-assigning an existing key keeps its + # position, so a waiting topic holds its place while its value ages up. + pending: dict[int, Any] = {} async def dispatcher() -> None: try: while True: await event.wait() event.clear() - if not slot["has_value"]: - continue - msg = slot["value"] - slot["value"] = None - slot["has_value"] = False - try: - await async_handler(msg) - except asyncio.CancelledError: - raise - except BaseException as e: - self._log_async_handler_exception(e) + while pending: + key = next(iter(pending)) + msg = pending.pop(key) + try: + await async_handler(key, msg) + except asyncio.CancelledError: + raise + except BaseException as e: + self._log_async_handler_exception(e) except asyncio.CancelledError: return - return event, slot, asyncio.create_task(dispatcher()) + return event, pending, asyncio.create_task(dispatcher()) - event, slot, task = asyncio.run_coroutine_threadsafe(_bootstrap(), loop).result(timeout=5.0) + event, pending, task = asyncio.run_coroutine_threadsafe(_bootstrap(), loop).result( + timeout=5.0 + ) - def on_msg(msg: Any) -> None: + def on_msg(key: int, msg: Any) -> None: loop_now = self._loop if loop_now is None or not loop_now.is_running(): return def _set() -> None: - slot["value"] = msg - slot["has_value"] = True + pending[key] = msg event.set() loop_now.call_soon_threadsafe(_set) @@ -811,7 +953,7 @@ def __str__(self) -> str: @rpc def set_transport(self, stream_name: str, transport: Transport) -> bool: # type: ignore[type-arg] - stream = getattr(self, stream_name, None) + stream = self._funnel_streams.get(stream_name) or getattr(self, stream_name, None) if not stream: raise ValueError(f"{stream_name} not found in {self.__class__.__name__}") diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 84fea6b685..94153205b5 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -52,7 +52,7 @@ class MyCppModule(NativeModule): import sys import threading import time -from typing import IO, Any +from typing import IO, Any, TypedDict, cast from pydantic import Field, model_validator @@ -60,10 +60,24 @@ class MyCppModule(NativeModule): from dimos.core.core import rpc from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig +from dimos.core.transport import PubSubTransport from dimos.core.transport_factory import session_config from dimos.protocol.service.spec import SessionConfig from dimos.utils.logging_config import setup_logger + +class TopicEntry(TypedDict): + """A funnel entry carrying that stream's info, mirroring the native `TopicEntry`.""" + + topic: str + info: dict[str, Any] + + +# The wire channel(s) one port is launched with. A funnel port carries a list, +# and an entry with per-topic info is an object rather than a bare string. +TopicValue = str | list[str | TopicEntry] +TopicsMap = dict[str, TopicValue] + if sys.platform.startswith("linux"): import ctypes from ctypes.util import find_library @@ -201,6 +215,9 @@ class NativeModule(Module): ``DIMOS_TRANSPORT`` env var. With ``stdin_config``, the topics, config, publisher QoS and session settings also arrive as one JSON line on stdin. + A port fanned in by ``.remappings()`` gets a list of channels instead of one, + so several same-typed streams reach a single native handler. + The native process should parse whichever it uses and pub/sub on the given topics directly. On ``stop()``, the process receives SIGTERM. """ @@ -258,16 +275,22 @@ def _session(self) -> SessionConfig: # A blueprint builds its config before global config is settled. return pinned.rebased() - def _argv(self, topics: dict[str, str]) -> list[str]: + def _argv(self, topics: TopicsMap) -> list[str]: """The command line the native process is spawned with.""" cmd = [self.config.executable] - for name, topic_str in topics.items(): - cmd.extend([f"--{name}", topic_str]) + for name, value in topics.items(): + if isinstance(value, list): + joined = ",".join( + entry["topic"] if isinstance(entry, dict) else entry for entry in value + ) + else: + joined = value + cmd.extend([f"--{name}", joined]) cmd.extend(self.config.to_cli_args()) cmd.extend(self.config.extra_args) return cmd - def _stdin_blob(self, topics: dict[str, str]) -> bytes: + def _stdin_blob(self, topics: TopicsMap) -> bytes: """The JSON line the native process reads its launch from.""" config_dict = self.config.to_config_dict() blob: dict[str, Any] = { @@ -513,8 +536,8 @@ def _maybe_build(self) -> None: duration_sec=round(build_elapsed, 3), ) - def _collect_topics(self) -> dict[str, str]: - topics: dict[str, str] = {} + def _collect_topics(self) -> TopicsMap: + topics: TopicsMap = {} for name in list(self.inputs) + list(self.outputs) + list(self.ios): stream = getattr(self, name, None) if stream is None: @@ -525,6 +548,13 @@ def _collect_topics(self) -> dict[str, str]: channel = getattr(transport, "channel", None) if channel is not None: topics[name] = channel + for port, funnel in self.config.topic_funnels.items(): + entries: list[str | TopicEntry] = [] + for name in funnel.names: + channel = cast("PubSubTransport[Any]", self._funnel_streams[name].transport).channel + info = funnel.info.get(name) + entries.append(channel if info is None else TopicEntry(topic=channel, info=info)) + topics[port] = entries return topics def _collect_output_qos(self) -> dict[str, dict[str, str]]: diff --git a/dimos/core/test_async_module_topic_funnels.py b/dimos/core/test_async_module_topic_funnels.py new file mode 100644 index 0000000000..a0fa520dcb --- /dev/null +++ b/dimos/core/test_async_module_topic_funnels.py @@ -0,0 +1,152 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from queue import Queue + +import pytest + +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.module import Module, TopicFunnel, TopicMetadata +from dimos.core.stream import In, Out +from dimos.core.transport_factory import make_transport + +DELIVERY_TIMEOUT = 1.0 + + +class FanInModule(Module): + """One handler for two same-typed streams; echoes which one it came from.""" + + sensors: In[int] + tagged: Out[int] + named: Out[str] + scaled: Out[float] + + async def handle_sensors(self, value: int, meta: TopicMetadata) -> None: + self.tagged.publish(meta.index * 100 + value) + self.named.publish(meta.name) + self.scaled.publish(value * meta.info.get("scale", 1.0)) + + +@pytest.fixture +def start_fan_in_module(each_transport): + blueprint = FanInModule.blueprint().remappings( + [(FanInModule, "sensors", {"s0": {"scale": 0.5}, "s1": {}})] + ) + coordinator = ModuleCoordinator.build(blueprint) + yield + coordinator.stop() + + +@pytest.fixture +def group_transports(each_transport): + transports = [ + make_transport("s0"), + make_transport("s1"), + make_transport("tagged"), + make_transport("named"), + make_transport("scaled"), + ] + for transport in transports: + transport.start() + yield transports + for transport in transports: + transport.stop() + + +def test_topic_funnel_tags_each_message_with_its_stream(start_fan_in_module, group_transports): + s0, s1, tagged, named, scaled = group_transports + queue: Queue[int] = Queue() + names: Queue[str] = Queue() + scales: Queue[float] = Queue() + tagged.subscribe(queue.put) + named.subscribe(names.put) + scaled.subscribe(scales.put) + + s0.publish(7) + assert queue.get(timeout=DELIVERY_TIMEOUT) == 7 + assert names.get(timeout=DELIVERY_TIMEOUT) == "s0" + assert scales.get(timeout=DELIVERY_TIMEOUT) == 3.5 + + s1.publish(7) + assert queue.get(timeout=DELIVERY_TIMEOUT) == 107 + assert names.get(timeout=DELIVERY_TIMEOUT) == "s1" + assert scales.get(timeout=DELIVERY_TIMEOUT) == 7.0 + + +@pytest.fixture +def start_remapped_fan_in_module(each_transport): + blueprint = FanInModule.blueprint().remappings( + [(FanInModule, "sensors", ["s0", "s1"]), (FanInModule, "s0", "alt0")] + ) + coordinator = ModuleCoordinator.build(blueprint) + yield + coordinator.stop() + + +def test_topic_funnel_entries_follow_remappings(start_remapped_fan_in_module, each_transport): + """A remapped entry keeps its index and declared name, but listens on the + remapped stream.""" + alt0, tagged, named = make_transport("alt0"), make_transport("tagged"), make_transport("named") + for transport in (alt0, tagged, named): + transport.start() + try: + queue: Queue[int] = Queue() + names: Queue[str] = Queue() + tagged.subscribe(queue.put) + named.subscribe(names.put) + alt0.publish(7) + assert queue.get(timeout=DELIVERY_TIMEOUT) == 7 + assert names.get(timeout=DELIVERY_TIMEOUT) == "s0" + finally: + for transport in (alt0, tagged, named): + transport.stop() + + +def test_fanning_in_an_undeclared_port_is_an_error(): + with pytest.raises(ValueError, match="no such In/IO stream"): + FanInModule.blueprint().remappings([(FanInModule, "nope", ["s0"])]) + + +def test_a_funnel_entry_cannot_collide_with_a_declared_stream(): + with pytest.raises(ValueError, match="collides"): + FanInModule(topic_funnels={"sensors": TopicFunnel(["tagged"])}) + + +class PlainFanInModule(Module): + """A funnel handler that doesn't ask for metadata just gets the message.""" + + sensors: In[int] + echoed: Out[int] + + async def handle_sensors(self, value: int) -> None: + self.echoed.publish(value) + + +def test_a_funnel_handler_without_meta_gets_just_the_message(each_transport): + blueprint = PlainFanInModule.blueprint().remappings( + [(PlainFanInModule, "sensors", ["s0", "s1"])] + ) + coordinator = ModuleCoordinator.build(blueprint) + s1, echoed = make_transport("s1"), make_transport("echoed") + for transport in (s1, echoed): + transport.start() + try: + queue: Queue[int] = Queue() + echoed.subscribe(queue.put) + s1.publish(7) + assert queue.get(timeout=DELIVERY_TIMEOUT) == 7 + finally: + for transport in (s1, echoed): + transport.stop() + coordinator.stop() diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index a58742c3c5..9490140dea 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -35,7 +35,7 @@ from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, TransportBackend -from dimos.core.module import Module +from dimos.core.module import Module, TopicFunnel from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig from dimos.core.stream import IO, In, Out from dimos.core.transport import LCMTransport, ZenohTransport @@ -105,6 +105,12 @@ class StubIoModule(NativeModule): tf: IO[TFMessage] +class StubFunnelModule(NativeModule): + config: StubNativeConfig + cams: In[Imu] + cmd_vel: In[Twist] + + class StubConsumer(Module): pointcloud: In[PointCloud2] imu: In[Imu] @@ -211,6 +217,80 @@ def test_tf_topic_comes_from_the_declared_port_only() -> None: transport.stop() +def test_a_topic_funnel_reaches_the_native_process_as_a_list(monkeypatch) -> None: + """The funnelled port carries every entry's channel instead of one of its own.""" + monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") + module = StubFunnelModule( + executable=_ECHO, + topic_funnels={"cams": TopicFunnel(["cam0/imu", "cam1/imu"])}, + ) + try: + topics = module._collect_topics() + assert topics["cams"] == ["/cam0/imu#sensor_msgs.Imu", "/cam1/imu#sensor_msgs.Imu"] + assert module._argv(topics)[1:3] == [ + "--cams", + "/cam0/imu#sensor_msgs.Imu,/cam1/imu#sensor_msgs.Imu", + ] + finally: + module.stop() + + +def test_funnel_info_rides_the_launch_line_but_not_the_argv(monkeypatch) -> None: + """Per-topic info becomes a {topic, info} entry on stdin; argv keeps only topics.""" + monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") + module = StubFunnelModule( + executable=_ECHO, + topic_funnels={"cams": TopicFunnel({"cam0/imu": {"rectified": True}, "cam1/imu": {}})}, + ) + try: + topics = module._collect_topics() + assert topics["cams"] == [ + {"topic": "/cam0/imu#sensor_msgs.Imu", "info": {"rectified": True}}, + "/cam1/imu#sensor_msgs.Imu", + ] + assert module._argv(topics)[1:3] == [ + "--cams", + "/cam0/imu#sensor_msgs.Imu,/cam1/imu#sensor_msgs.Imu", + ] + finally: + module.stop() + + +def test_a_wired_topic_funnel_entry_uses_its_transport() -> None: + """Remapping/pins arrive as set_transport on the entry, and the launch line follows.""" + module = StubFunnelModule( + executable=_ECHO, + topic_funnels={"cams": TopicFunnel(["cam0/imu"])}, + ) + transport = LCMTransport("/remapped/imu", Imu) + try: + module.set_transport("cam0/imu", transport) + assert module._collect_topics()["cams"] == ["/remapped/imu#sensor_msgs.Imu"] + finally: + module.stop() + with contextlib.suppress(Exception): + transport.stop() + + +def test_a_topic_funnel_needs_a_declared_port() -> None: + """The funnel replaces a port's wiring, so it has to have one to replace.""" + with pytest.raises(ValueError, match="not an In or IO stream"): + StubFunnelModule(executable=_ECHO, topic_funnels={"nope": TopicFunnel(["cam0/imu"])}) + + +def test_a_topic_funnel_is_not_a_native_config_field() -> None: + """The funnel is wiring, so it belongs in `topics`, not the config struct.""" + module = StubFunnelModule( + executable=_ECHO, + topic_funnels={"cams": TopicFunnel(["cam0/imu"])}, + ) + try: + assert "topic_funnels" not in module.config.to_config_dict() + assert "--topic_funnels" not in module._argv({}) + finally: + module.stop() + + def test_io_port_publisher_qos_reaches_the_native_process() -> None: module = StubIoModule(executable=_ECHO) transport = ZenohTransport(ZenohTopic("/tf", TFMessage, qos=QOS_NEVER_DROP)) diff --git a/native/rust/README.md b/native/rust/README.md index 4c31fc1d81..efc9004b88 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -61,7 +61,7 @@ Every transport is compiled into the binary. `run_with_transport` opens the one - `#[derive(Module)]`: on the struct. Required. - `#[module(setup = fn, teardown = fn)]`: on the struct. Both optional. Names methods on `Self`. `setup` runs once before the input dispatch loop starts (use it to spawn background tasks or initialize resources); `teardown` runs once after the loop exits (use it for cleanup). -- `#[input(decode = fn, handler = fn)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. +- `#[input(decode = fn, handler = fn, meta)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. The `meta` flag makes the handler take `(msg, meta: TopicMetadata)`, where `meta` names the topic the message arrived on — useful when the port is a topic funnel (see [Topic funnels](#topic-funnels)). - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. - `#[io(decode = fn, encode = fn, handler = fn)]`: on a field of type `Io`, a port that publishes to and subscribes on one topic. `decode` and `encode` are required; `handler` defaults to `handle_`. The transports deliver a message back to its own sender, so the handler also sees what the module publishes. Use `#[output]` instead when the module only publishes. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. @@ -108,6 +108,52 @@ At runtime `run()` enforces the mapping on the Python payload: deserialization r Field name = port name. Ports map to topics via the stdin JSON; unmapped ports fall back to `/{port}`. +## Topic funnels + +A rig with N identical sensors would otherwise need N ports and N near-identical handlers. A topic funnel is one `Input` port wired to a list of topics that all carry `T`, delivered to one handler in arrival order. Any input accepts a funnel — on the launch line the port's value is an array of topics rather than a string, and nothing in the module changes. A handler that needs to tell the sources apart opts in with the `meta` flag and receives a `TopicMetadata` alongside each message: + +```rust +#[derive(Module)] +struct MultiCam { + #[input(decode = Image::decode, meta)] + cameras: Input, +} + +impl MultiCam { + async fn handle_cameras(&mut self, image: Image, meta: TopicMetadata) { + // meta.index: position in the topic list; meta.topic: the topic itself; + // meta.info: per-topic JSON the coordinator attached (Null when none) + } +} +``` + +Funnelling happens at the blueprint routing level: `.remappings()` takes a list of stream names where it would otherwise take one. + +```python +MultiCam.blueprint().remappings([ + (MultiCam, "cameras", ["left_cam", "right_cam"]), +]) +``` + +A dict instead of a list attaches arbitrary per-stream info — `{"left_cam": {"rectified": True}, "right_cam": {}}` — which reaches the handler as `TopicMetadata.info` (a dict in python, `serde_json::Value` in rust; on the launch line such an entry is a `{"topic": ..., "info": ...}` object instead of a plain string). + +The entries are stream names, not backend topics — a leading `/` is rejected. Each one takes the port's place as an `In` stream of the same type, so autoconnect matches it against producers' `Out` streams, `.namespace()` and further `.remappings()` rewrite it, and blueprint transport pins apply. The same blueprint runs unchanged over LCM or zenoh. + +The port itself stops being a stream of its own: python hands the wired entries' channels to the native process, which subscribes them directly. On the launch line the port's value is an array rather than a string. A funnel with no entries still claims its port but never yields. + +Fan-in is a `ModuleConfig` feature, so a plain Python `Module` funnels the same way. There the module subscribes the entries itself and dispatches to `async def handle_`. Opting into metadata is by signature — a one-parameter handler just gets the message, a two-parameter handler also gets a `TopicMetadata` with `index` (position in the list) and `name` (the stream name as declared, pre-remapping): + +```python +class MultiCam(Module): + cameras: In[Image] + + async def handle_cameras(self, image: Image, meta: TopicMetadata) -> None: ... +``` + +The same applies to any `handle_` for a plain `In` port, where the metadata is always `index=0, name=`. + +Either way the whole funnel shares one dispatcher, so the handler is never re-entered. What they do under load differs: python's mailbox holds the latest unprocessed message per stream and drains them oldest-waiting first, so a chatty camera cannot starve the others, while rust feeds every topic of a funnel into one bounded channel that drops the arriving message once it is full. + ## Transforms A `#[tf]` field gives a module a view of the transform graph, the Rust counterpart to Python's `tf.get()` and `tf.publish()`. It subscribes to the `tf` topic (mapped like any other port, default `/tf`), buffers each `parent -> child` edge it sees, and answers queries by composing transforms along the shortest path through the graph. diff --git a/native/rust/dimos-module-macros/src/lib.rs b/native/rust/dimos-module-macros/src/lib.rs index 94afe00b31..b7d8fcda51 100644 --- a/native/rust/dimos-module-macros/src/lib.rs +++ b/native/rust/dimos-module-macros/src/lib.rs @@ -184,6 +184,7 @@ enum FieldKind { Input { decode: Path, handler: Ident, + wants_meta: bool, }, Output { encode: Path, @@ -336,26 +337,33 @@ fn expand(input: DeriveInput) -> syn::Result { }); // Every port that receives messages gets an arm in the select! loop. - let handled_fields: Vec<(&Ident, &Ident)> = classified + let handle_arms: Vec = classified .iter() - .filter_map(|f| match &f.kind { - FieldKind::Input { handler, .. } | FieldKind::Io { handler, .. } => { - Some((f.name, handler)) + .filter_map(|f| { + let name = f.name; + match &f.kind { + FieldKind::Input { + handler, + wants_meta: true, + .. + } => Some(quote!( + ::core::option::Option::Some((msg, meta)) = self.#name.recv_meta() => { + self.#handler(msg, meta).await + } + )), + FieldKind::Input { handler, .. } | FieldKind::Io { handler, .. } => Some(quote!( + ::core::option::Option::Some(msg) = self.#name.recv() => { + self.#handler(msg).await + } + )), + _ => None, } - _ => None, }) .collect(); - let handle_body = if handled_fields.is_empty() { + let handle_body = if handle_arms.is_empty() { quote!(::std::future::pending::<()>().await) } else { - let handle_arms = handled_fields.iter().map(|(name, handler)| { - quote!( - ::core::option::Option::Some(msg) = self.#name.recv() => { - self.#handler(msg).await - } - ) - }); quote! { loop { ::tokio::select! { @@ -558,6 +566,7 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { } let mut decode: Option = None; let mut handler: Option = None; + let mut wants_meta = false; attr.parse_nested_meta(|meta| { if meta.path.is_ident("decode") { decode = Some(meta.value()?.parse()?); @@ -565,10 +574,12 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { handler = Some(meta.value()?.parse()?); } else if meta.path.is_ident("msg") { msg = Some(meta.value()?.parse::()?.value()); + } else if meta.path.is_ident("meta") { + wants_meta = true; } else { return Err(meta.error( "unrecognized #[input] argument; expected `decode = ...`, \ - `handler = ...` or `msg = ...`", + `handler = ...`, `msg = ...` or `meta`", )); } Ok(()) @@ -576,7 +587,11 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { let decode = decode .ok_or_else(|| syn::Error::new_spanned(attr, "#[input] requires `decode = ...`"))?; let handler = handler.unwrap_or_else(|| format_ident!("handle_{}", name)); - found = Some(FieldKind::Input { decode, handler }); + found = Some(FieldKind::Input { + decode, + handler, + wants_meta, + }); } else if path.is_ident("output") { if found.is_some() { return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); diff --git a/native/rust/dimos-module/src/host.rs b/native/rust/dimos-module/src/host.rs index cc7982b7e5..07d219950c 100644 --- a/native/rust/dimos-module/src/host.rs +++ b/native/rust/dimos-module/src/host.rs @@ -19,7 +19,6 @@ //! [`host_main`]. Everything a host needs beyond that lives here, so the //! generated crate stays a table of module entries and three `include_str!`s. -use std::collections::HashMap; use std::future::Future; use std::io; use std::pin::Pin; @@ -33,7 +32,7 @@ use tracing::{error, info, warn}; use crate::lcm::LcmTransport; use crate::module::{ init_tracing, log_wiring, parse_config_value, read_launch_config, run_module_core, - validate_config, Module, + validate_config, Module, Topics, }; use crate::transport::{SharedTransport, Transport}; use crate::zenoh::{ZenohTransport, SESSION_KEY}; @@ -52,7 +51,7 @@ type RunFn = Box, watch::Receiver) -> Modu /// once the transport is open. Produced before anything is spawned so a bad /// config kills the host instead of half-starting it. struct Prepared { - topics: HashMap, + topics: Topics, run: RunFn, } @@ -450,7 +449,7 @@ fn run_host_fallible(spec: &HostSpec) -> io::Result<()> { let prepared = prepare_all(spec, &stdin)?; let known: Vec = prepared .iter() - .flat_map(|p| p.topics.values().cloned()) + .flat_map(|p| p.topics.channels().cloned()) .collect(); let suppress = resolve_suppress(spec, &stdin, &known)?; if let Ok(transport) = std::env::var("DIMOS_TRANSPORT") { @@ -587,6 +586,7 @@ fn supervise( #[cfg(test)] mod tests { use super::*; + use std::collections::HashMap; #[test] fn thread_name_is_prefixed_and_capped() { diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index a1a6f6efd6..7307b1145e 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -29,7 +29,9 @@ pub mod zenoh; pub use dimos_module_macros::{native_config, Module}; pub use host::{host_main, HostSpec, ModuleEntry}; pub use lcm::LcmTransport; -pub use module::{run, Builder, Input, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output}; +pub use module::{ + run, Builder, Input, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output, TopicMetadata, +}; pub use tf::{Lookup, Tf, Transform}; pub use transport::{SharedTransport, Transport}; pub use workers::worker_pool; diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index e9a5f014fa..6e5070d391 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -74,9 +74,13 @@ pub(crate) trait Route: Send + Sync { fn try_dispatch(&self, data: &[u8]); } +/// Decodes a frame into whatever the port's channel carries. Boxed because an +/// input route tags the message with the index of the topic it arrived on. +type Decode = Box io::Result + Send + Sync>; + struct TypedRoute { topic: String, - decode: fn(&[u8]) -> io::Result, + decode: Decode, sender: mpsc::Sender, drop_count: AtomicU64, last_log_ns: AtomicU64, @@ -109,14 +113,55 @@ impl Route for TypedRoute { } } } -pub struct Input { +/// Which topic of an input a message arrived on, for handlers that opt in +/// with `#[input(meta)]`. `info` is whatever the coordinator attached to that +/// topic on the launch line (`Null` when nothing was). +#[derive(Clone, Debug)] +pub struct TopicMetadata { + pub index: usize, pub topic: String, - receiver: mpsc::Receiver, + pub info: serde_json::Value, +} + +/// A port wired to one topic, or to every topic the launch line lists under +/// it — a topic funnel. Every message carries the index of the topic it +/// arrived on; `recv` drops it, `recv_meta` hands it over. A port given an +/// empty array never yields. +pub struct Input { + pub topics: Vec, + pub infos: Vec, + receiver: mpsc::Receiver<(usize, T)>, } impl Input { pub async fn recv(&mut self) -> Option { - self.receiver.recv().await + self.receiver.recv().await.map(|(_, msg)| msg) + } + + pub async fn recv_meta(&mut self) -> Option<(T, TopicMetadata)> { + let (index, msg) = self.receiver.recv().await?; + let meta = TopicMetadata { + index, + topic: self.topics[index].clone(), + info: self.infos[index].clone(), + }; + Some((msg, meta)) + } + + pub fn topic(&self, index: usize) -> &str { + &self.topics[index] + } + + pub fn info(&self, index: usize) -> &serde_json::Value { + &self.infos[index] + } + + pub fn len(&self) -> usize { + self.topics.len() + } + + pub fn is_empty(&self) -> bool { + self.topics.is_empty() } } @@ -164,19 +209,100 @@ pub(crate) async fn publish_encoded( .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "background task gone")) } +/// One source of a funneled input: its topic plus whatever per-topic info the +/// coordinator attached on the launch line. +#[derive(Clone, Debug)] +pub(crate) struct TopicEntry { + topic: String, + info: serde_json::Value, +} + +impl TopicEntry { + fn bare(topic: impl Into) -> Self { + Self { + topic: topic.into(), + info: serde_json::Value::Null, + } + } +} + +/// The port-to-topic wiring from the launch line. A port names one topic, or an +/// array of same-typed topics that one handler sees. +#[derive(Clone, Debug, Default)] +pub(crate) struct Topics { + single: HashMap, + grouped: HashMap>, +} + +impl Topics { + fn ports(&self) -> BTreeSet<&String> { + self.single.keys().chain(self.grouped.keys()).collect() + } + + /// Every wire channel the module touches, groups flattened. + pub(crate) fn channels(&self) -> impl Iterator { + self.single + .values() + .chain(self.grouped.values().flatten().map(|entry| &entry.topic)) + } +} + +fn parse_topics(json: &serde_json::Value) -> io::Result { + let mut topics = Topics::default(); + let Some(table) = json.get("topics").and_then(|v| v.as_object()) else { + return Ok(topics); + }; + for (port, value) in table { + let invalid = |detail: &str| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("topic for port '{port}' {detail}"), + ) + }; + match value { + serde_json::Value::String(topic) => { + topics.single.insert(port.clone(), topic.clone()); + } + serde_json::Value::Array(items) => { + let group = items + .iter() + .map(|item| match item { + serde_json::Value::String(topic) => Ok(TopicEntry::bare(topic)), + serde_json::Value::Object(fields) => { + let topic = + fields + .get("topic") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + invalid("has an entry object without a string 'topic'") + })?; + Ok(TopicEntry { + topic: topic.to_string(), + info: fields + .get("info") + .cloned() + .unwrap_or(serde_json::Value::Null), + }) + } + _ => Err(invalid( + "is a group, so every entry must be a string or a {topic, info} object", + )), + }) + .collect::>>()?; + topics.grouped.insert(port.clone(), group); + } + _ => return Err(invalid("must be a string or an array of strings")), + } + } + Ok(topics) +} + /// Extract `(topics, config)` from an already-parsed config object. `run` /// parses the line once and also reads `qos` from it, so this takes the value. pub(crate) fn parse_config_value( json: &serde_json::Value, -) -> io::Result<(HashMap, C)> { - let mut topics = HashMap::new(); - if let Some(t) = json.get("topics").and_then(|v| v.as_object()) { - for (port, topic) in t { - if let Some(s) = topic.as_str() { - topics.insert(port.clone(), s.to_string()); - } - } - } +) -> io::Result<(Topics, C)> { + let topics = parse_topics(json)?; let config_value = json.get("config").ok_or_else(|| { io::Error::new( @@ -305,7 +431,7 @@ pub trait Module: Sized + Send + 'static { } pub struct Builder { - topics: HashMap, + topics: Topics, // Every port the module asked for a topic, matched against topics after build. requested: BTreeSet, routes: HashMap>>, @@ -315,7 +441,7 @@ pub struct Builder { } impl Builder { - pub(crate) fn new(topics: HashMap) -> Self { + pub(crate) fn new(topics: Topics) -> Self { Self { topics, requested: BTreeSet::new(), @@ -328,15 +454,29 @@ impl Builder { fn topic_for(&mut self, port: &str) -> String { self.requested.insert(port.to_string()); self.topics + .single .get(port) .cloned() .unwrap_or_else(|| format!("/{port}")) } + /// One topic, or every topic an array under `port` lists. An unnamed port + /// falls back to `/{port}`; an empty array is a port with no sources. + fn input_entries_for(&mut self, port: &str) -> Vec { + self.requested.insert(port.to_string()); + if let Some(topic) = self.topics.single.get(port) { + vec![TopicEntry::bare(topic)] + } else if let Some(group) = self.topics.grouped.get(port) { + group.clone() + } else { + vec![TopicEntry::bare(format!("/{port}"))] + } + } + // A mismatch is dead wiring: an unclaimed topic reaches no port, and an // unsent one leaves the port on a fallback name nothing else publishes to. pub(crate) fn enforce_topics_match_ports(&self) -> io::Result<()> { - let provided: BTreeSet<&String> = self.topics.keys().collect(); + let provided = self.topics.ports(); let requested: BTreeSet<&String> = self.requested.iter().collect(); if provided == requested { return Ok(()); @@ -351,22 +491,31 @@ impl Builder { )) } - fn add_route( + fn push_route( &mut self, topic: &str, - decode: fn(&[u8]) -> io::Result, - ) -> mpsc::Receiver { - let (tx, rx) = mpsc::channel(INPUT_CHANNEL_CAPACITY); + decode: Decode, + sender: mpsc::Sender, + ) { self.routes .entry(topic.to_string()) .or_default() .push(Box::new(TypedRoute { topic: topic.to_string(), decode, - sender: tx, + sender, drop_count: AtomicU64::new(0), last_log_ns: AtomicU64::new(0), })); + } + + fn add_route( + &mut self, + topic: &str, + decode: fn(&[u8]) -> io::Result, + ) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(INPUT_CHANNEL_CAPACITY); + self.push_route(topic, Box::new(decode), tx); rx } @@ -381,9 +530,21 @@ impl Builder { port: &str, decode: fn(&[u8]) -> io::Result, ) -> Input { - let topic = self.topic_for(port); - let receiver = self.add_route(&topic, decode); - Input { topic, receiver } + let entries = self.input_entries_for(port); + let (sender, receiver) = mpsc::channel(INPUT_CHANNEL_CAPACITY); + for (index, entry) in entries.iter().enumerate() { + let tag = Box::new(move |bytes: &[u8]| decode(bytes).map(|msg| (index, msg))); + self.push_route(&entry.topic, tag, sender.clone()); + } + let (topics, infos) = entries + .into_iter() + .map(|entry| (entry.topic, entry.info)) + .unzip(); + Input { + topics, + infos, + receiver, + } } pub fn output(&mut self, port: &str, encode: fn(&T) -> Vec) -> Output { @@ -514,7 +675,7 @@ where /// and by the baked host, which drives several of these on one transport. pub(crate) async fn run_module_core( transport: Arc, - topics: HashMap, + topics: Topics, config: M::Config, mut shutdown: watch::Receiver, ) -> io::Result<()> @@ -551,10 +712,13 @@ where /// Log the resolved wiring of a module, tagged with whatever the operator sees /// in `ps`: the executable for a lone module, the module id inside a host. -pub(crate) fn log_wiring(exe: &str, topics: &HashMap, config: &C) { - for (port, topic) in topics { +pub(crate) fn log_wiring(exe: &str, topics: &Topics, config: &C) { + for (port, topic) in &topics.single { info!(exe = %exe, port = %port, topic = %topic, "topic mapping"); } + for (port, group) in &topics.grouped { + info!(exe = %exe, port = %port, topics = ?group, "topic group mapping"); + } info!(exe = %exe, config = ?config, "config loaded"); } @@ -599,9 +763,7 @@ mod tests { /// Parse a raw config line the way `run` does, for exercising /// `parse_config_value` from the string form the coordinator sends. - fn parse_config_json( - line: &str, - ) -> io::Result<(HashMap, C)> { + fn parse_config_json(line: &str) -> io::Result<(Topics, C)> { parse_config_value(&parse_launch_config(line)?) } @@ -726,8 +888,8 @@ mod tests { fn parses_topics_and_config() { let json = r#"{"topics": {"data": "/foo/data", "confirm": "/foo/confirm"}, "config": {"value": 42, "name": "hello"}}"#; let (topics, config) = parse_config_json::(json).unwrap(); - assert_eq!(topics["data"], "/foo/data"); - assert_eq!(topics["confirm"], "/foo/confirm"); + assert_eq!(topics.single["data"], "/foo/data"); + assert_eq!(topics.single["confirm"], "/foo/confirm"); assert_eq!( config, TestConfig { @@ -798,7 +960,7 @@ mod tests { fn missing_topics_field_gives_empty_map() { let json = r#"{"config": {"value": 1, "name": "x"}}"#; let (topics, _config) = parse_config_json::(json).unwrap(); - assert!(topics.is_empty()); + assert!(topics.ports().is_empty()); } #[test] @@ -887,11 +1049,24 @@ mod tests { // topic_for fallback - fn topics(pairs: &[(&str, &str)]) -> HashMap { - pairs - .iter() - .map(|(p, t)| (p.to_string(), t.to_string())) - .collect() + fn topics(pairs: &[(&str, &str)]) -> Topics { + Topics { + single: pairs + .iter() + .map(|(port, topic)| (port.to_string(), topic.to_string())) + .collect(), + grouped: HashMap::new(), + } + } + + fn grouped_topics(port: &str, group: &[&str]) -> Topics { + Topics { + single: HashMap::new(), + grouped: HashMap::from([( + port.to_string(), + group.iter().map(|t| TopicEntry::bare(*t)).collect(), + )]), + } } fn builder_with_topics(pairs: &[(&str, &str)]) -> Builder { @@ -914,14 +1089,14 @@ mod tests { fn input_uses_mapped_topic() { let mut builder = builder_with_topics(&[("data", "/test/data")]); let input = builder.input("data", |b| Ok(b.to_vec())); - assert_eq!(input.topic, "/test/data"); + assert_eq!(input.topics, ["/test/data"]); } #[test] fn input_falls_back_to_slash_port_when_unmapped() { let mut builder = builder_with_topics(&[]); let input = builder.input("data", |b| Ok(b.to_vec())); - assert_eq!(input.topic, "/data"); + assert_eq!(input.topics, ["/data"]); } #[test] @@ -931,6 +1106,83 @@ mod tests { assert_eq!(output.topic, "/robot/cmd_vel"); } + // topic funnels: a port wired to an array of topics + + #[test] + fn a_port_given_an_array_of_topics_becomes_a_funnel() { + let json = r#"{"topics": {"cams": ["/cam0", "/cam1"], "odom": "/odom"}, "config": null}"#; + let (topics, _config) = parse_config_json::<()>(json).unwrap(); + let cams: Vec<&str> = topics.grouped["cams"].iter().map(|e| &*e.topic).collect(); + assert_eq!(cams, ["/cam0", "/cam1"]); + assert_eq!(topics.single["odom"], "/odom"); + } + + #[test] + fn a_funnel_entry_that_is_not_a_string_is_rejected() { + let json = r#"{"topics": {"cams": ["/cam0", 7]}, "config": null}"#; + let err = parse_config_json::<()>(json).expect_err("a non-string funnel entry is invalid"); + assert!(err.to_string().contains("cams"), "{err}"); + } + + #[test] + fn a_funnel_entry_object_carries_per_topic_info() { + let json = r#"{"topics": {"cams": + ["/cam0", {"topic": "/cam1", "info": {"rectified": true}}]}, "config": null}"#; + let (topics, _config) = parse_config_json::<()>(json).unwrap(); + let cams = &topics.grouped["cams"]; + assert_eq!(cams[0].info, serde_json::Value::Null); + assert_eq!(cams[1].topic, "/cam1"); + assert_eq!(cams[1].info["rectified"], true); + } + + #[test] + fn a_funnel_entry_object_without_a_topic_is_rejected() { + let json = r#"{"topics": {"cams": [{"info": {"rectified": true}}]}, "config": null}"#; + let err = parse_config_json::<()>(json).expect_err("an entry object needs a topic"); + assert!(err.to_string().contains("'topic'"), "{err}"); + } + + #[test] + fn a_funneled_input_subscribes_every_topic_it_was_given() { + let mut builder = Builder::new(grouped_topics("cams", &["/cam0", "/cam1"])); + let input = builder.input("cams", |b| Ok(b.to_vec())); + assert_eq!(input.topics, ["/cam0", "/cam1"]); + assert_eq!(builder.routes.get("/cam0").map(Vec::len), Some(1)); + assert_eq!(builder.routes.get("/cam1").map(Vec::len), Some(1)); + builder.enforce_topics_match_ports().expect("port claimed"); + } + + /// One handler sees every topic, so the metadata is the only thing that + /// says which camera of a rig a frame came from. + #[tokio::test] + async fn recv_meta_names_the_topic_a_message_arrived_on() { + let mut topics = grouped_topics("cams", &["/cam0", "/cam1"]); + topics.grouped.get_mut("cams").unwrap()[1].info = serde_json::json!({"rectified": true}); + let mut builder = Builder::new(topics); + let mut input = builder.input("cams", |b| Ok(b.to_vec())); + + builder.routes["/cam1"][0].try_dispatch(b"second"); + builder.routes["/cam0"][0].try_dispatch(b"first"); + + let (msg, meta) = input.recv_meta().await.expect("cam1 frame"); + assert_eq!(msg, b"second"); + assert_eq!(meta.index, 1); + assert_eq!(meta.topic, "/cam1"); + assert_eq!(meta.info["rectified"], true); + + // recv drops the tag for handlers that treat the funnel as one stream. + assert_eq!(input.recv().await.expect("cam0 frame"), b"first"); + } + + #[test] + fn an_empty_funnel_still_claims_its_port() { + let mut builder = Builder::new(grouped_topics("cams", &[])); + let input = builder.input("cams", |b| Ok(b.to_vec())); + assert!(input.is_empty()); + assert!(builder.routes.is_empty()); + builder.enforce_topics_match_ports().expect("port claimed"); + } + #[test] fn topics_matching_ports_exactly_pass() { let mut builder = builder_with_topics(&[("cmd", "/robot/cmd"), ("odom", "/robot/odom")]); @@ -1241,7 +1493,7 @@ mod tests { let (tx, _rx) = mpsc::channel::>(1); let route = TypedRoute { topic: "/test".to_string(), - decode: |b| Ok(b.to_vec()), + decode: Box::new(|b: &[u8]| Ok(b.to_vec())), sender: tx, drop_count: AtomicU64::new(0), last_log_ns: AtomicU64::new(0), @@ -1284,6 +1536,41 @@ mod tests { } } + #[derive(crate::Module)] + struct Rig { + #[input(decode = decode, meta)] + cams: crate::Input, + #[output(encode = encode)] + seen: crate::Output, + } + + impl Rig { + async fn handle_cams(&mut self, msg: Msg, meta: crate::TopicMetadata) { + let mut tagged = vec![meta.index as u8]; + tagged.extend(msg.0); + self.seen.publish(&Msg(tagged)).await.expect("publish"); + } + } + + #[tokio::test] + async fn a_meta_input_hands_its_handler_the_topic_index() { + let mut builder = Builder::new(Topics { + single: HashMap::from([("seen".to_string(), "/seen".to_string())]), + grouped: HashMap::from([( + "cams".to_string(), + vec![TopicEntry::bare("/cam0"), TopicEntry::bare("/cam1")], + )]), + }); + let mut rig = Rig::build(&mut builder, NoConfig); + + builder.routes["/cam1"][0].try_dispatch(b"frame"); + builder.routes.clear(); + rig.handle().await; + + let (_, rx) = &mut builder.outputs[0]; + assert_eq!(rx.recv().await.expect("handler output"), b"\x01frame"); + } + #[tokio::test] async fn io_field_is_wired_to_its_handler_and_can_publish() { let mut builder = Builder::new(topics(&[("cmd", "/robot/cmd")]));