add tcp keepalive and option to disable websocket keepalive

Author: Sheeproid

src/robusta/core/model/env_vars.py

@@ -82,6 +82,15 @@ def load_bool(env_var, default: bool):
 # Timeout for the ping response, before killing the connection. Must be smaller than the interval
 WEBSOCKET_PING_TIMEOUT = int(os.environ.get("WEBSOCKET_PING_TIMEOUT", 30))
 
+# TCP keepalive configuration (disabled by default)
+WEBSOCKET_TCP_KEEPALIVE_ENABLED = os.environ.get("WEBSOCKET_TCP_KEEPALIVE_ENABLED", "false").lower() == "true"
+# Time in seconds before sending the first keepalive probe (Linux: TCP_KEEPIDLE, macOS: TCP_KEEPALIVE)
+WEBSOCKET_TCP_KEEPALIVE_IDLE = int(os.environ.get("WEBSOCKET_TCP_KEEPALIVE_IDLE", 2))
+# Interval in seconds between keepalive probes
+WEBSOCKET_TCP_KEEPALIVE_INTERVAL = int(os.environ.get("WEBSOCKET_TCP_KEEPALIVE_INTERVAL", 2))
+# Number of failed probes before connection is considered dead
+WEBSOCKET_TCP_KEEPALIVE_COUNT = int(os.environ.get("WEBSOCKET_TCP_KEEPALIVE_COUNT", 5))
+
 TRACE_INCOMING_REQUESTS = load_bool("TRACE_INCOMING_REQUESTS", False)
 TRACE_INCOMING_ALERTS = load_bool("TRACE_INCOMING_ALERTS", False)
 

src/robusta/integrations/receiver.py

@@ -4,6 +4,8 @@
 import json
 import logging
 import os
+import socket
+import sys
 import time
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import nullcontext
@@ -24,6 +26,10 @@
     SENTRY_ENABLED,
     WEBSOCKET_PING_INTERVAL,
     WEBSOCKET_PING_TIMEOUT,
+    WEBSOCKET_TCP_KEEPALIVE_COUNT,
+    WEBSOCKET_TCP_KEEPALIVE_ENABLED,
+    WEBSOCKET_TCP_KEEPALIVE_IDLE,
+    WEBSOCKET_TCP_KEEPALIVE_INTERVAL, WEBSOCKET_APP_KEEPALIVE_ENABLED,
 )
 from robusta.core.playbooks.playbook_utils import to_safe_str
 from robusta.core.playbooks.playbooks_event_handler import PlaybooksEventHandler
@@ -42,6 +48,22 @@
 WEBSOCKET_THREADPOOL_SIZE = int(os.environ.get("WEBSOCKET_THREADPOOL_SIZE", 10))
 
 
+def _get_tcp_keepalive_options() -> tuple:
+    """Build TCP keepalive socket options tuple for run_forever(sockopt=...)."""
+    # TCP_KEEPIDLE is Linux-only; macOS uses TCP_KEEPALIVE (0x10) for the same purpose
+    if sys.platform == "darwin":
+        tcp_keepalive_idle = 0x10
+    else:
+        tcp_keepalive_idle = socket.TCP_KEEPIDLE
+
+    return (
+        (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+        (socket.IPPROTO_TCP, tcp_keepalive_idle, WEBSOCKET_TCP_KEEPALIVE_IDLE),
+        (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, WEBSOCKET_TCP_KEEPALIVE_INTERVAL),
+        (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, WEBSOCKET_TCP_KEEPALIVE_COUNT),
+    )
+
+
 class ValidationResponse(BaseModel):
     http_code: int = 200
     error_code: Optional[int] = None
@@ -114,11 +136,22 @@ def start_receiver(self):
 
     def run_forever(self):
         logging.info("starting relay receiver")
+        sockopt = None
+        if WEBSOCKET_TCP_KEEPALIVE_ENABLED:
+            sockopt = _get_tcp_keepalive_options()
+            logging.info(
+                f"TCP keepalive enabled: idle={WEBSOCKET_TCP_KEEPALIVE_IDLE}s, "
+                f"interval={WEBSOCKET_TCP_KEEPALIVE_INTERVAL}s, count={WEBSOCKET_TCP_KEEPALIVE_COUNT}"
+            )
         while self.active:
+            # Handles WEBSOCKET_PING_INTERVAL == 0
+            ping_timeout = WEBSOCKET_PING_TIMEOUT if WEBSOCKET_PING_INTERVAL else None
+            logging.info("relay websocket starting")
             self.ws.run_forever(
                 ping_interval=WEBSOCKET_PING_INTERVAL,
                 ping_payload="p",
-                ping_timeout=WEBSOCKET_PING_TIMEOUT,
+                ping_timeout=ping_timeout,
+                sockopt=sockopt,
             )
             logging.info("relay websocket closed")
             time.sleep(INCOMING_WEBSOCKET_RECONNECT_DELAY_SEC)