fix: make sure leftover traces are flushed before exiting

Author: Essoz

traincheck/config/config.py

@@ -6,7 +6,7 @@
 
 # trace dumper configs:
 BUFFER_SIZE = 1000  # number of events to buffer before dumping
-FLUSH_INTERVAL = 5  # seconds
+FLUSH_INTERVAL = 0.5  # seconds
 
 # runner configs
 RUNNER_DEFAULT_ENV = {

traincheck/instrumentor/dumper.py

@@ -43,7 +43,6 @@
 
 logger = logging.getLogger(__name__)
 
-
 def serialize(obj_dict: dict[str, object | str]) -> str:
     try:
         return orjson.dumps(obj_dict).decode("utf-8")
@@ -55,20 +54,23 @@ def serialize(obj_dict: dict[str, object | str]) -> str:
 def monitor_main_thread(main_thread, stop_event):
     main_thread.join()  # Wait for the main thread to finish
     print("Main thread has finished or encountered an exception")
+    print("Flushing all buffers to the trace log file")
     stop_event.set()  # Signal the logging threads to stop
 
 
 def trace_dumper(task_queue: Queue, trace_file_name: str, stop_event: threading.Event):
     with open(trace_file_name, "w") as f:
         while True:
             try:
-                trace = task_queue.get(timeout=0.5)
+                trace = task_queue.get(timeout=FLUSH_INTERVAL * 2) # wait for 2x the flush interval, this is an arbitrary number, as long as it is larger than the flush interval, it should be fine.
             except Empty:
                 if stop_event.is_set():
+                    print("Trace dumper thread has stopped.")
                     break
                 continue
             f.write(f"{trace}\n")
             task_queue.task_done()
+    print("Trace dumper thread has finished normally...")
 
 
 def get_trace_API_dumper_queue():
@@ -178,6 +180,7 @@ def _flush_periodically(self):
         while not stop_event.is_set():
             time.sleep(self.flush_interval)
             self._flush()
+        self._flush()  # flush any remaining traces before exiting
 
 
 # THREAD_DATA.api_trace_buffer = TraceBuffer(get_trace_API_dumper_queue)