add second fix

Author: paul-nechifor

dimos/protocol/rpc/pubsubrpc.py

@@ -169,16 +169,18 @@ def stop(self) -> None:
         """
         self._shutdown_thread_pool()
 
-        # Cleanup shared response subscriptions
+        # Stop the LCM loop BEFORE unsubscribing. Doing it the other way
+        # triggers a NULL-deref in LCM's C code.  This order is safer because
+        # unsub() is either a no-op (LCM destroyed) or safe (no concurrent
+        # dispatch).
+        if hasattr(super(), "stop"):
+            super().stop()  # type: ignore[misc]
+
         with self._response_subs_lock:
             for unsub, _ in self._response_subs.values():
                 unsub()
             self._response_subs.clear()
 
-        # Call parent stop if it exists
-        if hasattr(super(), "stop"):
-            super().stop()  # type: ignore[misc]
-
     def call(self, name: str, arguments: Args, cb: Callable | None):  # type: ignore[no-untyped-def, type-arg]
         if cb is None:
             return self.call_nowait(name, arguments)
@@ -209,9 +211,7 @@ def shared_response_handler(msg: MsgT, _: TopicT) -> None:
                     if res_id is None:
                         return
 
-                    # Look up callback for this msg_id
-                    with self._response_subs_lock:
-                        callback = callbacks_dict.pop(res_id, None)
+                    callback = callbacks_dict.pop(res_id, None)
 
                     if callback is None:
                         return  # No callback registered (already handled or timed out)

dimos/protocol/service/lcmservice.py

@@ -63,7 +63,7 @@ class LCMService(Service):
     config: LCMConfig
     l: lcm_mod.LCM | None
     _stop_event: threading.Event
-    _l_lock: threading.Lock
+    _l_lock: threading.RLock
     _thread: threading.Thread | None
     _call_thread_pool: ThreadPoolExecutor | None = None
     _call_thread_pool_lock: threading.RLock = threading.RLock()
@@ -77,7 +77,7 @@ def __init__(self, **kwargs: Any) -> None:
         else:
             self.l = lcm_mod.LCM(self.config.url) if self.config.url else lcm_mod.LCM()
 
-        self._l_lock = threading.Lock()
+        self._l_lock = threading.RLock()
         self._stop_event = threading.Event()
         self._thread = None
 
@@ -100,7 +100,7 @@ def __setstate__(self, state) -> None:  # type: ignore[no-untyped-def]
         self.l = None
         self._stop_event = threading.Event()
         self._thread = None
-        self._l_lock = threading.Lock()
+        self._l_lock = threading.RLock()
         self._call_thread_pool = None
         self._call_thread_pool_lock = threading.RLock()
 
@@ -120,31 +120,44 @@ def start(self) -> None:
                 return
 
             # Reinitialize LCM if it's None (e.g., after unpickling)
-            if self.l is None:
+            l = self.l
+            if l is None:
                 if self.config.lcm:
-                    self.l = self.config.lcm
+                    l = self.config.lcm
                 else:
-                    self.l = lcm_mod.LCM(self.config.url) if self.config.url else lcm_mod.LCM()
+                    l = lcm_mod.LCM(self.config.url) if self.config.url else lcm_mod.LCM()
+                self.l = l
+
+            # Pre-warm LCM recv setup (including the up-to-10s self-test) now,
+            # so _lcm_loop doesn't do it later while holding _l_lock, which
+            # would block external publish/subscribe/unsubscribe calls.
+            try:
+                l.fileno()
+            except Exception as e:
+                logger.warning(f"LCM fileno pre-warm failed: {e}")
 
             self._stop_event.clear()
             self._thread = threading.Thread(target=self._lcm_loop, daemon=True)
             self._thread.start()
 
     def _lcm_loop(self) -> None:
-        """LCM message handling loop."""
+        """LCM message handling loop.
+
+        Holds _l_lock during handle_timeout so that external callers
+        (publish/subscribe/unsubscribe on this LCMService from other threads)
+        are serialized with dispatch. This prevents a data race in the LCM
+        Python binding that segfaults otherwise.
+        """
         try:
             while not self._stop_event.is_set():
                 with self._l_lock:
-                    l = self.l
-                    if l is None:
+                    if self.l is None:
                         break
-                try:
-                    # This doesn't have to be under a lock because the C
-                    # library has its own locking for this.
-                    l.handle_timeout(_LCM_LOOP_TIMEOUT)
-                except Exception as e:
-                    stack_trace = traceback.format_exc()
-                    print(f"Error in LCM handling: {e}\n{stack_trace}")
+                    try:
+                        self.l.handle_timeout(_LCM_LOOP_TIMEOUT)
+                    except Exception as e:
+                        stack_trace = traceback.format_exc()
+                        print(f"Error in LCM handling: {e}\n{stack_trace}")
         finally:
             self._cleanup_owned_lcm()
             with self._l_lock:

dimos/protocol/service/test_lcmservice.py

@@ -52,6 +52,12 @@ def __init__(self) -> None:
         self.unsubscribe_called = threading.Event()
         self.subscription = BlockingSubscription()
 
+    def fileno(self) -> int:
+        # Called by LCMService.start() to pre-warm LCM recv setup. A real LCM
+        # sets up sockets and runs a self-test here; for the fake we just
+        # return a sentinel fd since nothing will select() on it.
+        return -1
+
     def handle_timeout(self, _timeout: int) -> None:
         self.handle_entered.set()
         self.release_handle.wait(timeout=1.0)
@@ -323,7 +329,7 @@ def fake_pub_sub(fake_lcm):
     pubsub.stop()
 
 
-def test_publish_proceeds_during_handle_loop(fake_lcm, fake_pub_sub):
+def test_publish_waits_for_handle_loop(fake_lcm, fake_pub_sub):
     assert fake_lcm.handle_entered.wait(timeout=0.5)
 
     publisher = threading.Thread(
@@ -332,14 +338,20 @@ def test_publish_proceeds_during_handle_loop(fake_lcm, fake_pub_sub):
     )
     publisher.start()
 
-    assert fake_lcm.publish_called.wait(timeout=0.5)
-    publisher.join(timeout=1.0)
-    assert not publisher.is_alive()
+    # Loop holds _l_lock inside handle_timeout; publish waits for the lock.
+    assert not fake_lcm.publish_called.wait(timeout=0.1)
+    assert publisher.is_alive()
 
+    # Releasing handle_timeout drops _l_lock; publish now proceeds.
     fake_lcm.release_handle.set()
 
+    assert fake_lcm.publish_called.wait(timeout=1.0)
+    publisher.join(timeout=1.0)
+    assert not publisher.is_alive()
 
-def test_subscribe_proceeds_during_handle_loop(fake_lcm, fake_pub_sub):
+
+def test_subscribe_waits_for_handle_loop(fake_lcm, fake_pub_sub):
+    """subscribe() must block while the loop thread is inside handle_timeout()."""
     assert fake_lcm.handle_entered.wait(timeout=0.5)
 
     subscriber = threading.Thread(
@@ -348,16 +360,27 @@ def test_subscribe_proceeds_during_handle_loop(fake_lcm, fake_pub_sub):
     )
     subscriber.start()
 
-    assert fake_lcm.subscribe_called.wait(timeout=0.5)
+    assert not fake_lcm.subscribe_called.wait(timeout=0.1)
+    assert subscriber.is_alive()
+
+    fake_lcm.release_handle.set()
+
+    assert fake_lcm.subscribe_called.wait(timeout=1.0)
     subscriber.join(timeout=1.0)
     assert not subscriber.is_alive()
     assert fake_lcm.subscription.queue_capacity == 10000
 
-    fake_lcm.release_handle.set()
 
+def test_unsubscribe_waits_for_handle_loop(fake_lcm, fake_pub_sub):
+    """unsubscribe() must block while the loop thread is inside handle_timeout().
 
-def test_unsubscribe_proceeds_during_handle_loop(fake_lcm, fake_pub_sub):
+    This is the specific race whose resolution fixes the segfault in
+    pylcm.c. Unsubscribing from another thread while dispatch is running
+    would set subs_obj->lcm_obj = NULL under the nose of pylcm_msg_handler.
+    """
+    # Let the first handle_timeout complete so we can subscribe cleanly.
     assert fake_lcm.handle_entered.wait(timeout=0.5)
+    fake_lcm.release_handle.set()
 
     unsubscribe_holder: dict[str, object] = {}
 
@@ -366,20 +389,29 @@ def do_subscribe() -> None:
 
     subscriber = threading.Thread(target=do_subscribe, daemon=True)
     subscriber.start()
-    assert fake_lcm.subscribe_called.wait(timeout=0.5)
+    assert fake_lcm.subscribe_called.wait(timeout=1.0)
     subscriber.join(timeout=1.0)
     assert not subscriber.is_alive()
 
+    # Reset gates so the next handle_timeout iteration blocks again.
+    fake_lcm.handle_entered.clear()
+    fake_lcm.release_handle.clear()
+    assert fake_lcm.handle_entered.wait(timeout=1.0)
+
     unsubscribe = unsubscribe_holder["fn"]
     unsub_thread = threading.Thread(target=unsubscribe, daemon=True)  # type: ignore[arg-type]
     unsub_thread.start()
 
-    assert fake_lcm.unsubscribe_called.wait(timeout=0.5)
-    unsub_thread.join(timeout=1.0)
-    assert not unsub_thread.is_alive()
+    # Loop holds _l_lock; unsubscribe waits for the lock.
+    assert not fake_lcm.unsubscribe_called.wait(timeout=0.1)
+    assert unsub_thread.is_alive()
 
     fake_lcm.release_handle.set()
 
+    assert fake_lcm.unsubscribe_called.wait(timeout=1.0)
+    unsub_thread.join(timeout=1.0)
+    assert not unsub_thread.is_alive()
+
 
 def test_stop_from_within_lcm_thread(mocker):
     """stop() called from inside handle_timeout must not deadlock and must
@@ -391,6 +423,9 @@ class SelfStoppingLCM:
         def __init__(self) -> None:
             self.done = threading.Event()
 
+        def fileno(self) -> int:
+            return -1
+
         def handle_timeout(self, _timeout: int) -> None:
             if not self.done.is_set():
                 captured["thread"] = threading.current_thread()
@@ -420,3 +455,57 @@ def unsubscribe(self, *_args: object) -> None:
     assert not thread.is_alive()
     assert service.l is None
     assert service._thread is None
+
+
+def test_handler_can_publish_via_rlock_reentry(mocker):
+    """A message handler dispatched from handle_timeout runs on the loop
+    thread while it already holds _l_lock. Reentry must work so the handler
+    can call self.publish/subscribe/unsubscribe. This is why _l_lock is an
+    RLock rather than a plain Lock.
+    """
+    publish_calls: list[tuple[str, bytes]] = []
+    handler_done = threading.Event()
+
+    class ReentrantLCM:
+        def __init__(self) -> None:
+            self._handler = None
+            self._dispatched = False
+            self._subscription = BlockingSubscription()
+
+        def fileno(self) -> int:
+            return -1
+
+        def handle_timeout(self, _timeout: int) -> None:
+            # Dispatch one fake message on the first call after a handler
+            # has been registered. The handler will call self.publish, which
+            # must reenter _l_lock recursively on the loop thread.
+            if self._handler is not None and not self._dispatched:
+                self._dispatched = True
+                self._handler("/req", b"req-payload")
+                handler_done.set()
+
+        def publish(self, channel: str, message: bytes) -> None:
+            publish_calls.append((channel, message))
+
+        def subscribe(self, _channel, handler) -> BlockingSubscription:
+            self._handler = handler
+            return self._subscription
+
+        def unsubscribe(self, _subscription) -> None:
+            pass
+
+    fake = ReentrantLCM()
+    mocker.patch("dimos.protocol.service.lcmservice.lcm_mod.LCM", return_value=fake)
+
+    pubsub = LCMPubSubBase()
+    pubsub.start()
+
+    def reentrant_callback(_msg: bytes, _topic: Topic) -> None:
+        pubsub.publish(Topic("/res"), b"res-payload")
+
+    pubsub.subscribe(Topic("/req"), reentrant_callback)
+
+    assert handler_done.wait(timeout=2.0)
+    pubsub.stop()
+
+    assert ("/res", b"res-payload") in publish_calls