From c93324da552f0af11cc5559719844efd6b283bc1 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Wed, 7 Aug 2024 07:59:56 +0900
Subject: [PATCH] try-catch

---
 mixnet/framework/usim.py      |  2 +-
 mixnet/queuesim/queuesim.py   | 22 +++++++++++++++++-----
 mixnet/queuesim/simulation.py |  5 -----
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/mixnet/framework/usim.py b/mixnet/framework/usim.py
index 749cbb0..c576297 100644
--- a/mixnet/framework/usim.py
+++ b/mixnet/framework/usim.py
@@ -1,7 +1,7 @@
 from typing import Any, Awaitable, Coroutine, TypeVar
 
 import usim
-from usim._primitives.task import Task, TaskCancelled
+from usim._primitives.task import Task
 
 from framework import framework
 
diff --git a/mixnet/queuesim/queuesim.py b/mixnet/queuesim/queuesim.py
index e1d3dce..b28ca63 100644
--- a/mixnet/queuesim/queuesim.py
+++ b/mixnet/queuesim/queuesim.py
@@ -1,8 +1,8 @@
 import concurrent.futures
-import itertools
 import os
 import random
 import time
+import traceback
 from collections import defaultdict
 from copy import deepcopy
 from datetime import datetime, timedelta
@@ -122,8 +122,15 @@ def run_session(
         paramsets_done: set[int] = set()
         for future in concurrent.futures.as_completed(future_map):
             paramset_id, paramset, iter_idx = future_map[future]
+            try:
+                dissemination_times = future.result()
+            except BaseException as e:
+                print(f"Error occurred in ParamSet-{paramset_id}, Iter-{iter_idx}: {e}")
+                traceback.print_exc()
+                raise
+
             paramset_results[paramset_id][0].add(iter_idx)
-            paramset_results[paramset_id][1].extend(future.result())
+            paramset_results[paramset_id][1].extend(dissemination_times)
             # If all iterations of the paramset are done, process the results
             if len(paramset_results[paramset_id][0]) == paramset.num_iterations:
                 paramsets_done.add(paramset_id)
@@ -204,9 +211,14 @@ def __run_iteration(cfg: Config) -> list[float]:
     Run a single iteration of a certain parameter set.
     The iteration uses the independent uSim instance.
     """
-    sim = Simulation(cfg)
-    usim.run(sim.run())
-    return sim.dissemination_times
+    try:
+        sim = Simulation(cfg)
+        usim.run(sim.run())
+        return sim.dissemination_times
+    except BaseException as e:
+        print(f"Error in running the iteration: {e}")
+        traceback.print_exc()
+        raise
 
 
 def __process_paramset_result(
diff --git a/mixnet/queuesim/simulation.py b/mixnet/queuesim/simulation.py
index d40fd1d..a2cc3a6 100644
--- a/mixnet/queuesim/simulation.py
+++ b/mixnet/queuesim/simulation.py
@@ -1,7 +1,5 @@
-import random
 import struct
 from dataclasses import dataclass
-from random import Random
 from typing import Counter, Self
 
 import usim
@@ -9,11 +7,8 @@ import usim
 from framework.framework import Queue
 from framework.usim import Framework
 from protocol.connection import LocalSimplexConnection, SimplexConnection
-from protocol.nomssip import NomssipConfig
-from protocol.temporalmix import TemporalMixConfig, TemporalMixType
 from queuesim.config import Config
 from queuesim.node import Node
-from sim.config import LatencyConfig
 from sim.connection import RemoteSimplexConnection
 from sim.topology import build_full_random_topology