Merge pull request #76 from sourcebots/kch-support

Add support for KCH leds

Author: WillB97

.github/workflows/ci.yml

@@ -10,10 +10,10 @@ jobs:
       MYPY: venv/bin/mypy
     steps:
     - uses: actions/checkout@v3
-    - name: Set up Python 3.7
+    - name: Set up Python 3.8
       uses: actions/setup-python@v3
       with:
-        python-version: 3.7
+        python-version: 3.8
     - name: Set up virtualenv
       run: python -m venv venv
     - name: Install dependencies

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'
 
 
@@ -86,37 +88,84 @@ class Mountpoint(NamedTuple):
 )
 
 
+class LedStatus(Enum):
+    NoUSB = (False, False, False)  # Off
+    Running = (False, False, True)  # Blue
+    Killed = (True, False, True)  # Magenta
+    Finished = (False, True, False)  # Green
+    Crashed = (True, False, False)  # Red
+
+
 class LEDController():
     @unique
     class LEDs(IntEnum):
-        RED = 2
-        YELLOW = 3
-        GREEN = 4
+        BOOT_100 = 13
+        CODE = 11
+        COMP = 16
+        WIFI = 8
+        STATUS_RED = 26
+        STATUS_GREEN = 20
+        STATUS_BLUE = 21
 
     def __init__(self) -> None:
         if IS_PI:
             LOGGER.debug("Configuring LED controller")
+            self._register_exit()
             atexit.register(GPIO.cleanup)  # type: ignore[attr-defined]
             GPIO.setmode(GPIO.BCM)
             GPIO.setup([led.value for led in self.LEDs], GPIO.OUT, initial=GPIO.LOW)
 
-    def red(self) -> None:
+    def _register_exit(self) -> None:
+        """
+        Ensure `atexit` triggers on `SIGTERM`.
+
+        > The functions registered via [`atexit`] are not called when the program is
+        killed by a signal not handled by Python
+        """
+
+        if signal.getsignal(signal.SIGTERM) != signal.SIG_DFL:
+            # If a signal handler is already present for SIGTERM,
+            # this is sufficient for `atexit` to trigger, so do nothing.
+            return
+
+        def handle_signal(handled_signum: int, frame) -> None:
+            """Semi-default signal handler for SIGTERM, enough for atexit."""
+            USERCODE_LOGGER.error(signal.strsignal(handled_signum))
+            exit(128 + handled_signum)  # 143 for SIGTERM
+
+        # Add the null-ish signal handler
+        signal.signal(signal.SIGTERM, handle_signal)
+
+    def mark_start(self) -> None:
+        if IS_PI:
+            GPIO.output(self.LEDs.BOOT_100, GPIO.HIGH)
+
+    def set_comp(self, value: bool) -> None:
+        if IS_PI:
+            GPIO.output(self.LEDs.COMP, GPIO.HIGH if value else GPIO.LOW)
+
+    def set_code(self, value: bool) -> None:
         if IS_PI:
-            GPIO.output(self.LEDs.RED, GPIO.HIGH)
-            GPIO.output(self.LEDs.YELLOW, GPIO.LOW)
-            GPIO.output(self.LEDs.GREEN, GPIO.LOW)
+            GPIO.output(self.LEDs.CODE, GPIO.HIGH if value else GPIO.LOW)
 
-    def yellow(self) -> None:
+    def set_wifi(self, value: bool) -> None:
         if IS_PI:
-            GPIO.output(self.LEDs.RED, GPIO.LOW)
-            GPIO.output(self.LEDs.YELLOW, GPIO.HIGH)
-            GPIO.output(self.LEDs.GREEN, GPIO.LOW)
+            GPIO.output(self.LEDs.WIFI, GPIO.HIGH if value else GPIO.LOW)
 
-    def green(self) -> None:
+    def set_status(self, value: LedStatus) -> None:
         if IS_PI:
-            GPIO.output(self.LEDs.RED, GPIO.LOW)
-            GPIO.output(self.LEDs.YELLOW, GPIO.LOW)
-            GPIO.output(self.LEDs.GREEN, GPIO.HIGH)
+            GPIO.output(self.LEDs.STATUS_RED, GPIO.HIGH if value.value[0] else GPIO.LOW)
+            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()
@@ -198,7 +247,9 @@ def close(self) -> None:
 class RobotUSBHandler(USBHandler):
     def __init__(self, mountpoint_path: str) -> None:
         self._setup_logging(mountpoint_path)
-        LED_CONTROLLER.yellow()
+        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:
@@ -223,15 +274,19 @@ def __init__(self, mountpoint_path: str) -> None:
             target=self._log_output, args=(self.process.stdout,))
         self.log_thread.start()
 
-    def close(self) -> None:
+    def cleanup(self) -> None:
         self._send_signal(signal.SIGTERM)
         try:
             # Wait for the process to exit
             self.process.communicate(timeout=5)
         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()
+        LED_CONTROLLER.set_status(LedStatus.NoUSB)
+        LED_CONTROLLER.set_code(False)
         USERCODE_LOGGER.removeHandler(self.handler)
 
     def _send_signal(self, sig: int) -> None:
@@ -245,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
 
@@ -256,7 +313,7 @@ def _watch_process(self) -> None:
             time.sleep(1 - process_lifetime)
 
         # Start clean-up
-        self.close()
+        self.cleanup()
 
     def _setup_logging(self, log_dir: str) -> None:
         self._rotate_old_logs(log_dir)
@@ -284,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.green()
-        else:
-            LED_CONTROLLER.red()
-
     def _rotate_old_logs(self, log_dir: str) -> None:
         """
         Add a suffix to the old log file, if it exists.
@@ -309,10 +360,12 @@ def _rotate_old_logs(self, log_dir: str) -> None:
 
 class MetadataUSBHandler(USBHandler):
     def __init__(self, mountpoint_path: str) -> None:
-        pass  # Nothing to do.
+        # NOTE the comp LED just represents the presence of a comp mode USB
+        # not whether comp mode is enabled
+        LED_CONTROLLER.set_comp(True)
 
     def close(self) -> None:
-        pass  # Nothing to do.
+        LED_CONTROLLER.set_comp(False)
 
 
 class AutorunProcessRegistry(object):
@@ -414,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)
@@ -432,7 +485,11 @@ def setup_usercode_logging() -> None:
                 username=mqtt_config.username,
                 password=mqtt_config.password,
                 connected_topic=f"{mqtt_config.topic_prefix}/connected",
+                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(
@@ -452,6 +509,9 @@ 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())
 

setup.cfg

@@ -11,11 +11,11 @@ url = https://github.com/sourcebots/runusb
 python_requires = >=3.8
 packages = find:
 install_requires =
-    logger-extras==0.4.0
+    logger-extras==0.4.1
     rpi.GPIO==0.7.1
 
 [options.extras_require]
-mqtt = logger-extras[mqtt]==0.4.0
+mqtt = logger-extras[mqtt]==0.4.1
 
 [options.entry_points]
 console_scripts =