Packet IO dump file mode, automatically exits when all sessions are released

2 files changed, 46 insertions, 2 deletions

diff --git a/src/stellar/main.cpp b/src/stellar/main.cpp
index aba22be..f16edc3 100644
--- a/src/stellar/main.cpp
+++ b/src/stellar/main.cpp
@@ -42,6 +42,27 @@ static void signal_handler(int signo)
     }
 }
 
+static int all_session_have_freed(void)
+{
+    for (int i = 0; i < config->io_opts.nr_threads; i++)
+    {
+        struct session_manager *sess_mgr = runtime->threads[i].sess_mgr;
+        struct session_manager_stat *sess_stat = session_manager_stat(sess_mgr);
+
+        if (ATOMIC_READ(&sess_stat->curr_nr_tcp_sess_used) != 0)
+        {
+            return 0;
+        }
+
+        if (ATOMIC_READ(&sess_stat->curr_nr_udp_sess_used) != 0)
+        {
+            return 0;
+        }
+    }
+
+    return 1;
+}
+
 int main(int argc, char **argv)
 {
     timestamp_update();
@@ -116,7 +137,8 @@ int main(int argc, char **argv)
         }
         usleep(1000); // 1ms
 
-        if (packet_io_wait_exit(runtime->packet_io))
+        // Only available in dump file mode, automatically exits when all sessions have been released
+        if (packet_io_wait_exit(runtime->packet_io) && all_session_have_freed())
         {
             ATOMIC_SET(&runtime->need_exit, 1);
         }
diff --git a/test/packet_injector.cpp b/test/packet_injector.cpp
index 922fda9..d174fa0 100644
--- a/test/packet_injector.cpp
+++ b/test/packet_injector.cpp
@@ -209,6 +209,27 @@ static void signal_handler(int signo)
     }
 }
 
+static int all_session_have_freed(void)
+{
+    for (int i = 0; i < config->io_opts.nr_threads; i++)
+    {
+        struct session_manager *sess_mgr = runtime->threads[i].sess_mgr;
+        struct session_manager_stat *sess_stat = session_manager_stat(sess_mgr);
+
+        if (ATOMIC_READ(&sess_stat->curr_nr_tcp_sess_used) != 0)
+        {
+            return 0;
+        }
+
+        if (ATOMIC_READ(&sess_stat->curr_nr_udp_sess_used) != 0)
+        {
+            return 0;
+        }
+    }
+
+    return 1;
+}
+
 static void usage(char *cmd)
 {
     printf("Usage: %s [options]\n\n", cmd);
@@ -427,7 +448,8 @@ int main(int argc, char **argv)
         }
         usleep(1000); // 1ms
 
-        if (packet_io_wait_exit(runtime->packet_io))
+        // Only available in dump file mode, automatically exits when all sessions have been released
+        if (packet_io_wait_exit(runtime->packet_io) && all_session_have_freed())
         {
             ATOMIC_SET(&runtime->need_exit, 1);
         }