Support keep-alive mode in proxy.

Author: notsyncing

README.md

@@ -69,7 +69,10 @@ The `wolsocketproxy.conf` has the following structure:
             "online_check_timeout": 300,                                // Timeout when checking machine is online, seconds, optional, default is 60
             "online_check_http_expected_code": 200,                     // Expected HTTP status code when using "http" method, optional, default is 200
             "ipmi_force_reset_if_power_up_failed": true,                // Use IPMI power reset if this machine is not online after timeout, optional, default is false
-            "ipmi_max_reset_try_count": 3                               // Max IPMI power reset retry count before giving-up, optional, default is 3
+            "ipmi_max_reset_try_count": 3,                              // Max IPMI power reset retry count before giving-up, optional, default is 3
+            "keep_alive_mode": true,                                    // Indicates whether this machine has a keep-alive daemon, optional, default is false
+                                                                        // If you enable this, wolsocketproxy will send a request when there is any traffic
+            "keep_alive_mode_base_url": "http://192.168.1.124:8080"     // Keep-alive daemon provided URL, optional
         },
         // ... more machines ...
     },
@@ -131,3 +134,5 @@ The config file `wolsocketproxy.conf` should look like the following:
 You need to periodically send an HTTP GET request to `/watchdog/feed` at the `listen_port` in `watchdog_feed_interval` time,
 or the special process will be killed.
 You can combine this mode with `circadian`'s `process_block` config to keep it from auto-suspending.
+
+If you enable `keep_alive_mode` in proxy mode's config, the proxy will send requests to `/watchdog/feed` of this machine whenever there is any traffic towards this machine through it.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "wolsocketproxy"
-version = "0.3.0"
+version = "0.3.1"
 description = "A socket proxy with wake-on-lan feature."
 authors = [
     {name = "Song Fuchang", email = "song.fc@gmail.com"}
@@ -72,4 +72,5 @@ lint.ignore = [
     "ISC001",
     "TRY003",
     "TRY201",
+    "PLR0913",
 ]

src/wolsocketproxy/common.py

@@ -0,0 +1 @@
+URL_WATCHDOG_FEED = "/watchdog/feed"

src/wolsocketproxy/keepalive.py

@@ -12,6 +12,8 @@
 from aiohttp.web import Request, Response
 from setproctitle import setproctitle
 
+from wolsocketproxy.common import URL_WATCHDOG_FEED
+
 
 @dataclass
 class KeepAliveConfig:
@@ -41,11 +43,7 @@ def __init__(self, config: KeepAliveConfig) -> None:
 
         self._web_app = web.Application()
 
-        self._web_app.add_routes(
-            [
-                web.get("/watchdog/feed", self._handle_watchdog_feed)
-            ]
-        )
+        self._web_app.add_routes([web.get(URL_WATCHDOG_FEED, self._handle_watchdog_feed)])
 
     async def _watchdog_timer(self) -> None:
         while True:
@@ -97,8 +95,7 @@ def _start_special_process(self) -> bool:
             return True
 
         self._log.info(
-            "Started special process with name %s, PID %d",
-            self._config.special_process_name, self._special_process_id
+            "Started special process with name %s, PID %d", self._config.special_process_name, self._special_process_id
         )
 
         return False
@@ -123,7 +120,9 @@ def _start_loop() -> None:
 
         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
+            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)

src/wolsocketproxy/proxy.py

@@ -3,11 +3,14 @@
 import logging
 from dataclasses import dataclass
 from logging import Logger
+from threading import Thread
 from typing import Any, Literal, override
 
+import aiohttp
 import wakeonlan
 from redfish.rest.v1 import HttpClient, redfish_client
 
+from wolsocketproxy.common import URL_WATCHDOG_FEED
 from wolsocketproxy.monitor import Monitor, MonitorConfig
 from wolsocketproxy.utils import perform_ipmi_action
 
@@ -29,6 +32,10 @@ class MachineConfig:
     online_check_http_expected_code: int = 200
     online_check_timeout: int = 60
 
+    keep_alive_mode: bool = False
+    keep_alive_mode_base_url: str | None = None
+    keep_alive_min_interval: int = 1
+
 
 @dataclass
 class ProxyRoute:
@@ -56,23 +63,67 @@ class ProxyConfig:
     ipmi_configs: list[IPMIConfig] | None = None
 
 
+class TargetKeepAliveSender:
+    _target_url: str
+    _keep_alive_min_interval: int
+    _loop: asyncio.AbstractEventLoop
+    _queue: asyncio.Queue
+
+    def __init__(self, target_base_url: str, keep_alive_min_interval: int) -> None:
+        self._target_url = target_base_url.removesuffix("/") + URL_WATCHDOG_FEED
+        self._keep_alive_min_interval = keep_alive_min_interval
+
+        self._loop = asyncio.new_event_loop()
+        self._queue = asyncio.Queue(1)
+
+        def _loop() -> None:
+            asyncio.set_event_loop(self._loop)
+            self._loop.run_until_complete(self._send_worker())
+
+        Thread(target=_loop, daemon=True).start()
+
+    async def _send_worker(self) -> None:
+        while True:
+            await self._queue.get()
+
+            async with aiohttp.request("GET", self._target_url) as resp:
+                await resp.json()
+
+            await asyncio.sleep(self._keep_alive_min_interval)
+
+    def schedule_send(self) -> None:
+        def _no_exception_put() -> None:
+            with contextlib.suppress(asyncio.QueueFull):
+                self._queue.put_nowait(1)
+
+        self._loop.call_soon_threadsafe(_no_exception_put)
+
+
 class ProxyUdpProtocol(asyncio.DatagramProtocol):
     _proxy: "Proxy"
     _monitor: Monitor
     _transport: asyncio.transports.DatagramTransport
     _target_machine_name: str
     _target_address: str
     _target_port: int
+    _target_keep_alive_sender: TargetKeepAliveSender | None = None
     _target_pair: tuple[str, int]
 
     def __init__(
-        self, proxy: "Proxy", monitor: Monitor, target_machine_name: str, target_address: str, target_port: int
+        self,
+        proxy: "Proxy",
+        monitor: Monitor,
+        target_machine_name: str,
+        target_address: str,
+        target_port: int,
+        target_keep_alive_sender: TargetKeepAliveSender | None = None,
     ) -> None:
         self._proxy = proxy
         self._monitor = monitor
         self._target_machine_name = target_machine_name
         self._target_address = target_address
         self._target_port = target_port
+        self._target_keep_alive_sender = target_keep_alive_sender
         self._target_pair = (target_address, target_port)
 
     @override
@@ -92,6 +143,9 @@ async def handle_datagram(self, data: bytes, addr: tuple[str | Any, int]) -> Non
 
         self._transport.sendto(data, self._target_pair)
 
+        if self._target_keep_alive_sender is not None:
+            self._target_keep_alive_sender.schedule_send()
+
 
 class Proxy:
     _log: Logger = logging.getLogger()
@@ -180,26 +234,55 @@ def __create_route(self, route: ProxyRoute) -> None:
 
     def __create_tcp_route(self, route: ProxyRoute) -> None:
         assert route.target_machine_name is not None
+        machine_config = self._machines[route.target_machine_name]
+        target_keep_alive_sender = None
+
+        if machine_config.keep_alive_mode:
+            assert machine_config.keep_alive_mode_base_url is not None
+
+            target_keep_alive_sender = TargetKeepAliveSender(
+                machine_config.keep_alive_mode_base_url, machine_config.keep_alive_min_interval
+            )
 
         cr = asyncio.start_server(
-            self.__make_tcp_route_handler(route.target_machine_name, route.target_address, route.target_port),
+            self.__make_tcp_route_handler(
+                route.target_machine_name,
+                route.target_address,
+                route.target_port,
+                target_keep_alive_sender,
+            ),
             route.local_address,
             route.local_port,
         )
 
         loop = asyncio.get_event_loop()
         loop.run_until_complete(cr)
 
-    async def __pipe(self, target_address: str, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+    async def __pipe(
+        self,
+        target_address: str,
+        reader: asyncio.StreamReader,
+        writer: asyncio.StreamWriter,
+        target_keep_alive_sender: TargetKeepAliveSender | None = None,
+    ) -> None:
         try:
             while not reader.at_eof():
                 writer.write(await reader.read(2048))
+
+                if target_keep_alive_sender is not None:
+                    target_keep_alive_sender.schedule_send()
         except ConnectionResetError:
             self._log.warning("Connection reset by target %s", target_address)
         finally:
             writer.close()
 
-    def __make_tcp_route_handler(self, target_machine_name: str, target_address: str, target_port: int) -> Any:  # noqa: ANN401
+    def __make_tcp_route_handler(
+        self,
+        target_machine_name: str,
+        target_address: str,
+        target_port: int,
+        target_keep_alive_sender: TargetKeepAliveSender | None = None,
+    ) -> Any:  # noqa: ANN401
         async def handler(local_reader: asyncio.StreamReader, local_writer: asyncio.StreamWriter) -> None:
             if not self._monitor.is_available(target_machine_name):
                 await self._wake_up_target(target_machine_name)
@@ -211,8 +294,8 @@ async def handler(local_reader: asyncio.StreamReader, local_writer: asyncio.Stre
                 self._log.error("Unable to open connection to %s:%d", target_address, target_port)
                 raise e
 
-            send_pipe = self.__pipe(target_address, local_reader, target_writer)
-            recv_pipe = self.__pipe(target_address, target_reader, local_writer)
+            send_pipe = self.__pipe(target_address, local_reader, target_writer, target_keep_alive_sender)
+            recv_pipe = self.__pipe(target_address, target_reader, local_writer, target_keep_alive_sender)
             await asyncio.gather(send_pipe, recv_pipe)
 
         return handler
@@ -221,10 +304,27 @@ def __create_udp_route(self, route: ProxyRoute) -> None:
         target_machine_name = route.target_machine_name
         assert target_machine_name is not None
 
+        machine_config = self._machines[target_machine_name]
+        target_keep_alive_sender = None
+
+        if machine_config.keep_alive_mode:
+            assert machine_config.keep_alive_mode_base_url is not None
+
+            target_keep_alive_sender = TargetKeepAliveSender(
+                machine_config.keep_alive_mode_base_url, machine_config.keep_alive_min_interval
+            )
+
         loop = asyncio.get_event_loop()
 
         cr = loop.create_datagram_endpoint(
-            lambda: ProxyUdpProtocol(self, self._monitor, target_machine_name, route.target_address, route.target_port),
+            lambda: ProxyUdpProtocol(
+                self,
+                self._monitor,
+                target_machine_name,
+                route.target_address,
+                route.target_port,
+                target_keep_alive_sender,
+            ),
             local_addr=(route.local_address, route.local_port),
         )
 

wolsocketproxy.conf.example

@@ -9,7 +9,9 @@
             "online_check_timeout": 300,
             "online_check_http_expected_code": 200,
             "ipmi_force_reset_if_power_up_failed": true,
-            "ipmi_max_reset_try_count": 3
+            "ipmi_max_reset_try_count": 3,
+            "keep_alive_mode": true,
+            "keep_alive_mode_base_url": "http://192.168.1.124:8080"
         }
     },
     "routes": [