Fix silent live session stalls

jesse/services/failure.py

@@ -2,57 +2,85 @@
 from jesse.services import logger as jesse_logger
 import threading
 import traceback
+import os
 from jesse.services.redis import sync_publish
 from jesse.repositories import live_session_repository
 from jesse.store import store
 from jesse.enums import live_session_statuses
 
 
+def _terminal_debug(message: str) -> None:
+    try:
+        jh.terminal_debug(message)
+    except Exception:
+        pass
+
+
 def register_custom_exception_handler() -> None:
     # other threads
     def handle_thread_exception(args) -> None:
         if args.exc_type == SystemExit:
             return
 
+        formatted_traceback = ''.join(
+            traceback.format_exception(args.exc_type, args.exc_value, args.exc_traceback)
+        )
+
         if args.exc_type.__name__ == 'Termination':
             sync_publish('termination', {})
             jh.terminate_app()
         else:
             # send notifications if it's a live session
             if jh.is_live():
-                jesse_logger.error(
-                    f'{args.exc_type.__name__}: {args.exc_value}'
-                )
-                jesse_logger.info(
-                    str(traceback.format_exc())
-                )
-
+                try:
+                    jesse_logger.error(
+                        f'{args.exc_type.__name__}: {args.exc_value}'
+                    )
+                    jesse_logger.info(formatted_traceback)
+                except Exception as e:
+                    _terminal_debug(
+                        f'Error logging uncaught thread exception: {type(e).__name__}: {e}\n{formatted_traceback}'
+                    )
+
                 # Store exception in live session
                 try:
                     live_session_repository.store_live_session_exception(
                         store.app.session_id,
                         f"{args.exc_type.__name__}: {str(args.exc_value)}",
-                        str(traceback.format_exc())
+                        formatted_traceback
                     )
                     live_session_repository.update_live_session_status(store.app.session_id, live_session_statuses.STOPPED)
                     live_session_repository.update_live_session_finished(store.app.session_id)
                 except Exception as e:
-                    jh.debug(f'Error storing live session exception: {e}')
+                    _terminal_debug(f'Error storing live session exception: {type(e).__name__}: {e}')
 
-            sync_publish('exception', {
-                'error': f"{args.exc_type.__name__}: {str(args.exc_value)}",
-                'traceback': str(traceback.format_exc())
-            })
-            terminate_session()
+            try:
+                sync_publish('exception', {
+                    'error': f"{args.exc_type.__name__}: {str(args.exc_value)}",
+                    'traceback': formatted_traceback
+                })
+            finally:
+                terminate_session()
 
     threading.excepthook = handle_thread_exception
 
 
 def terminate_session():
-    sync_publish('unexpectedTermination', {
-        'message': "Session terminated as the result of an uncaught exception",
-    })
+    try:
+        sync_publish('unexpectedTermination', {
+            'message': "Session terminated as the result of an uncaught exception",
+        })
+    except Exception as e:
+        _terminal_debug(f'Error publishing unexpected session termination: {type(e).__name__}: {e}')
 
-    jesse_logger.error('Session terminated as the result of an uncaught exception')
+    try:
+        jesse_logger.error('Session terminated as the result of an uncaught exception')
+    except Exception as e:
+        _terminal_debug(f'Error logging unexpected session termination: {type(e).__name__}: {e}')
 
-    jh.terminate_app()
+    try:
+        jh.terminate_app()
+    except BaseException as e:
+        _terminal_debug(f'Error closing resources during session termination: {type(e).__name__}: {e}')
+    finally:
+        os._exit(1)

jesse/services/multiprocessing.py

@@ -14,6 +14,13 @@
 mp.set_start_method('spawn', force=True)
 
 
+def _terminal_debug(message: str) -> None:
+    try:
+        jh.terminal_debug(message)
+    except Exception:
+        pass
+
+
 class Process(mp.Process):
     def __init__(self, *args, **kwargs):
         mp.Process.__init__(self, *args, **kwargs)
@@ -45,6 +52,7 @@ def __init__(self):
         self._workers: List[Process] = []
         self._pid_to_client_id_map = {}
         self.client_id_to_pid_to_map = {}
+        self._pending_worker_removals: set[str] = set()
         self._workers_lock = threading.Lock()
         try:
             port = ENV_VALUES.get('APP_PORT', '9000')
@@ -56,11 +64,23 @@ def __init__(self):
         self._cleanup_thread.start()
 
     def _reset(self):
+        client_ids = {
+            jh.string_after_character(prefixed_client_id, '|')
+            for prefixed_client_id in self.client_id_to_pid_to_map
+        }
+        self._pending_worker_removals.update(client_ids)
         self._workers = []
         self._pid_to_client_id_map = {}
         self.client_id_to_pid_to_map = {}
         # clear all process status
-        sync_redis.delete(self._active_workers_key)
+        try:
+            sync_redis.delete(self._active_workers_key)
+        except Exception as e:
+            _terminal_debug(
+                f'Error clearing active workers from Redis; cleanup will retry: {type(e).__name__}: {e}'
+            )
+        else:
+            self._pending_worker_removals.clear()
 
     @staticmethod
     def _prefixed_pid(pid):
@@ -84,6 +104,9 @@ def add_task(self, function, *args):
             prefixed_client_id = self._prefixed_client_id(client_id)
             self._pid_to_client_id_map[prefixed_pid] = prefixed_client_id
             self.client_id_to_pid_to_map[prefixed_client_id] = prefixed_pid
+            # A new worker owns this Redis marker now, so an older deferred
+            # cleanup must never remove it after a reconnect.
+            self._pending_worker_removals.discard(client_id)
             self._add_process(client_id)
 
     def get_client_id(self, pid):
@@ -119,17 +142,59 @@ def flush(self):
 
             self._reset()
 
+    def _remove_active_worker(self, client_id: str) -> bool:
+        """Remove a finished worker's Redis marker, or retain it for retry."""
+        was_pending = client_id in self._pending_worker_removals
+        try:
+            sync_redis.srem(self._active_workers_key, client_id)
+        except Exception as e:
+            self._pending_worker_removals.add(client_id)
+            if not was_pending:
+                _terminal_debug(
+                    f'Error removing finished worker {client_id} from Redis; cleanup will retry: '
+                    f'{type(e).__name__}: {e}'
+                )
+            return False
+        else:
+            self._pending_worker_removals.discard(client_id)
+            return True
+
+    def _retry_pending_worker_removals(self) -> None:
+        """Retry Redis cleanup without removing markers owned by newer workers."""
+        for client_id in tuple(self._pending_worker_removals):
+            if self._prefixed_client_id(client_id) in self.client_id_to_pid_to_map:
+                self._pending_worker_removals.discard(client_id)
+                continue
+            if self._remove_active_worker(client_id):
+                _terminal_debug(f"Removed deferred worker {client_id} from active workers")
+            else:
+                # One timeout is enough to treat Redis as unavailable for this cycle. Stopping here
+                # bounds how long cleanup holds the worker lock while Redis remains unreachable.
+                break
+
     def _cleanup_finished_workers(self):
         while True:
             try:
                 with self._workers_lock:
+                    self._retry_pending_worker_removals()
                     for w in self._workers[:]:  # Create a copy of the list to avoid modification during iteration
                         if not w.is_alive():
                             try:
                                 prefixed_pid = self._prefixed_pid(w.pid)
                                 prefixed_client_id = self._pid_to_client_id_map.get(prefixed_pid)
 
                                 w.join(timeout=1)
+                                exit_code = w.exitcode
+                                worker_pid = w.pid
+                                client_id = (
+                                    jh.string_after_character(prefixed_client_id, '|')
+                                    if prefixed_client_id
+                                    else 'unknown'
+                                )
+                                _terminal_debug(
+                                    f'Worker {client_id} (PID {worker_pid}) exited with code {exit_code}'
+                                )
+
                                 w.close()
                                 self._workers.remove(w)
                                 self._pid_to_client_id_map.pop(prefixed_pid, None)
@@ -139,13 +204,12 @@ def _cleanup_finished_workers(self):
                                     and self.client_id_to_pid_to_map.get(prefixed_client_id) == prefixed_pid
                                 ):
                                     self.client_id_to_pid_to_map.pop(prefixed_client_id, None)
-                                    client_id = jh.string_after_character(prefixed_client_id, '|')
-                                    sync_redis.srem(self._active_workers_key, client_id)
-                                    jh.debug(f"Removed finished worker {client_id} from active workers")
+                                    if self._remove_active_worker(client_id):
+                                        _terminal_debug(f"Cleaned up finished worker {client_id}")
                             except Exception as e:
-                                jh.debug(f"Error during worker cleanup: {str(e)}")
+                                _terminal_debug(f"Error during worker cleanup: {type(e).__name__}: {e}")
             except Exception as e:
-                jh.debug(f"Error in cleanup thread: {str(e)}")
+                _terminal_debug(f"Error in cleanup thread: {type(e).__name__}: {e}")
             time.sleep(5)
 
     @property

jesse/services/redis.py

@@ -1,7 +1,12 @@
 import aioredis
 import redis as sync_redis_lib
+from redis.exceptions import (
+    ConnectionError as RedisConnectionError,
+    TimeoutError as RedisTimeoutError,
+)
 import simplejson as json
 import asyncio
+import time
 import jesse.helpers as jh
 from jesse.libs.custom_json import NpEncoder
 import os
@@ -19,12 +24,16 @@ async def init_redis():
 
 async_redis = None
 sync_redis = None
+_last_active_check_error_at = 0
 if jh.is_jesse_project():
     if not jh.is_notebook():
         async_redis = asyncio.run(init_redis())
         sync_redis = sync_redis_lib.Redis(
             host=ENV_VALUES['REDIS_HOST'], port=ENV_VALUES['REDIS_PORT'], db=int(ENV_VALUES.get('REDIS_DB') or 0),
-            password=ENV_VALUES['REDIS_PASSWORD'] if ENV_VALUES['REDIS_PASSWORD'] else None
+            password=ENV_VALUES['REDIS_PASSWORD'] if ENV_VALUES['REDIS_PASSWORD'] else None,
+            socket_connect_timeout=1,
+            socket_timeout=1,
+            health_check_interval=30,
         )
 
 
@@ -105,7 +114,23 @@ def get_live_charts_snapshot(session_id: str) -> dict:
 
 
 def is_process_active(client_id: str) -> bool:
+    global _last_active_check_error_at
+
     if jh.is_unit_testing():
         return False
 
-    return sync_redis.sismember(f"{ENV_VALUES['APP_PORT']}|active-processes", client_id)
+    try:
+        is_active = sync_redis.sismember(f"{ENV_VALUES['APP_PORT']}|active-processes", client_id)
+        _last_active_check_error_at = 0
+        return is_active
+    except (RedisConnectionError, RedisTimeoutError, OSError) as e:
+        now = time.monotonic()
+        if _last_active_check_error_at == 0 or now - _last_active_check_error_at >= 30:
+            try:
+                jh.terminal_debug(
+                    f'Redis active-process check failed for {client_id}; keeping the worker active: {type(e).__name__}: {e}'
+                )
+            except Exception:
+                pass
+            _last_active_check_error_at = now
+        return True