From a6187311cdadb893eb24f1193516bd371cb4a522 Mon Sep 17 00:00:00 2001
From: Mateo <mateo.lostanlen@gmail.com>
Date: Tue, 30 Jul 2024 15:30:45 +0200
Subject: [PATCH] timeout heartbeat (#225)

---
 pyroengine/engine.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/pyroengine/engine.py b/pyroengine/engine.py
index a0412a9..ff5fe16 100644
--- a/pyroengine/engine.py
+++ b/pyroengine/engine.py
@@ -9,6 +9,7 @@
 import logging
 import os
 import shutil
+import signal
 import time
 from collections import deque
 from datetime import datetime, timedelta, timezone
@@ -30,6 +31,23 @@
 logging.basicConfig(format="%(asctime)s | %(levelname)s: %(message)s", level=logging.INFO, force=True)
 
 
+def handler(signum, frame):
+    raise TimeoutError("Heartbeat check timed out")
+
+
+def heartbeat_with_timeout(api_instance, cam_id, timeout=1):
+    signal.signal(signal.SIGALRM, handler)
+    signal.alarm(timeout)
+    try:
+        api_instance.heartbeat(cam_id)
+    except TimeoutError:
+        logging.warning(f"Heartbeat check timed out for {cam_id}")
+    except ConnectionError:
+        logging.warning(f"Unable to reach the pyro-api with {cam_id}")
+    finally:
+        signal.alarm(0)
+
+
 class Engine:
     """This implements an object to manage predictions and API interactions for wildfire alerts.
 
@@ -253,10 +271,7 @@ def predict(self, frame: Image.Image, cam_id: Optional[str] = None) -> float:
 
         # Heartbeat
         if len(self.api_client) > 0 and isinstance(cam_id, str):
-            try:
-                self.heartbeat(cam_id)
-            except ConnectionError:
-                logging.warning(f"Unable to reach the pyro-api with {cam_id}")
+            heartbeat_with_timeout(self, cam_id, timeout=1)
 
         cam_key = cam_id or "-1"
         # Reduce image size to save bandwidth