From 30c8d3f1fb0de3e4fe8978fa585c2e166ca88a7b Mon Sep 17 00:00:00 2001 From: Alexis Asseman Date: Wed, 29 Mar 2023 14:17:53 -0700 Subject: [PATCH] fix: add program abort on k8s_service_watcher fail Signed-off-by: Alexis Asseman --- autoagora/k8s_service_watcher.py | 3 +++ autoagora/misc.py | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+) create mode 100644 autoagora/misc.py diff --git a/autoagora/k8s_service_watcher.py b/autoagora/k8s_service_watcher.py index c71df1b..60631c6 100644 --- a/autoagora/k8s_service_watcher.py +++ b/autoagora/k8s_service_watcher.py @@ -8,6 +8,8 @@ from kubernetes.client.api_client import ApiClient from kubernetes.client.rest import ApiException +from autoagora.misc import async_exit_on_exception + class K8SServiceEndpointsWatcher: def __init__(self, service_name: str) -> None: @@ -47,6 +49,7 @@ def __init__(self, service_name: str) -> None: # Starts the async _loop immediately self._future = aio.ensure_future(self._watch_loop()) + @async_exit_on_exception() async def _watch_loop(self) -> None: """Restarts the k8s watch on expiration.""" while True: diff --git a/autoagora/misc.py b/autoagora/misc.py new file mode 100644 index 0000000..8998618 --- /dev/null +++ b/autoagora/misc.py @@ -0,0 +1,26 @@ +import functools +import logging + + +def async_exit_on_exception(exit_code: int = -1): + """Returns decorator that logs any exception and exits the program immediately. + + The goal of this function is to easily trigger an immediate program abort from any + asynchronous function. + + Args: + exit_code (int, optional): Self explanatory. Defaults to -1. + """ + + def decorator(func): + @functools.wraps(func) + async def wrapper(*args, **kwargs): + try: + return await func(*args, **kwargs) + except: + logging.exception("exit_on_exception triggered") + exit(exit_code) + + return wrapper + + return decorator