feat: graceful shutdown with store persistence and restart overlay

- Add /api/v1/system/shutdown endpoint that triggers clean uvicorn exit - Persist all 15 stores to disk during shutdown via _save_all_stores() - Add force parameter to BaseJsonStore._save() to bypass restore freeze - Restart script now requests graceful shutdown via API (15s timeout), falls back to force-kill only if server doesn't exit in time - Broadcast server_restarting event over WebSocket before shutdown - Frontend shows "Server restarting..." overlay instantly on WS event, replacing the old dynamically-created overlay from settings.ts - Add server_ref module to share uvicorn Server + TrayManager refs - Add i18n keys for restart overlay (en/ru/zh)
2026-03-24 15:50:32 +03:00
parent 73947eb6cb
commit 9b4dbac088
14 changed files with 242 additions and 72 deletions
--- a/server/restart.ps1
+++ b/server/restart.ps1
@@ -1,12 +1,76 @@
 # Restart the WLED Screen Controller server
-# Stop any running instance
+# Uses graceful shutdown first (lets the server persist data to disk),
+# then force-kills as a fallback.
+
+$serverRoot = 'c:\Users\Alexei\Documents\wled-screen-controller\server'
+
+# Read API key from config for authenticated shutdown request
+$configPath = Join-Path $serverRoot 'config\default_config.yaml'
+$apiKey = $null
+if (Test-Path $configPath) {
+    $inKeys = $false
+    foreach ($line in Get-Content $configPath) {
+        if ($line -match '^\s*api_keys:') { $inKeys = $true; continue }
+        if ($inKeys -and $line -match '^\s+\w+:\s*"(.+)"') {
+            $apiKey = $Matches[1]; break
+        }
+        if ($inKeys -and $line -match '^\S') { break }  # left the api_keys block
+    }
+}
+
+# Find running server processes
 $procs = Get-CimInstance Win32_Process -Filter "Name='python.exe'" |
    Where-Object { $_.CommandLine -like '*wled_controller*' -and $_.CommandLine -notlike '*demo*' -and $_.CommandLine -notlike '*vscode*' -and $_.CommandLine -notlike '*isort*' }
-foreach ($p in $procs) {
+
+if ($procs) {
+    # Step 1: Request graceful shutdown via API (triggers lifespan shutdown + store save)
+    $shutdownOk = $false
+    if ($apiKey) {
+        Write-Host "Requesting graceful shutdown..."
+        try {
+            $headers = @{ Authorization = "Bearer $apiKey" }
+            Invoke-RestMethod -Uri 'http://localhost:8080/api/v1/system/shutdown' `
+                -Method Post -Headers $headers -TimeoutSec 5 -ErrorAction Stop | Out-Null
+            $shutdownOk = $true
+        } catch {
+            Write-Host "  API shutdown failed ($($_.Exception.Message)), falling back to process kill"
+        }
+    }
+
+    if ($shutdownOk) {
+        # Step 2: Wait for the server to exit gracefully (up to 15 seconds)
+        # The server needs time to stop processors, disconnect devices, and persist stores.
+        Write-Host "Waiting for graceful shutdown..."
+        $waited = 0
+        while ($waited -lt 15) {
+            Start-Sleep -Seconds 1
+            $waited++
+            $still = Get-CimInstance Win32_Process -Filter "Name='python.exe'" |
+                Where-Object { $_.CommandLine -like '*wled_controller*' -and $_.CommandLine -notlike '*demo*' -and $_.CommandLine -notlike '*vscode*' -and $_.CommandLine -notlike '*isort*' }
+            if (-not $still) {
+                Write-Host "  Server exited cleanly after ${waited}s"
+                break
+            }
+        }
+        # Step 3: Force-kill stragglers
+        $still = Get-CimInstance Win32_Process -Filter "Name='python.exe'" |
+            Where-Object { $_.CommandLine -like '*wled_controller*' -and $_.CommandLine -notlike '*demo*' -and $_.CommandLine -notlike '*vscode*' -and $_.CommandLine -notlike '*isort*' }
+        if ($still) {
+            Write-Host "  Force-killing remaining processes..."
+            foreach ($p in $still) {
+                Stop-Process -Id $p.ProcessId -Force -ErrorAction SilentlyContinue
+            }
+            Start-Sleep -Seconds 1
+        }
+    } else {
+        # No API key or API call failed — force-kill directly
+        foreach ($p in $procs) {
            Write-Host "Stopping server (PID $($p.ProcessId))..."
            Stop-Process -Id $p.ProcessId -Force -ErrorAction SilentlyContinue
+        }
+        Start-Sleep -Seconds 2
+    }
 }
-if ($procs) { Start-Sleep -Seconds 2 }

 # Merge registry PATH with current PATH so newly-installed tools (e.g. scrcpy) are visible
 $regUser = [Environment]::GetEnvironmentVariable('PATH', 'User')
@@ -19,15 +83,16 @@ if ($regUser) {
    }
 }

-# Start server detached
+# Start server detached (set WLED_RESTART=1 to skip browser open)
 Write-Host "Starting server..."
+$env:WLED_RESTART = "1"
 $pythonExe = (Get-Command python -ErrorAction SilentlyContinue).Source
 if (-not $pythonExe) {
    # Fallback to known install location
    $pythonExe = "$env:LOCALAPPDATA\Programs\Python\Python313\python.exe"
 }
 Start-Process -FilePath $pythonExe -ArgumentList '-m', 'wled_controller' `
-    -WorkingDirectory 'c:\Users\Alexei\Documents\wled-screen-controller\server' `
+    -WorkingDirectory $serverRoot `
    -WindowStyle Hidden

 Start-Sleep -Seconds 3
--- a/server/src/wled_controller/main.py
+++ b/server/src/wled_controller/main.py
@@ -15,6 +15,7 @@ from pathlib import Path
 import uvicorn

 from wled_controller.config import get_config
+from wled_controller.server_ref import set_server, set_tray
 from wled_controller.tray import PYSTRAY_AVAILABLE, TrayManager
 from wled_controller.utils import setup_logging, get_logger

@@ -52,6 +53,7 @@ def main() -> None:
        log_level=config.server.log_level.lower(),
    )
    server = uvicorn.Server(uv_config)
+    set_server(server)

    use_tray = PYSTRAY_AVAILABLE and (
        sys.platform == "win32" or _force_tray()
@@ -80,6 +82,7 @@ def main() -> None:
            port=config.server.port,
            on_exit=lambda: _request_shutdown(server),
        )
+        set_tray(tray)
        tray.run()

        # Tray exited — wait for server to finish its graceful shutdown
--- a/server/src/wled_controller/api/routes/backup.py
+++ b/server/src/wled_controller/api/routes/backup.py
@@ -228,10 +228,25 @@ def backup_config(_: AuthRequired):
@router.post("/api/v1/system/restart", tags=["System"])
 def restart_server(_: AuthRequired):
    """Schedule a server restart and return immediately."""
+    from wled_controller.server_ref import _broadcast_restarting
+    _broadcast_restarting()
    _schedule_restart()
    return {"status": "restarting"}


+@router.post("/api/v1/system/shutdown", tags=["System"])
+def shutdown_server(_: AuthRequired):
+    """Gracefully shut down the server.
+
+    Signals uvicorn to exit, which triggers the lifespan shutdown handler
+    (persists all stores to disk, stops processors, etc.).
+    Used by the restart script to ensure data is saved before the process exits.
+    """
+    from wled_controller.server_ref import request_shutdown
+    request_shutdown()
+    return {"status": "shutting_down"}
+
+
@router.post("/api/v1/system/restore", response_model=RestoreResponse, tags=["System"])
 async def restore_config(
    _: AuthRequired,
--- a/server/src/wled_controller/main.py
+++ b/server/src/wled_controller/main.py
@@ -92,6 +92,29 @@ processor_manager = ProcessorManager(
 )


+def _save_all_stores() -> None:
+    """Persist every store to disk.
+
+    Called during graceful shutdown to ensure in-memory data survives
+    restarts even if no CRUD happened during the session.
+    """
+    all_stores = [
+        device_store, template_store, pp_template_store,
+        picture_source_store, output_target_store, pattern_template_store,
+        color_strip_store, audio_source_store, audio_template_store,
+        value_source_store, automation_store, scene_preset_store,
+        sync_clock_store, cspt_store, gradient_store,
+    ]
+    saved = 0
+    for store in all_stores:
+        try:
+            store._save(force=True)
+            saved += 1
+        except Exception as e:
+            logger.error(f"Failed to save {store._json_key} on shutdown: {e}")
+    logger.info(f"Shutdown save: persisted {saved}/{len(all_stores)} stores to disk")
+
+
@asynccontextmanager
 async def lifespan(app: FastAPI):
    """Application lifespan manager.
@@ -230,6 +253,11 @@ async def lifespan(app: FastAPI):
    # Shutdown
    logger.info("Shutting down LED Grab")

+    # Persist all stores to disk before stopping anything.
+    # This ensures in-memory data survives force-kills and restarts
+    # where no CRUD happened during the session.
+    _save_all_stores()
+
    # Stop auto-backup engine
    try:
        await auto_backup_engine.stop()
--- a/server/src/wled_controller/server_ref.py
+++ b/server/src/wled_controller/server_ref.py
@@ -0,0 +1,59 @@
+"""Module-level holder for the uvicorn Server and TrayManager references.
+
+Allows the shutdown API endpoint to trigger graceful shutdown via
+``server.should_exit = True`` + ``tray.stop()``, which is the same
+mechanism the system tray "Shutdown" menu item uses.
+"""
+
+from typing import Any, Optional
+
+_server: Optional[Any] = None  # uvicorn.Server
+_tray: Optional[Any] = None    # TrayManager
+
+
+def set_server(server: Any) -> None:
+    """Store the uvicorn Server instance (called from __main__)."""
+    global _server
+    _server = server
+
+
+def set_tray(tray: Any) -> None:
+    """Store the TrayManager instance (called from __main__)."""
+    global _tray
+    _tray = tray
+
+
+def request_shutdown() -> None:
+    """Signal uvicorn + tray to perform a graceful shutdown.
+
+    Broadcasts a ``server_restarting`` event so the frontend can show
+    a restart indicator, persists all stores to disk, then sets
+    ``should_exit = True`` on the uvicorn Server and stops the tray.
+    """
+    # Notify connected clients that a restart is in progress
+    _broadcast_restarting()
+
+    # Persist stores before signaling shutdown.
+    # The lifespan shutdown handler also saves, but it may not run
+    # reliably when uvicorn is in a daemon thread.
+    try:
+        from wled_controller.main import _save_all_stores
+        _save_all_stores()
+    except Exception:
+        pass  # best-effort; lifespan handler is the backup
+
+    if _server is not None:
+        _server.should_exit = True
+    if _tray is not None:
+        _tray.stop()
+
+
+def _broadcast_restarting() -> None:
+    """Push a server_restarting event to all connected WebSocket clients."""
+    try:
+        from wled_controller.api.dependencies import _deps
+        pm = _deps.get("processor_manager")
+        if pm is not None:
+            pm.fire_event({"type": "server_restarting"})
+    except Exception:
+        pass
--- a/server/src/wled_controller/static/js/core/api.ts
+++ b/server/src/wled_controller/static/js/core/api.ts
@@ -6,6 +6,7 @@ import { apiKey, setApiKey, authRequired, setAuthRequired, refreshInterval, setR
 import { t } from './i18n.ts';
 import { showToast } from './ui.ts';
 import { getEl, queryEl } from './dom-utils.ts';
+import { serverRestarting, clearRestartingFlag } from './events-ws.ts';

 export const API_BASE = '/api/v1';

@@ -168,6 +169,30 @@ export function handle401Error() {
 let _connCheckTimer: ReturnType<typeof setInterval> | null = null;
 let _serverOnline: boolean | null = null;  // null = unknown, true/false

+/** Toggle which message block is visible inside the connection overlay. */
+function _setOverlayMode(restarting: boolean) {
+    const msgOffline = document.getElementById('conn-msg-offline');
+    const msgRestarting = document.getElementById('conn-msg-restarting');
+    if (msgOffline) msgOffline.style.display = restarting ? 'none' : '';
+    if (msgRestarting) msgRestarting.style.display = restarting ? '' : 'none';
+}
+
+/**
+ * Show the restart overlay immediately (called when server_restarting
+ * event arrives via WebSocket, before the connection actually drops).
+ */
+export function showRestartingOverlay() {
+    _serverOnline = false;
+    const banner = document.getElementById('connection-overlay');
+    const badge = document.getElementById('server-status');
+    if (banner) {
+        (banner as HTMLElement).style.display = 'flex';
+        banner.setAttribute('aria-hidden', 'false');
+    }
+    _setOverlayMode(true);
+    if (badge) badge.className = 'status-badge offline';
+}
+
 function _setConnectionState(online: boolean) {
    const changed = _serverOnline !== online;
    _serverOnline = online;
@@ -176,8 +201,14 @@ function _setConnectionState(online: boolean) {
    if (online) {
        if (banner) { (banner as HTMLElement).style.display = 'none'; banner.setAttribute('aria-hidden', 'true'); }
        if (badge) badge.className = 'status-badge online';
+        // Clear the restarting flag once the server is back
+        if (serverRestarting) clearRestartingFlag();
    } else {
-        if (banner) { (banner as HTMLElement).style.display = 'flex'; banner.setAttribute('aria-hidden', 'false'); }
+        if (banner) {
+            (banner as HTMLElement).style.display = 'flex';
+            banner.setAttribute('aria-hidden', 'false');
+        }
+        _setOverlayMode(serverRestarting);
        if (badge) badge.className = 'status-badge offline';
    }
    return changed;
--- a/server/src/wled_controller/static/js/core/events-ws.ts
+++ b/server/src/wled_controller/static/js/core/events-ws.ts
@@ -10,6 +10,14 @@
 */

 import { apiKey, authRequired } from './state.ts';
+import { showRestartingOverlay } from './api.ts';
+
+/** True when the server has signalled it is restarting (not crashed). */
+export let serverRestarting = false;
+
+export function clearRestartingFlag() {
+    serverRestarting = false;
+}

 let _ws: WebSocket | null = null;
 let _reconnectTimer: ReturnType<typeof setTimeout> | null = null;
@@ -32,6 +40,10 @@ export function startEventsWS() {
        _ws.onmessage = (event) => {
            try {
                const data = JSON.parse(event.data);
+                if (data.type === 'server_restarting') {
+                    serverRestarting = true;
+                    showRestartingOverlay();
+                }
                document.dispatchEvent(new CustomEvent(`server:${data.type}`, { detail: data }));
            } catch {}
        };
--- a/server/src/wled_controller/static/js/features/settings.ts
+++ b/server/src/wled_controller/static/js/features/settings.ts
@@ -347,9 +347,6 @@ export async function handleRestoreFileSelected(input: HTMLInputElement): Promis
        showToast(data.message || t('settings.restore.success'), 'success');
        settingsModal.forceClose();

-        if (data.restart_scheduled) {
-            showRestartOverlay();
-        }
    } catch (err) {
        console.error('Restore failed:', err);
        showToast(t('settings.restore.error') + ': ' + err.message, 'error');
@@ -369,63 +366,12 @@ export async function restartServer(): Promise<void> {
            throw new Error(err.detail || `HTTP ${resp.status}`);
        }
        settingsModal.forceClose();
-        showRestartOverlay(t('settings.restarting'));
    } catch (err) {
        console.error('Server restart failed:', err);
        showToast(t('settings.restore.error') + ': ' + err.message, 'error');
    }
 }

-// ─── Restart overlay ───────────────────────────────────────
-
-function showRestartOverlay(message?: string): void {
-    const msg = message || t('settings.restore.restarting');
-    const overlay = document.createElement('div');
-    overlay.id = 'restart-overlay';
-    overlay.style.cssText =
-        'position:fixed;inset:0;z-index:100000;display:flex;flex-direction:column;' +
-        'align-items:center;justify-content:center;background:rgba(0,0,0,0.85);color:#fff;font-size:1.2rem;';
-    overlay.innerHTML =
-        '<div class="spinner" style="width:48px;height:48px;border:4px solid rgba(255,255,255,0.3);' +
-        'border-top-color:#fff;border-radius:50%;animation:spin 0.8s linear infinite;margin-bottom:1rem;"></div>' +
-        `<div id="restart-msg">${msg}</div>`;
-
-    // Add spinner animation if not present
-    if (!document.getElementById('restart-spinner-style')) {
-        const style = document.createElement('style');
-        style.id = 'restart-spinner-style';
-        style.textContent = '@keyframes spin{to{transform:rotate(360deg)}}';
-        document.head.appendChild(style);
-    }
-
-    document.body.appendChild(overlay);
-    pollHealth();
-}
-
-function pollHealth(): void {
-    const start = Date.now();
-    const maxWait = 30000;
-    const interval = 1500;
-
-    const check = async () => {
-        if (Date.now() - start > maxWait) {
-            const msg = document.getElementById('restart-msg');
-            if (msg) msg.textContent = t('settings.restore.restart_timeout');
-            return;
-        }
-        try {
-            const resp = await fetch('/health', { signal: AbortSignal.timeout(3000) });
-            if (resp.ok) {
-                window.location.reload();
-                return;
-            }
-        } catch { /* server still down */ }
-        setTimeout(check, interval);
-    };
-    // Wait a moment before first check to let the server shut down
-    setTimeout(check, 2000);
-}
-
 // ─── Auto-Backup settings ─────────────────────────────────

 export async function loadAutoBackupSettings(): Promise<void> {
@@ -567,9 +513,6 @@ export async function restoreSavedBackup(filename: string): Promise<void> {
        showToast(data.message || t('settings.restore.success'), 'success');
        settingsModal.forceClose();

-        if (data.restart_scheduled) {
-            showRestartOverlay();
-        }
    } catch (err) {
        console.error('Restore from saved backup failed:', err);
        showToast(t('settings.restore.error') + ': ' + err.message, 'error');
@@ -687,9 +630,6 @@ export async function handlePartialImportFileSelected(input: HTMLInputElement):
        showToast(data.message || t('settings.partial.import_success'), 'success');
        settingsModal.forceClose();

-        if (data.restart_scheduled) {
-            showRestartOverlay();
-        }
    } catch (err) {
        console.error('Partial import failed:', err);
        showToast(t('settings.partial.import_error') + ': ' + err.message, 'error');
--- a/server/src/wled_controller/static/locales/en.json
+++ b/server/src/wled_controller/static/locales/en.json
@@ -4,6 +4,8 @@
    "app.api_docs": "API Documentation",
    "app.connection_lost": "Server unreachable",
    "app.connection_retrying": "Attempting to reconnect…",
+    "app.server_restarting": "Server restarting…",
+    "app.server_restarting_sub": "Please wait, the server will be back shortly.",
    "demo.badge": "DEMO",
    "demo.banner": "You're in demo mode — all devices and data are virtual. No real hardware is used.",
    "theme.toggle": "Toggle theme",
--- a/server/src/wled_controller/static/locales/ru.json
+++ b/server/src/wled_controller/static/locales/ru.json
@@ -4,6 +4,8 @@
    "app.api_docs": "Документация API",
    "app.connection_lost": "Сервер недоступен",
    "app.connection_retrying": "Попытка переподключения…",
+    "app.server_restarting": "Сервер перезапускается…",
+    "app.server_restarting_sub": "Пожалуйста, подождите, сервер скоро вернётся.",
    "demo.badge": "ДЕМО",
    "demo.banner": "Вы в демо-режиме — все устройства и данные виртуальные. Реальное оборудование не используется.",
    "theme.toggle": "Переключить тему",
--- a/server/src/wled_controller/static/locales/zh.json
+++ b/server/src/wled_controller/static/locales/zh.json
@@ -4,6 +4,8 @@
    "app.api_docs": "API 文档",
    "app.connection_lost": "服务器不可达",
    "app.connection_retrying": "正在尝试重新连接…",
+    "app.server_restarting": "服务器正在重启…",
+    "app.server_restarting_sub": "请稍候，服务器即将恢复。",
    "demo.badge": "演示",
    "demo.banner": "您正处于演示模式 — 所有设备和数据均为虚拟。未使用任何真实硬件。",
    "theme.toggle": "切换主题",
--- a/server/src/wled_controller/storage/base_store.py
+++ b/server/src/wled_controller/storage/base_store.py
@@ -98,15 +98,18 @@ class BaseJsonStore(Generic[T]):
            f"{self._entity_name} store initialized with {len(self._items)} items"
        )

-    def _save(self) -> None:
+    def _save(self, *, force: bool = False) -> None:
        """Persist all items to disk atomically.

        Note: This is synchronous blocking I/O. When called from async route
        handlers, it briefly blocks the event loop (typically < 5ms for small
        stores). Acceptable for user-initiated CRUD; not suitable for hot loops.
        Callers must hold ``self._lock``.
+
+        Args:
+            force: bypass the ``_saves_frozen`` flag (used by shutdown save).
        """
-        if _saves_frozen:
+        if _saves_frozen and not force:
            logger.warning(f"Save blocked (frozen after restore): {self._json_key}")
            return
        try:
--- a/server/src/wled_controller/templates/index.html
+++ b/server/src/wled_controller/templates/index.html
@@ -25,9 +25,15 @@
        <div id="connection-overlay" class="connection-overlay" style="display:none" aria-hidden="true">
            <div class="connection-overlay-content">
                <div class="connection-spinner-lg"></div>
+                <div id="conn-msg-offline">
                    <h2 data-i18n="app.connection_lost">Server unreachable</h2>
                    <p data-i18n="app.connection_retrying">Attempting to reconnect…</p>
                </div>
+                <div id="conn-msg-restarting" style="display:none">
+                    <h2 data-i18n="app.server_restarting">Server restarting…</h2>
+                    <p data-i18n="app.server_restarting_sub">Please wait, the server will be back shortly.</p>
+                </div>
+            </div>
        </div>
        <header>
            <div class="header-title">
--- a/server/src/wled_controller/tray.py
+++ b/server/src/wled_controller/tray.py
@@ -63,6 +63,8 @@ class TrayManager:
    def _restart(self, icon: "pystray.Icon", item: "pystray.MenuItem") -> None:
        if not messagebox.askyesno("LED Grab", "Restart the server?"):
            return
+        from wled_controller.server_ref import _broadcast_restarting
+        _broadcast_restarting()
        self._icon.stop()
        self._on_exit()
        os.environ["WLED_RESTART"] = "1"