From f1e35909b424a75c80b6daa7f3c891cd23efa342 Mon Sep 17 00:00:00 2001
From: Tian Gao <gaogaotiantian@hotmail.com>
Date: Sat, 16 Nov 2024 10:38:05 -0800
Subject: [PATCH] Free the memory/reference when module is freed (#517)

---
 src/viztracer/modules/quicktime.c | 14 ++++++++++++--
 src/viztracer/modules/quicktime.h |  1 +
 src/viztracer/modules/snaptrace.c | 26 +++++++++++++++++---------
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/viztracer/modules/quicktime.c b/src/viztracer/modules/quicktime.c
index 66bd91a8..13df3db1 100644
--- a/src/viztracer/modules/quicktime.c
+++ b/src/viztracer/modules/quicktime.c
@@ -25,8 +25,8 @@ mach_timebase_info_data_t timebase_info;
 double ts_to_ns_factor = 1.0;
 int64_t system_base_time = 0;
 int64_t system_base_ts = 0;
-int64_t start_ts[CALIBRATE_SIZE] = {0};
-int64_t start_ns[CALIBRATE_SIZE] = {0};
+int64_t* start_ts = NULL;
+int64_t* start_ns = NULL;
 int64_t t0_ts = 0;
 int64_t t0_ns = 0;
 bool calibrated = false;
@@ -105,6 +105,12 @@ int64_t dur_ts_to_ns(int64_t dur)
     return dur * ts_to_ns_factor;
 }
 
+void quicktime_free()
+{
+    free(start_ts);
+    free(start_ns);
+}
+
 void quicktime_init()
 {
 #if _WIN32
@@ -112,6 +118,10 @@ void quicktime_init()
 #elif defined(__APPLE__)
     mach_timebase_info(&timebase_info);
 #endif
+
+    start_ts = (int64_t*)malloc(sizeof(int64_t) * CALIBRATE_SIZE);
+    start_ns = (int64_t*)malloc(sizeof(int64_t) * CALIBRATE_SIZE);
+
     int64_t diff_ns[CALIBRATE_SIZE] = {0};
 
     t0_ts = 0;
diff --git a/src/viztracer/modules/quicktime.h b/src/viztracer/modules/quicktime.h
index 28174c93..920d81f1 100644
--- a/src/viztracer/modules/quicktime.h
+++ b/src/viztracer/modules/quicktime.h
@@ -21,6 +21,7 @@ extern double ts_to_ns_factor;
 extern int64_t system_base_time;
 
 void quicktime_init();
+void quicktime_free();
 double system_ts_to_us(int64_t ts);
 int64_t system_ts_to_ns(int64_t ts);
 double dur_ts_to_us(int64_t dur);
diff --git a/src/viztracer/modules/snaptrace.c b/src/viztracer/modules/snaptrace.c
index 49fb8010..d475acad 100644
--- a/src/viztracer/modules/snaptrace.c
+++ b/src/viztracer/modules/snaptrace.c
@@ -46,6 +46,7 @@ static PyObject* snaptrace_setcurrstack(TracerObject* self, PyObject* stack_dept
 static PyObject* snaptrace_setignorestackcounter(TracerObject* self, PyObject* value);
 static void snaptrace_flush_unfinished(TracerObject* self, int flush_as_finish);
 static void snaptrace_threaddestructor(void* key);
+static void snaptrace_free(void*);
 static struct ThreadInfo* snaptrace_createthreadinfo(TracerObject* self);
 static void log_func_args(struct FunctionNode* node, PyFrameObject* frame, PyObject* log_func_repr);
 static int64_t get_ts(struct ThreadInfo*);
@@ -238,20 +239,15 @@ static PyMethodDef Tracer_methods[] = {
     {NULL, NULL, 0, NULL}
 };
 
-static PyMethodDef Snaptrace_methods[] = {
-    {NULL, NULL, 0, NULL}
-};
-
 // ================================================================
 // Python interface
 // ================================================================
 
 static struct PyModuleDef snaptracemodule = {
-    PyModuleDef_HEAD_INIT,
-    "viztracer.snaptrace",
-    NULL,
-    -1,
-    Snaptrace_methods
+    .m_base = PyModuleDef_HEAD_INIT,
+    .m_name = "viztracer.snaptrace",
+    .m_size = -1,
+    .m_free = snaptrace_free,
 };
 
 // =============================================================================
@@ -1768,6 +1764,18 @@ static PyTypeObject TracerType = {
     .tp_getset = Tracer_getsetters,
 };
 
+void snaptrace_free(void* Py_UNUSED(unused)) {
+    quicktime_free();
+    Py_CLEAR(threading_module);
+    Py_CLEAR(multiprocessing_module);
+    Py_CLEAR(asyncio_module);
+    Py_CLEAR(asyncio_tasks_module);
+    Py_CLEAR(curr_task_getters[0]);
+    Py_CLEAR(trio_lowlevel_module);
+    Py_CLEAR(curr_task_getters[1]);
+    Py_CLEAR(json_module);
+}
+
 PyMODINIT_FUNC
 PyInit_snaptrace(void) 
 {