Merge pull request #77 from sourcebots/mqtt-publish-state

Publish robot USB status to MQTT

Author: WillB97

runusb/__main__.py

@@ -48,6 +48,8 @@
 # This will be populated if we have the config file
 # url format: mqtt[s]://[<username>[:<password>]@]<host>[:<port>]/<topic_root>
 MQTT_URL = None
+MQTT_TOPIC_ROOT = ''
+MQTT_CLIENT = None
 MQTT_CONFIG_FILE = '/etc/sbot/mqtt.conf'
 
 
@@ -156,6 +158,15 @@ def set_status(self, value: LedStatus) -> None:
             GPIO.output(self.LEDs.STATUS_GREEN, GPIO.HIGH if value.value[1] else GPIO.LOW)
             GPIO.output(self.LEDs.STATUS_BLUE, GPIO.HIGH if value.value[2] else GPIO.LOW)
 
+        # Also send the status over MQTT
+        if MQTT_CLIENT is not None:
+            MQTT_CLIENT.publish(
+                f'{MQTT_TOPIC_ROOT}/state',
+                json.dumps({"state": value.name}),
+                qos=1,
+                retain=True,
+            )
+
 
 LED_CONTROLLER = LEDController()
 
@@ -238,6 +249,7 @@ def __init__(self, mountpoint_path: str) -> None:
         self._setup_logging(mountpoint_path)
         LED_CONTROLLER.set_code(True)
         LED_CONTROLLER.set_status(LedStatus.Running)
+
         env = dict(os.environ)
         env["SBOT_METADATA_PATH"] = MOUNTPOINT_DIR
         if MQTT_URL is not None:
@@ -270,7 +282,6 @@ def cleanup(self) -> None:
         except subprocess.TimeoutExpired:
             # The process did not exit after 5 seconds, so kill it.
             self._send_signal(signal.SIGKILL)
-        self._set_leds()
 
     def close(self) -> None:
         self.cleanup()
@@ -289,8 +300,10 @@ def _watch_process(self) -> None:
         self.process.wait()
         if self.process.returncode != 0:
             USERCODE_LOGGER.warning(f"Process exited with code {self.process.returncode}")
+            LED_CONTROLLER.set_status(LedStatus.Crashed)
         else:
             USERCODE_LOGGER.info("Your code finished successfully.")
+            LED_CONTROLLER.set_status(LedStatus.Finished)
 
         process_lifetime = time.time() - self.process_start_time
 
@@ -328,12 +341,6 @@ def _log_output(self, pipe: IO[str]) -> None:
             USERCODE_LOGGER.log(USERCODE_LEVEL, line.rstrip('\n'))
         LOGGER.info('Process output finished')
 
-    def _set_leds(self) -> None:
-        if self.process.returncode == 0:
-            LED_CONTROLLER.set_status(LedStatus.Finished)
-        else:
-            LED_CONTROLLER.set_status(LedStatus.Crashed)
-
     def _rotate_old_logs(self, log_dir: str) -> None:
         """
         Add a suffix to the old log file, if it exists.
@@ -460,7 +467,7 @@ def read_mqtt_config_file() -> MQTTVariables | None:
 
 
 def setup_usercode_logging() -> None:
-    global REL_TIME_FILTER
+    global REL_TIME_FILTER, MQTT_CLIENT, MQTT_TOPIC_ROOT
     REL_TIME_FILTER = RelativeTimeFilter()
     USERCODE_LOGGER.addFilter(REL_TIME_FILTER)
     USERCODE_LOGGER.setLevel(logging.DEBUG)
@@ -481,6 +488,8 @@ def setup_usercode_logging() -> None:
                 connected_callback=lambda: LED_CONTROLLER.set_wifi(True),
                 disconnected_callback=lambda: LED_CONTROLLER.set_wifi(False),
             )
+            MQTT_CLIENT = handler.mqtt
+            MQTT_TOPIC_ROOT = mqtt_config.topic_prefix
 
             handler.setLevel(logging.INFO)
             handler.setFormatter(TieredFormatter(
@@ -501,6 +510,8 @@ def main():
     registry = AutorunProcessRegistry()
 
     LED_CONTROLLER.mark_start()
+    LED_CONTROLLER.set_status(LedStatus.NoUSB)
+
     # Initial pass (in case an autorun FS is already mounted)
     registry.update_filesystems(fstab_reader.read())