Add keep-alive mode on target machine.

Author: notsyncing

.dockerignore

@@ -6,5 +6,6 @@ __pycache__/
 .ruff_cache/
 .mypy_cache/
 wolsocketproxy.conf
+wolsocketproxy-keepalive.conf
 *.egg-info/
 .vscode/

.gitignore

@@ -6,5 +6,6 @@ __pycache__/
 .ruff_cache/
 .mypy_cache/
 wolsocketproxy.conf
+wolsocketproxy-keepalive.conf
 *.egg-info/
 .vscode/

docker/docker-compose.yml

@@ -11,3 +11,11 @@ services:
     # If you only use IPMI, then you do not need this
     network_mode: host
 
+  # Keep-alive daemon on target machine
+  wolsocketproxy-keepalive:
+    image: ghcr.io/notsyncing/wol-socket-proxy:main
+    volumes:
+      - /mnt/data/wolsocketproxy/config:/app/config:Z
+    restart: always
+    # Required to override original command
+    command: ["python", "-m", "wolsocketproxy", "-k", "-c", "/app/config/wolsocketproxy.conf"]

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
     "dataclass-wizard==0.39.1",
     "icmplib==3.0.4",
     "redfish==3.3.5",
+    "setproctitle==1.3.7",
     "wakeonlan==3.1.0",
 ]
 
@@ -52,6 +53,7 @@ src = ["src"]
 target-version = "py312"
 lint.select = ["ALL"]
 lint.ignore = [
+    "ASYNC110",
     "COM812",
     "D100",
     "D101",
@@ -62,6 +64,7 @@ lint.ignore = [
     "D203",
     "D211",
     "D213",
+    "DTZ005",
     "S101",
     "EM102",
     "FBT001",

src/wolsocketproxy/__init__.py

@@ -2,18 +2,40 @@
 import logging
 from argparse import ArgumentParser
 from pathlib import Path
+from sys import stdout
 
 import dataclass_wizard
 
+from wolsocketproxy.keepalive import KeepAliveConfig, KeepAliveDaemon
 from wolsocketproxy.proxy import Proxy, ProxyConfig
 
 
 def main() -> None:
-    logging.basicConfig(level=logging.INFO)
+    logging.basicConfig(
+        format="%(asctime)s %(levelname)-8s %(name)s %(message)s",
+        level=logging.INFO,
+        datefmt="%Y-%m-%d %H:%M:%S",
+        handlers=[
+            logging.StreamHandler(
+                stream=stdout,
+            ),
+        ],
+        force=True,
+    )
+
     log = logging.getLogger()
 
     parser = ArgumentParser(prog="wolsocketproxy", description="A socket proxy with wake-on-lan feature.")
 
+    parser.add_argument(
+        "-k",
+        "--keep-alive",
+        dest="keep_alive_mode",
+        default=False,
+        action="store_true",
+        help="Start in keep-alive daemon mode",
+    )
+
     parser.add_argument(
         "-c",
         "--config",
@@ -36,9 +58,18 @@ def main() -> None:
         return
 
     config_data = json.loads(config_path.read_text())
-    config = dataclass_wizard.fromdict(ProxyConfig, config_data)
 
-    log.info("Loaded config from %s", config_path)
+    if args.keep_alive_mode:
+        config = dataclass_wizard.fromdict(KeepAliveConfig, config_data)
+
+        log.info("Loaded keep-alive mode config from %s", config_path)
+
+        keep_alive_daemon = KeepAliveDaemon(config)
+        keep_alive_daemon.start()
+    else:
+        config = dataclass_wizard.fromdict(ProxyConfig, config_data)
+
+        log.info("Loaded proxy mode config from %s", config_path)
 
-    proxy = Proxy(config)
-    proxy.start()
+        proxy = Proxy(config)
+        proxy.start()

src/wolsocketproxy/keepalive.py

@@ -0,0 +1,153 @@
+import asyncio
+import logging
+import os
+import signal
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from threading import Thread
+from typing import Literal
+
+from aiohttp import web
+from aiohttp.web import Request, Response
+from setproctitle import setproctitle
+
+
+@dataclass
+class KeepAliveConfig:
+    listen_address: str = "127.0.0.1"
+    listen_port: int = 18080
+
+    watchdog_feed_interval: int = 600
+
+    keep_alive_method: Literal["special_process"] = "special_process"
+
+    special_process_name: str | None = "wsp-keepalive"
+
+
+class KeepAliveDaemon:
+    _log: logging.Logger = logging.getLogger(__name__)
+
+    _config: KeepAliveConfig
+    _web_app: web.Application
+
+    _watchdog_last_feed_time: datetime
+    _watchdog_is_hungry: bool = False
+
+    _special_process_id: int = -1
+
+    def __init__(self, config: KeepAliveConfig) -> None:
+        self._config = config
+
+        self._web_app = web.Application()
+
+        self._web_app.add_routes(
+            [
+                web.get("/watchdog/feed", self._handle_watchdog_feed)
+            ]
+        )
+
+    async def _watchdog_timer(self) -> None:
+        while True:
+            previous_is_hungry = self._watchdog_is_hungry
+            last_feed_interval = datetime.now() - self._watchdog_last_feed_time
+
+            if last_feed_interval.total_seconds() > self._config.watchdog_feed_interval:
+                self._watchdog_is_hungry = True
+            else:
+                self._watchdog_is_hungry = False
+
+            if not previous_is_hungry and self._watchdog_is_hungry:
+                Thread(target=self._on_watchdog_hungry, daemon=True).start()
+            elif previous_is_hungry and not self._watchdog_is_hungry:
+                Thread(target=self._on_watchdog_feed, daemon=True).start()
+
+            await asyncio.sleep(1)
+
+    def _on_watchdog_hungry(self) -> None:
+        self._log.warning("Watchdog is hungry!")
+
+        if self._special_process_id > 0:
+            os.kill(self._special_process_id, signal.SIGKILL)
+            os.waitpid(self._special_process_id, 0)
+            self._log.warning("Killed special process PID %d", self._special_process_id)
+            self._special_process_id = -1
+
+    def _on_watchdog_feed(self) -> None:
+        self._log.info("Watchdog is feed.")
+
+        if self._config.keep_alive_method == "special_process" and self._special_process_id < 0:
+            self._start_special_process()
+
+    def _special_process(self) -> None:
+        assert self._config.special_process_name is not None
+        setproctitle(self._config.special_process_name)
+
+        while True:
+            time.sleep(60)
+
+            if os.getppid() == 1:
+                os._exit(0)
+
+    def _start_special_process(self) -> bool:
+        self._special_process_id = os.fork()
+
+        if self._special_process_id == 0:
+            self._special_process()
+            return True
+
+        self._log.info(
+            "Started special process with name %s, PID %d",
+            self._config.special_process_name, self._special_process_id
+        )
+
+        return False
+
+    def start(self) -> None:
+        self._watchdog_last_feed_time = datetime.now()
+        self._watchdog_is_hungry = False
+
+        if self._config.keep_alive_method == "special_process":  # noqa: SIM102
+            if self._start_special_process():
+                return
+
+        loop = asyncio.new_event_loop()
+
+        def _start_loop() -> None:
+            asyncio.set_event_loop(loop)
+            loop.run_forever()
+
+        Thread(target=_start_loop, daemon=True).start()
+
+        watchdog_task = asyncio.run_coroutine_threadsafe(self._watchdog_timer(), loop)
+
+        self._log.info(
+            "Keep-alive daemon started at %s:%d, watchdog feed interval %ds.",
+            self._config.listen_address, self._config.listen_port, self._config.watchdog_feed_interval
+        )
+
+        web.run_app(self._web_app, host=self._config.listen_address, port=self._config.listen_port)
+
+        self._log.info("Stopping...")
+
+        watchdog_task.cancel()
+        loop.stop()
+
+        if self._special_process_id > 0:
+            os.kill(self._special_process_id, signal.SIGKILL)
+            os.waitpid(self._special_process_id, 0)
+
+        self._log.info("Stopped.")
+
+    async def _handle_watchdog_feed(self, _: Request) -> Response:
+        last_feed = self._watchdog_last_feed_time
+        self._watchdog_last_feed_time = datetime.now()
+        interval = (self._watchdog_last_feed_time - last_feed).total_seconds()
+
+        return web.json_response(
+            {
+                "status": "ok",
+                "last_feed": last_feed.strftime("%Y-%m-%d %H:%M:%S"),
+                "interval": interval,
+            }
+        )