fix(utils): Improve exec_in_proc to handle more failure modes

Author: marcelklehr

context_chat_backend/utils.py

@@ -9,6 +9,7 @@
 import sys
 import traceback
 from collections.abc import Callable
+from contextlib import suppress
 from functools import partial, wraps
 from multiprocessing.connection import Connection
 from time import perf_counter_ns
@@ -72,31 +73,95 @@ def JSONResponse(
 
 
 class SubprocessKilledError(RuntimeError):
-	"""Raised when a subprocess exits with a non-zero exit code (likely OOM kill or unhandled signal)."""
+	"""Raised when a subprocess is terminated by a signal (for example SIGKILL)."""
 
 	def __init__(self, pid: int, target_name: str, exitcode: int):
 		super().__init__(
-			f'Subprocess PID {pid} for {target_name} exited with non-zero exit code {exitcode}'
-			' — possible OOM kill or unhandled signal'
+			f'Subprocess PID {pid} for {target_name} exited with signal {abs(exitcode)} '
+			f'(raw exit code: {exitcode})'
 		)
 		self.exitcode = exitcode
 
 
+class SubprocessExecutionError(RuntimeError):
+	"""Raised when a subprocess exits non-zero without a recoverable Python exception payload."""
+
+	def __init__(self, pid: int, target_name: str, exitcode: int, details: str = ''):
+		msg = f'Subprocess PID {pid} for {target_name} exited with non-zero exit code {exitcode}'
+		if details:
+			msg = f'{msg}: {details}'
+		super().__init__(msg)
+		self.exitcode = exitcode
+
+
+_MAX_STD_CAPTURE_CHARS = 64 * 1024
+
+
+def _truncate_capture(text: str) -> tuple[str, bool]:
+	if len(text) <= _MAX_STD_CAPTURE_CHARS:
+		return text, False
+
+	head = _MAX_STD_CAPTURE_CHARS // 2
+	tail = _MAX_STD_CAPTURE_CHARS - head
+	omitted = len(text) - _MAX_STD_CAPTURE_CHARS
+	truncated = (
+		f'[truncated {omitted} chars]\n'
+		f'{text[:head]}\n'
+		'[...snip...]\n'
+		f'{text[-tail:]}'
+	)
+	return truncated, True
+
+
 def exception_wrap(fun: Callable | None, *args, resconn: Connection, stdconn: Connection, **kwargs):
 	stdout_capture = io.StringIO()
 	stderr_capture = io.StringIO()
+	orig_stdout = sys.stdout
+	orig_stderr = sys.stderr
 	sys.stdout = stdout_capture
 	sys.stderr = stderr_capture
 
 	try:
 		if fun is None:
-			return resconn.send({ 'value': None, 'error': None })
-		resconn.send({ 'value': fun(*args, **kwargs), 'error': None })
+			resconn.send({ 'value': None, 'error': None })
+		else:
+			resconn.send({ 'value': fun(*args, **kwargs), 'error': None })
 	except BaseException as e:
 		tb = traceback.format_exc()
-		resconn.send({ 'value': None, 'error': e, 'traceback': tb })
+		payload = {
+			'value': None,
+			'error': e,
+			'traceback': tb,
+			'error_type': type(e).__name__,
+			'error_module': type(e).__module__,
+			'error_message': str(e),
+		}
+		try:
+			resconn.send(payload)
+		except Exception as send_err:
+			# Fallback for unpicklable exceptions.
+			with suppress(Exception):
+				resconn.send({
+					'value': None,
+					'error': None,
+					'traceback': tb,
+					'error_type': type(e).__name__,
+					'error_module': type(e).__module__,
+					'error_message': str(e),
+					'send_error': str(send_err),
+				})
 	finally:
-		stdconn.send({'stdout': stdout_capture.getvalue(), 'stderr': stderr_capture.getvalue()})
+		sys.stdout = orig_stdout
+		sys.stderr = orig_stderr
+		stdout_text, stdout_truncated = _truncate_capture(stdout_capture.getvalue())
+		stderr_text, stderr_truncated = _truncate_capture(stderr_capture.getvalue())
+		with suppress(Exception):
+			stdconn.send({
+				'stdout': stdout_text,
+				'stderr': stderr_text,
+				'stdout_truncated': stdout_truncated,
+				'stderr_truncated': stderr_truncated,
+			})
 
 
 def exec_in_proc(group=None, target=None, name=None, args=(), kwargs={}, *, daemon=None):  # noqa: B006
@@ -117,30 +182,93 @@ def exec_in_proc(group=None, target=None, name=None, args=(), kwargs={}, *, daem
 	start = perf_counter_ns()
 	p.start()
 	_logger.debug('Subprocess PID %d started for %s, waiting for it to finish (no timeout)', p.pid, target_name)
+
+	result = None
+	stdobj = {
+		'stdout': '',
+		'stderr': '',
+		'stdout_truncated': False,
+		'stderr_truncated': False,
+	}
+	got_result = False
+	got_std = False
+
+	# Drain result/std pipes while child is still alive to avoid deadlock on full pipe buffers.
+	while p.is_alive() and (not got_result or not got_std):
+		if not got_result and pconn.poll(0.1):
+			with suppress(EOFError, OSError, BrokenPipeError):
+				result = pconn.recv()
+				got_result = True
+		if not got_std and std_pconn.poll():
+			with suppress(EOFError, OSError, BrokenPipeError):
+				stdobj = std_pconn.recv()
+				got_std = True
+
 	p.join()
 	elapsed_ms = (perf_counter_ns() - start) / 1e6
 	_logger.debug(
 		'Subprocess PID %d for %s finished in %.2f ms (exit code: %s)',
 		p.pid, target_name, elapsed_ms, p.exitcode,
 	)
-	stdobj = std_pconn.recv()
-	_logger.info(f'std info for {target_name}', extra={
-		'stdout': stdobj['stdout'],
-		'stderr': stdobj['stderr'],
-	})
-
-	result = pconn.recv()
-	if result['error'] is not None:
-		_logger.error('original traceback: %s', result['traceback'])
+
+	if not got_std:
+		with suppress(EOFError, OSError, BrokenPipeError):
+			if std_pconn.poll():
+				stdobj = std_pconn.recv()
+				got_std = True
+	if stdobj['stdout'] or stdobj['stderr']:
+		extra = {
+			'stdout': stdobj['stdout'],
+			'stderr': stdobj['stderr'],
+		}
+		if stdobj.get('stdout_truncated') or stdobj.get('stderr_truncated'):
+			extra['stdio_truncated'] = {
+				'stdout': bool(stdobj.get('stdout_truncated')),
+				'stderr': bool(stdobj.get('stderr_truncated')),
+			}
+		_logger.info('std info for %s', target_name, extra=extra)
+
+	if not got_result:
+		with suppress(EOFError, OSError, BrokenPipeError):
+			if pconn.poll():
+				result = pconn.recv()
+				got_result = True
+
+	if result is not None and result.get('error') is not None:
+		_logger.error('original traceback: %s', result.get('traceback', ''))
 		raise result['error']
 
-	if p.exitcode != 0:
+	if result is not None and result.get('error_type'):
+		details = (
+			f"{result.get('error_module', '')}.{result.get('error_type', '')}: "
+			f"{result.get('error_message', '')}"
+		)
+		if result.get('traceback'):
+			_logger.error('remote traceback: %s', result['traceback'])
+		raise SubprocessExecutionError(p.pid or 0, target_name, p.exitcode or 1, details)
+
+	if p.exitcode and p.exitcode < 0:
 		_logger.warning(
-			'Subprocess PID %d for %s exited with non-zero exit code %d after %.2f ms'
-			' — possible OOM kill or unhandled signal',
-			p.pid, target_name, p.exitcode, elapsed_ms,
+			'Subprocess PID %d for %s exited due to signal %d after %.2f ms',
+			p.pid, target_name, abs(p.exitcode), elapsed_ms,
+		)
+		raise SubprocessKilledError(p.pid or 0, target_name, p.exitcode)
+
+	if p.exitcode not in (None, 0):
+		raise SubprocessExecutionError(
+			p.pid or 0,
+			target_name,
+			p.exitcode,
+			'No structured exception payload received from child process',
+		)
+
+	if result is None:
+		raise SubprocessExecutionError(
+			p.pid or 0,
+			target_name,
+			0,
+			'Subprocess exited successfully but returned no result payload',
 		)
-		raise SubprocessKilledError(p.pid or 0, target_name, p.exitcode or -1)
 
 	return result['value']