From 81a3b704152b2f74a635d7cfaff346447f4282a0 Mon Sep 17 00:00:00 2001 From: Lloyd Date: Tue, 21 Apr 2026 12:07:08 +0100 Subject: [PATCH] feat: implement graceful shutdown handling and version cache optimizations --- repeater/data_acquisition/letsmesh_handler.py | 48 +++++++++-- repeater/main.py | 24 ++++-- repeater/web/update_endpoints.py | 81 ++++++++++++++++--- 3 files changed, 130 insertions(+), 23 deletions(-) diff --git a/repeater/data_acquisition/letsmesh_handler.py b/repeater/data_acquisition/letsmesh_handler.py index 44a45f4..88c0f6a 100644 --- a/repeater/data_acquisition/letsmesh_handler.py +++ b/repeater/data_acquisition/letsmesh_handler.py @@ -96,6 +96,7 @@ class _BrokerConnection: self._reconnect_timer = None self._max_reconnect_delay = 300 # 5 minutes max self._jwt_refresh_timer = None + self._shutdown_requested = False client_id = f"meshcore_{self.public_key}_{broker['host']}" self.client = mqtt.Client(client_id=client_id, transport="websockets") self.client.on_connect = self._on_connect @@ -163,6 +164,12 @@ class _BrokerConnection: was_running = self._running self._running = False + if self._shutdown_requested: + logger.info(f"Clean disconnect from {self.broker['name']}") + if self._on_disconnect_callback: + self._on_disconnect_callback(self.broker["name"]) + return + if rc != 0: # Unexpected disconnect error_msg = get_mqtt_error_message(rc, is_disconnect=True) logger.warning(f"Disconnected from {self.broker['name']} (rc={rc}): {error_msg}") @@ -176,6 +183,9 @@ class _BrokerConnection: def _schedule_reconnect(self, reason: str = "connection lost"): """Schedule reconnection with exponential backoff""" + if self._shutdown_requested: + return + if self._reconnect_timer: self._reconnect_timer.cancel() @@ -192,6 +202,9 @@ class _BrokerConnection: def _attempt_reconnect(self, reason: str = "connection lost"): """Attempt to reconnect to broker with fresh JWT""" + if self._shutdown_requested: + return + try: logger.info(f"Attempting reconnection to {self.broker['name']} (reason: {reason})...") @@ -227,6 +240,8 @@ class _BrokerConnection: def connect(self): """Establish connection to broker""" + self._shutdown_requested = False + # Conditional TLS setup if self.use_tls: import ssl @@ -252,6 +267,7 @@ class _BrokerConnection: def disconnect(self): """Disconnect from broker""" + self._shutdown_requested = True self._running = False self._loop_running = False @@ -407,7 +423,9 @@ class MeshCoreToMqttJwtPusher: self.stats_provider = stats_provider self._status_task = None self._running = False + self._shutdown_requested = False self._lock = threading.Lock() + self._connect_timers: List[threading.Timer] = [] # Create broker connections self.connections: List[_BrokerConnection] = [] @@ -431,6 +449,9 @@ class MeshCoreToMqttJwtPusher: def _on_broker_connected(self, broker_name: str): """Callback when a broker connects""" + if self._shutdown_requested: + return + # Publish initial status on first connection if not self._status_task and self.status_interval > 0: self._running = True @@ -455,6 +476,9 @@ class MeshCoreToMqttJwtPusher: def connect(self): """Establish connections to all configured brokers""" + self._shutdown_requested = False + self._connect_timers = [] + for idx, conn in enumerate(self.connections): try: if idx == 0: @@ -467,11 +491,15 @@ class MeshCoreToMqttJwtPusher: timer = threading.Timer(delay, lambda c=conn: self._delayed_connect(c)) timer.daemon = True timer.start() + self._connect_timers.append(timer) except Exception as e: logger.error(f"Failed to connect to {conn.broker['name']}: {e}") def _delayed_connect(self, conn): """Connect a broker after a delay (called by timer)""" + if self._shutdown_requested: + return + try: conn.connect() except Exception as e: @@ -479,15 +507,24 @@ class MeshCoreToMqttJwtPusher: def disconnect(self): """Disconnect from all brokers""" + self._shutdown_requested = True + + # Cancel any delayed connect timers first. + for timer in self._connect_timers: + try: + timer.cancel() + except Exception: + pass + self._connect_timers = [] + # Stop the heartbeat loop self._running = False # Publish offline status before disconnecting - self.publish_status(state="offline", origin=self.node_name, radio_config=self.radio_config) - - import time - - time.sleep(0.5) # Give time for messages to be sent + try: + self.publish_status(state="offline", origin=self.node_name, radio_config=self.radio_config) + except Exception: + pass # Disconnect all brokers for conn in self.connections: @@ -496,6 +533,7 @@ class MeshCoreToMqttJwtPusher: except Exception as e: logger.error(f"Error disconnecting from {conn.broker['name']}: {e}") + self._status_task = None logger.info("Disconnected from all brokers") def _status_heartbeat_loop(self): diff --git a/repeater/main.py b/repeater/main.py index 2309425..2910b28 100644 --- a/repeater/main.py +++ b/repeater/main.py @@ -54,6 +54,7 @@ class RepeaterDaemon: self.companion_bridges: dict[int, object] = {} self.companion_frame_servers: list = [] self._shutdown_started = False + self._main_task = None log_level = config.get("logging", {}).get("level", "INFO") logging.basicConfig( @@ -1026,6 +1027,9 @@ class RepeaterDaemon: return logger.info(f"Received signal {sig.name}, shutting down...") loop.create_task(self._shutdown()) + # Cancel run() so dispatcher.run_forever() unwinds cleanly. + if self._main_task and not self._main_task.done(): + self._main_task.cancel() async def _shutdown(self): """Best-effort shutdown: stop background services and release hardware.""" @@ -1059,7 +1063,9 @@ class RepeaterDaemon: # Stop HTTP server if self.http_server: try: - self.http_server.stop() + await asyncio.wait_for(asyncio.to_thread(self.http_server.stop), timeout=3) + except asyncio.TimeoutError: + logger.warning("Timeout stopping HTTP server") except Exception as e: logger.warning(f"Error stopping HTTP server: {e}") @@ -1073,7 +1079,11 @@ class RepeaterDaemon: # Close storage publishers (MQTT/LetsMesh) to stop their worker threads. try: if self.repeater_handler and self.repeater_handler.storage: - self.repeater_handler.storage.close() + await asyncio.wait_for( + asyncio.to_thread(self.repeater_handler.storage.close), timeout=5 + ) + except asyncio.TimeoutError: + logger.warning("Timeout closing storage publishers") except Exception as e: logger.warning(f"Error closing storage: {e}") @@ -1093,12 +1103,7 @@ class RepeaterDaemon: except Exception as e: logger.debug(f"CH341 reset skipped/failed: {e}") - # Stop the event loop so the process can exit cleanly - try: - loop = asyncio.get_running_loop() - loop.stop() - except RuntimeError: - pass + # Do not force-stop the event loop here; asyncio.run() owns loop lifecycle. @staticmethod def _detect_container() -> bool: @@ -1114,6 +1119,7 @@ class RepeaterDaemon: async def run(self): logger.info("Repeater daemon started") + self._main_task = asyncio.current_task() # Register signal handlers for graceful shutdown loop = asyncio.get_running_loop() @@ -1172,6 +1178,8 @@ class RepeaterDaemon: # Run dispatcher (handles RX/TX via pymc_core) try: await self.dispatcher.run_forever() + except asyncio.CancelledError: + logger.info("Dispatcher loop cancelled for shutdown") except KeyboardInterrupt: logger.info("Shutting down...") for frame_server in getattr(self, "companion_frame_servers", []): diff --git a/repeater/web/update_endpoints.py b/repeater/web/update_endpoints.py index 941fa78..b5171ee 100644 --- a/repeater/web/update_endpoints.py +++ b/repeater/web/update_endpoints.py @@ -45,6 +45,10 @@ PACKAGE_NAME = "pymc_repeater" CHECK_CACHE_TTL = 600 # 10 minutes _github_ssl_ctx: Optional[ssl.SSLContext] = None +_disk_version_mismatch_logged: Optional[tuple] = None +_DISK_VERSION_MISMATCH_LOG_TTL = 300 # seconds +_installed_version_cache: Optional[tuple] = None +_INSTALLED_VERSION_CACHE_TTL = 15 # seconds def _get_github_ssl_context() -> ssl.SSLContext: @@ -61,7 +65,7 @@ class _RateLimitError(Exception): self.reset_at = reset_at -def _get_installed_version() -> str: +def _get_installed_version(force_refresh: bool = False) -> str: """ Return the highest dist-info version found for pymc_repeater across all directories the running interpreter actually uses. @@ -79,6 +83,20 @@ def _get_installed_version() -> str: import site as _site import sys + global _installed_version_cache + now = time.time() + if ( + not force_refresh + and _installed_version_cache is not None + and (now - _installed_version_cache[1]) < _INSTALLED_VERSION_CACHE_TTL + ): + return _installed_version_cache[0] + + def _cache_and_return(value: str) -> str: + global _installed_version_cache + _installed_version_cache = (value, now) + return value + # -- 1. Collect candidate directories ---------------------------------- # dirs: list = [] try: @@ -142,9 +160,9 @@ def _get_installed_version() -> str: if disk_version is None: try: from repeater import __version__ - return __version__ + return _cache_and_return(__version__) except Exception: - return "unknown" + return _cache_and_return("unknown") # -- 5. Sanity check: never return a version older than what's running -- # # If the running process is already on a higher version than anything found @@ -153,17 +171,33 @@ def _get_installed_version() -> str: from repeater import __version__ as _running from packaging.version import Version if Version(_running) > Version(disk_version): - logger.debug( - f"[Update] Disk version {disk_version!r} < running {_running!r};" - " using running __version__ as installed version." - ) + # status() polls can call this frequently; throttle mismatch logs. + global _disk_version_mismatch_logged + now = time.time() + should_log = True + if _disk_version_mismatch_logged is not None: + last_disk, last_running, last_ts = _disk_version_mismatch_logged + if ( + last_disk == disk_version + and last_running == _running + and (now - last_ts) < _DISK_VERSION_MISMATCH_LOG_TTL + ): + should_log = False + + if should_log: + logger.debug( + f"[Update] Disk version {disk_version!r} < running {_running!r};" + " using running __version__ as installed version." + ) + _disk_version_mismatch_logged = (disk_version, _running, now) + # Strip PEP 440 local identifier (+gXXXXXX) – it only encodes # the git hash and causes spurious mismatches with GitHub versions. - return re.sub(r'\+[a-zA-Z0-9.]+$', '', _running) + return _cache_and_return(re.sub(r'\+[a-zA-Z0-9.]+$', '', _running)) except Exception: pass - return re.sub(r'\+[a-zA-Z0-9.]+$', '', disk_version) + return _cache_and_return(re.sub(r'\+[a-zA-Z0-9.]+$', '', disk_version)) # Channels file – persisted so the choice survives daemon restarts _CHANNELS_FILE = "/var/lib/pymc_repeater/.update_channel" @@ -476,7 +510,7 @@ def _parse_dev_number(version_str: str) -> Optional[int]: return int(m.group(1)) if m else None -def _cleanup_stale_dist_info() -> None: +def _cleanup_stale_dist_info(allow_sudo: bool = True) -> None: import glob import shutil import site as _site @@ -524,6 +558,7 @@ def _cleanup_stale_dist_info() -> None: except Exception: return # can't determine winner safely — leave everything alone + removed_any = False for path, ver in found.items(): if path == keep: continue @@ -531,7 +566,13 @@ def _cleanup_stale_dist_info() -> None: shutil.rmtree(path) logger.info(f"[Update] Removed stale dist-info: {path} (version {ver})") _state.append_line(f"[pyMC updater] Removed stale dist-info: {os.path.basename(path)}") + removed_any = True except PermissionError: + if not allow_sudo: + logger.debug( + f"[Update] Skipping stale dist-info cleanup without sudo permissions: {path}" + ) + continue # dist-info is root-owned (pip ran via sudo); use sudo to remove try: subprocess.run( @@ -540,11 +581,28 @@ def _cleanup_stale_dist_info() -> None: ) logger.info(f"[Update] Removed stale dist-info (sudo): {path} (version {ver})") _state.append_line(f"[pyMC updater] Removed stale dist-info: {os.path.basename(path)}") + removed_any = True except Exception as exc2: logger.warning(f"[Update] Could not remove stale dist-info {path}: {exc2}") except Exception as exc: logger.warning(f"[Update] Could not remove stale dist-info {path}: {exc}") + if removed_any: + global _installed_version_cache + _installed_version_cache = None + + +def _startup_dist_info_cleanup() -> None: + """Best-effort cleanup during startup without sudo escalation.""" + try: + _cleanup_stale_dist_info(allow_sudo=False) + fresh = _get_installed_version(force_refresh=True) + if fresh != "unknown": + with _state._lock: + _state.current_version = fresh + except Exception as exc: + logger.debug(f"[Update] Startup dist-info cleanup skipped: {exc}") + def _has_update(installed: str, latest: str) -> bool: """ @@ -814,6 +872,9 @@ def _do_install() -> None: _state.finish_install(False, "pip install failed – see progress log for details") +_startup_dist_info_cleanup() + + # --------------------------------------------------------------------------- # CherryPy Endpoint class # ---------------------------------------------------------------------------