From c17f5ae1ff2ae11aa2ca9c70ece30c4069f23ee3 Mon Sep 17 00:00:00 2001
From: Jonas Dittrich <58814480+Kakadus@users.noreply.github.com>
Date: Wed, 10 Apr 2024 11:17:33 +0200
Subject: [PATCH 1/3] fix: race condition when resetting environment

---
 .../gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py  | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py b/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py
index cbbbe39..9ef0605 100644
--- a/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py
+++ b/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py
@@ -306,6 +306,7 @@ def reset(self, seed=None, options=None):
         if not self.gameOver:
             self.rx_env_state()
             self.send_close_command()
+            self.exp.proc.wait(2)
 
         self.msgInterface = None
         self.newStateRx = False
@@ -332,6 +333,11 @@ def get_random_action(self):
         return act
 
     def close(self):
+        if not self.gameOver:
+            self.rx_env_state()
+            self.send_close_command()
+            self.exp.proc.wait(2)
+
         # environment is not needed anymore, so kill subprocess in a straightforward way
         self.exp.kill()
         # destroy the message interface and its shared memory segment

From 3077e3d65120f1977ffff17cdb8004b24044060c Mon Sep 17 00:00:00 2001
From: Jonas Dittrich <58814480+Kakadus@users.noreply.github.com>
Date: Wed, 10 Apr 2024 10:58:49 +0200
Subject: [PATCH 2/3] fix: add more timeout suppressions

---
 .../py/ns3ai_gym_env/envs/ns3_environment.py           | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py b/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py
index 9ef0605..4e312bb 100644
--- a/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py
+++ b/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py
@@ -1,8 +1,11 @@
-import numpy as np
+from contextlib import suppress
+from subprocess import TimeoutExpired
+
 import gymnasium as gym
-from gymnasium import spaces
 import messages_pb2 as pb
 import ns3ai_gym_msg_py as py_binding
+import numpy as np
+from gymnasium import spaces
 from ns3ai_utils import Experiment
 
 
@@ -336,7 +339,8 @@ def close(self):
         if not self.gameOver:
             self.rx_env_state()
             self.send_close_command()
-            self.exp.proc.wait(2)
+            with suppress(TimeoutExpired):
+                self.exp.proc.wait(2)
 
         # environment is not needed anymore, so kill subprocess in a straightforward way
         self.exp.kill()

From 2efd59274867db0f480ec0abed0a092452634ad5 Mon Sep 17 00:00:00 2001
From: Jonas Dittrich <58814480+Kakadus@users.noreply.github.com>
Date: Tue, 18 Jun 2024 14:17:42 +0200
Subject: [PATCH 3/3] restore what was lost

---
 model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py b/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py
index 4e312bb..784420d 100644
--- a/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py
+++ b/model/gym-interface/py/ns3ai_gym_env/envs/ns3_environment.py
@@ -309,7 +309,8 @@ def reset(self, seed=None, options=None):
         if not self.gameOver:
             self.rx_env_state()
             self.send_close_command()
-            self.exp.proc.wait(2)
+            with suppress(TimeoutExpired):
+                self.exp.proc.wait(2)
 
         self.msgInterface = None
         self.newStateRx = False