improve error management in async functions

Author: enriquetomasmb

nebula/core/engine.py

@@ -873,6 +873,11 @@ async def shutdown(self):
             # Wait for tasks to complete naturally with shorter timeout
             try:
                 await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=3)
+            except asyncio.CancelledError:
+                logging.warning(
+                    "Timeout reached during task cleanup (CancelledError); proceeding with shutdown anyway."
+                )
+                # Do not re-raise, just continue
             except TimeoutError:
                 logging.warning("Some tasks did not complete in time, forcing cancellation...")
                 for task in tasks:
@@ -881,9 +886,13 @@ async def shutdown(self):
                 # Wait a bit more for cancellations to take effect
                 try:
                     await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=2)
+                except asyncio.CancelledError:
+                    logging.warning(
+                        "Timeout reached during forced cancellation (CancelledError); proceeding with shutdown anyway."
+                    )
+                    # Do not re-raise, just continue
                 except TimeoutError:
                     logging.warning("Some tasks still not responding to cancellation")
-
                     # Final aggressive cleanup - cancel all remaining tasks
                     remaining_tasks = [
                         t for t in asyncio.all_tasks() if t is not asyncio.current_task() and not t.done()
@@ -894,8 +903,17 @@ async def shutdown(self):
                             task.cancel()
                         try:
                             await asyncio.wait_for(asyncio.gather(*remaining_tasks, return_exceptions=True), timeout=1)
+                        except asyncio.CancelledError:
+                            logging.warning(
+                                "Timeout reached during final forced cancellation (CancelledError); proceeding with shutdown anyway."
+                            )
+                            # Do not re-raise, just continue
                         except TimeoutError:
                             logging.exception("Some tasks still not responding to forced cancellation")
+            # Proceed anyway after all cancellation attempts
+            logging.warning("Proceeding with shutdown even if some tasks are still pending/cancelled.")
+        else:
+            logging.info("No remaining tasks to clean up.")
 
         logging.info("✅ Engine shutdown complete")
 

nebula/core/network/connection.py

@@ -169,19 +169,23 @@ async def _monitor_inactivity(self):
         periodically checking if the last activity exceeds the inactivity threshold.
         If inactive, marks the connection as inactive and logs a warning.
         """
-        while await self.is_running():
-            if self.direct:
-                break
-            await asyncio.sleep(self.INACTIVITY_DAEMON_SLEEP_TIME)
-            async with self._activity_lock:
-                time_since_last = time.time() - self._last_activity
-                if time_since_last > self.INACTIVITY_TIMER:
-                    if not self._inactivity:
-                        self._inactivity = True
-                        logging.info(f"[{self}] Connection marked as inactive.")
-                else:
-                    if self._inactivity:
-                        self._inactivity = False
+        try:
+            while await self.is_running():
+                if self.direct:
+                    break
+                await asyncio.sleep(self.INACTIVITY_DAEMON_SLEEP_TIME)
+                async with self._activity_lock:
+                    time_since_last = time.time() - self._last_activity
+                    if time_since_last > self.INACTIVITY_TIMER:
+                        if not self._inactivity:
+                            self._inactivity = True
+                            logging.info(f"[{self}] Connection marked as inactive.")
+                    else:
+                        if self._inactivity:
+                            self._inactivity = False
+        except asyncio.CancelledError:
+            logging.info("_monitor_inactivity cancelled during shutdown.")
+            return
 
     def get_federated_round(self):
         return self.federated_round
@@ -482,21 +486,19 @@ async def handle_incoming_message(self) -> None:
         try:
             while await self.is_running():
                 if self.pending_messages_queue.full():
-                    await asyncio.sleep(0.1)  # Wait a bit if the queue is full to create backpressure
+                    await asyncio.sleep(0.1)
                     continue
                 header = await self._read_exactly(self.HEADER_SIZE)
                 message_id, chunk_index, is_last_chunk = self._parse_header(header)
-
                 chunk_data = await self._read_chunk(reusable_buffer)
                 await self._update_activity()
                 self._store_chunk(message_id, chunk_index, chunk_data, is_last_chunk)
-                # logging.debug(f"Received chunk {chunk_index} of message {message_id.hex()} | size: {len(chunk_data)} bytes")
-                # Active connection without fails
                 self.incompleted_reconnections = 0
                 if is_last_chunk:
                     await self._process_complete_message(message_id)
-        except asyncio.CancelledError as e:
-            logging.exception(f"Message handling cancelled: {e}")
+        except asyncio.CancelledError:
+            logging.info("handle_incoming_message cancelled during shutdown.")
+            return
         except ConnectionError as e:
             logging.exception(f"Connection closed while reading: {e}")
         except Exception as e:
@@ -678,27 +680,23 @@ def _decompress(self, data: bytes, compression: str) -> bytes | None:
     async def process_message_queue(self) -> None:
         """
         Continuously processes messages from the pending queue.
-
-        Behavior:
-            - Retrieves messages from the queue one by one.
-            - Delegates the message to the appropriate handler based on its type.
-            - Ensures the queue is marked as processed.
-
-        Notes:
-            Runs indefinitely unless externally cancelled or stopped.
         """
-        while await self.is_running():
-            try:
-                if self.pending_messages_queue is None:
-                    logging.error("Pending messages queue is not initialized")
-                    return
-                data_type_prefix, message = await self.pending_messages_queue.get()
-                await self._handle_message(data_type_prefix, message)
-                self.pending_messages_queue.task_done()
-            except Exception as e:
-                logging.exception(f"Error processing message queue: {e}")
-            finally:
-                await asyncio.sleep(0)
+        try:
+            while await self.is_running():
+                try:
+                    if self.pending_messages_queue is None:
+                        logging.error("Pending messages queue is not initialized")
+                        return
+                    data_type_prefix, message = await self.pending_messages_queue.get()
+                    await self._handle_message(data_type_prefix, message)
+                    self.pending_messages_queue.task_done()
+                except Exception as e:
+                    logging.exception(f"Error processing message queue: {e}")
+                finally:
+                    await asyncio.sleep(0)
+        except asyncio.CancelledError:
+            logging.info("process_message_queue cancelled during shutdown.")
+            return
 
     async def _handle_message(self, data_type_prefix: bytes, message: bytes) -> None:
         """

nebula/core/network/forwarder.py

@@ -75,14 +75,17 @@ async def run_forwarder(self):
         if self.config.participant["scenario_args"]["federation"] == "CFL":
             logging.info("🔁  Federation is CFL. Forwarder is disabled...")
             return
-        while await self.is_running():
-            # logging.debug(f"🔁  Pending messages: {self.pending_messages.qsize()}")
-            start_time = time.time()
-            await self.pending_messages_lock.acquire_async()
-            await self.process_pending_messages(messages_left=self.number_forwarded_messages)
-            await self.pending_messages_lock.release_async()
-            sleep_time = max(0, self.interval - (time.time() - start_time))
-            await asyncio.sleep(sleep_time)
+        try:
+            while await self.is_running():
+                start_time = time.time()
+                await self.pending_messages_lock.acquire_async()
+                await self.process_pending_messages(messages_left=self.number_forwarded_messages)
+                await self.pending_messages_lock.release_async()
+                sleep_time = max(0, self.interval - (time.time() - start_time))
+                await asyncio.sleep(sleep_time)
+        except asyncio.CancelledError:
+            logging.info("run_forwarder cancelled during shutdown.")
+            return
 
     async def stop(self):
         self._running.clear()

nebula/core/node.py

@@ -235,6 +235,13 @@ def randomize_value(value, variability):
     if node.cm is not None:
         await node.cm.network_wait()
 
+    # Ensure shutdown is always called and awaited before main() returns
+    if hasattr(node, "shutdown") and callable(node.shutdown):
+        logging.info("Calling node.shutdown() for final cleanup and Docker removal...")
+        await node.shutdown()
+    else:
+        logging.warning("Node does not have a shutdown() method; skipping explicit shutdown.")
+
 
 if __name__ == "__main__":
     config_path = str(sys.argv[1])