Fix timeout on async locks

Author: enriquetomasmb

nebula/core/aggregation/aggregator.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from abc import ABC, abstractmethod
 from functools import partial
@@ -208,6 +209,10 @@ async def get_aggregation(self):
             await self._aggregation_done_lock.acquire_async(timeout=timeout)
         except TimeoutError:
             logging.exception("🔄  get_aggregation | Timeout reached for aggregation")
+        except asyncio.CancelledError:
+            logging.exception("🔄  get_aggregation | Lock acquisition was cancelled")
+        except Exception as e:
+            logging.exception(f"🔄  get_aggregation | Error acquiring lock: {e}")
         finally:
             await self._aggregation_done_lock.release_async()
 

nebula/core/utils/locker.py

@@ -70,7 +70,13 @@ async def acquire_async(self, *args, **kwargs):
                 )
             else:
                 logging.debug(f"🔒  Acquiring async lock [{self._name}] from {filename}:{lineno}")
-        await self._lock.acquire(*args, **kwargs)
+        if "timeout" in kwargs:
+            try:
+                await asyncio.wait_for(self._lock.acquire(), timeout=kwargs["timeout"])
+            except Exception as e:
+                raise e
+        else:
+            await self._lock.acquire()
 
     async def release_async(self, *args, **kwargs):
         caller = inspect.stack()[1]

nebula/frontend/app.py

@@ -1,7 +1,6 @@
 import argparse
 import asyncio
 import datetime
-from dotenv import load_dotenv
 import io
 import json
 import logging
@@ -12,6 +11,7 @@
 from urllib.parse import urlencode
 
 import requests
+from dotenv import load_dotenv
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))