From 073f656da12ba3ee573cd50e9d3318810f97b440 Mon Sep 17 00:00:00 2001
From: Christian Deacon <christianmdeacon@gmail.com>
Date: Wed, 26 Feb 2025 12:31:16 -0500
Subject: [PATCH] Improve stats output.

---
 src/loader/prog.c          | 13 +++++++++++++
 src/loader/utils/logging.c | 10 +++++++++-
 src/loader/utils/logging.h |  2 ++
 src/loader/utils/stats.c   |  7 +++++--
 4 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/loader/prog.c b/src/loader/prog.c
index 80e2a5f..c5628fc 100644
--- a/src/loader/prog.c
+++ b/src/loader/prog.c
@@ -19,6 +19,7 @@
 #include <loader/utils/helpers.h>
 
 int cont = 1;
+int doing_stats = 0;
 
 int main(int argc, char *argv[])
 {
@@ -194,6 +195,12 @@ int main(int argc, char *argv[])
 
     struct stat conf_stat;
 
+    // Check if we're doing stats.
+    if (!cfg.nostats)
+    {
+        doing_stats = 1;
+    }
+
     while (cont)
     {
         // Get current time.
@@ -225,6 +232,12 @@ int main(int argc, char *argv[])
 
                 // Update timer
                 last_config_check = time(NULL);
+
+                // Make sure we set doing stats if needed.
+                if (!cfg.nostats && !doing_stats)
+                {
+                    doing_stats = 1;
+                }
             }
 
             // Update last updated variable.
diff --git a/src/loader/utils/logging.c b/src/loader/utils/logging.c
index 95ed381..f27402c 100644
--- a/src/loader/utils/logging.c
+++ b/src/loader/utils/logging.c
@@ -42,7 +42,15 @@ static void LogMsgRaw(int req_lvl, int cur_lvl, int error, const char* log_path,
     char full_msg[len + 6 + 1];
     snprintf(full_msg, sizeof(full_msg), "[%d] %s", req_lvl, f_msg);
 
-    fprintf(pipe, "%s\n", full_msg);
+    // If we're calculating stats, we need to prepend a new line.
+    if (doing_stats)
+    {
+        fprintf(pipe, "\n%s\n", full_msg);
+    }
+    else
+    {
+        fprintf(pipe, "%s\n", full_msg);
+    }
 
     if (log_path != NULL)
     {
diff --git a/src/loader/utils/logging.h b/src/loader/utils/logging.h
index d3dd611..beabf54 100644
--- a/src/loader/utils/logging.h
+++ b/src/loader/utils/logging.h
@@ -14,5 +14,7 @@
 
 #define RB_TIMEOUT 100
 
+extern int doing_stats;
+
 void LogMsg(config__t* cfg, int req_lvl, int error, const char* msg, ...);
 int HandleRbEvent(void* ctx, void* data, size_t sz);
\ No newline at end of file
diff --git a/src/loader/utils/stats.c b/src/loader/utils/stats.c
index e19dc8c..f317a7e 100644
--- a/src/loader/utils/stats.c
+++ b/src/loader/utils/stats.c
@@ -40,9 +40,12 @@ int CalculateStats(int stats_map, int cpus)
         dropped += stats[i].dropped;
         passed += stats[i].passed;
     }
+    
+    printf("\r\033[1;32mAllowed:\033[0m %llu  |  ", allowed);
+    printf("\033[1;31mDropped:\033[0m %llu  |  ", dropped);
+    printf("\033[1;34mPassed:\033[0m %llu", passed);
 
-    fflush(stdout);
-    fprintf(stdout, "\rAllowed: %llu | Dropped: %llu | Passed: %llu", allowed, dropped, passed);
+    fflush(stdout);    
 
     return EXIT_SUCCESS;
 }
\ No newline at end of file