From a23ef7d6039cbbed94bd639f24f84eb8772cfb96 Mon Sep 17 00:00:00 2001
From: Blake Harnden <32446120+bharnden@users.noreply.github.com>
Date: Sat, 5 Dec 2020 09:01:53 -0800
Subject: [PATCH] daemon: properly go through the data collect state for grpc
 session shutdown, also check and avoid repeating data collect or shutdown
 when already past those states

---
 daemon/core/api/grpc/server.py       |  1 +
 daemon/core/emulator/coreemu.py      |  1 +
 daemon/core/emulator/enumerations.py |  3 +++
 daemon/core/emulator/session.py      | 21 ++++++++++++++-------
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/daemon/core/api/grpc/server.py b/daemon/core/api/grpc/server.py
index 3159776a..fcd3210d 100644
--- a/daemon/core/api/grpc/server.py
+++ b/daemon/core/api/grpc/server.py
@@ -315,6 +315,7 @@ class CoreGrpcServer(core_pb2_grpc.CoreApiServicer):
         """
         logging.debug("stop session: %s", request)
         session = self.get_session(request.session_id, context)
+        session.data_collect()
         session.shutdown()
         return core_pb2.StopSessionResponse(result=True)
 
diff --git a/daemon/core/emulator/coreemu.py b/daemon/core/emulator/coreemu.py
index c07d8c95..885fb431 100644
--- a/daemon/core/emulator/coreemu.py
+++ b/daemon/core/emulator/coreemu.py
@@ -144,6 +144,7 @@ class CoreEmu:
         result = False
         if session:
             logging.info("shutting session down: %s", _id)
+            session.data_collect()
             session.shutdown()
             result = True
         else:
diff --git a/daemon/core/emulator/enumerations.py b/daemon/core/emulator/enumerations.py
index 5cad0aa5..83e7bffd 100644
--- a/daemon/core/emulator/enumerations.py
+++ b/daemon/core/emulator/enumerations.py
@@ -106,6 +106,9 @@ class EventTypes(Enum):
     def should_start(self) -> bool:
         return self.value > self.DEFINITION_STATE.value
 
+    def already_collected(self) -> bool:
+        return self.value >= self.DATACOLLECT_STATE.value
+
 
 class ExceptionLevels(Enum):
     """
diff --git a/daemon/core/emulator/session.py b/daemon/core/emulator/session.py
index e3363f7c..d98564ed 100644
--- a/daemon/core/emulator/session.py
+++ b/daemon/core/emulator/session.py
@@ -770,10 +770,11 @@ class Session:
         """
         Shutdown all session nodes and remove the session directory.
         """
+        if self.state == EventTypes.SHUTDOWN_STATE:
+            logging.info("session(%s) state(%s) already shutdown", self.id, self.state)
+            return
         logging.info("session(%s) state(%s) shutting down", self.id, self.state)
-        if self.state != EventTypes.SHUTDOWN_STATE:
-            self.set_state(EventTypes.DATACOLLECT_STATE, send_event=True)
-            self.set_state(EventTypes.SHUTDOWN_STATE, send_event=True)
+        self.set_state(EventTypes.SHUTDOWN_STATE, send_event=True)
         # clear out current core session
         self.clear()
         # shutdown sdt
@@ -1258,6 +1259,14 @@ class Session:
 
         :return: nothing
         """
+        if self.state.already_collected():
+            logging.info(
+                "session(%s) state(%s) already data collected", self.id, self.state
+            )
+            return
+        logging.info("session(%s) state(%s) data collection", self.id, self.state)
+        self.set_state(EventTypes.DATACOLLECT_STATE, send_event=True)
+
         # stop event loop
         self.event_loop.stop()
 
@@ -1279,10 +1288,8 @@ class Session:
         self.update_control_iface_hosts(remove=True)
 
         # remove all four possible control networks
-        self.add_remove_control_net(0, remove=True)
-        self.add_remove_control_net(1, remove=True)
-        self.add_remove_control_net(2, remove=True)
-        self.add_remove_control_net(3, remove=True)
+        for i in range(4):
+            self.add_remove_control_net(i, remove=True)
 
     def short_session_id(self) -> str:
         """