From 95aeda4cdda2431c20ed9fa3facb241b142ae773 Mon Sep 17 00:00:00 2001
From: daanx <daan@effp.org>
Date: Sat, 21 Dec 2024 10:53:34 -0800
Subject: [PATCH] merge subproc stats on delete

---
 include/mimalloc/internal.h |  1 +
 src/init.c                  |  4 ++++
 src/stats.c                 | 23 +++++++++++------------
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/include/mimalloc/internal.h b/include/mimalloc/internal.h
index 7774b378..e316de94 100644
--- a/include/mimalloc/internal.h
+++ b/include/mimalloc/internal.h
@@ -203,6 +203,7 @@ void        _mi_heap_page_reclaim(mi_heap_t* heap, mi_page_t* page);
 
 // "stats.c"
 void        _mi_stats_done(mi_stats_t* stats);
+void        _mi_stats_merge_from(mi_stats_t* to, mi_stats_t* from);
 mi_msecs_t  _mi_clock_now(void);
 mi_msecs_t  _mi_clock_end(mi_msecs_t start);
 mi_msecs_t  _mi_clock_start(void);
diff --git a/src/init.c b/src/init.c
index 5159941a..3af4f4ef 100644
--- a/src/init.c
+++ b/src/init.c
@@ -382,6 +382,10 @@ void mi_subproc_delete(mi_subproc_id_t subproc_id) {
     mi_lock_release(&subproc->os_pages_lock);
   }
   if (!safe_to_delete) return;
+
+  // merge stats back into the main subproc?
+  _mi_stats_merge_from(&_mi_subproc_main()->stats, &subproc->stats);
+
   // safe to release
   // todo: should we refcount subprocesses?
   mi_lock_done(&subproc->os_pages_lock);
diff --git a/src/stats.c b/src/stats.c
index 2a395ed5..102373ec 100644
--- a/src/stats.c
+++ b/src/stats.c
@@ -411,14 +411,6 @@ static mi_stats_t* mi_get_tld_stats(void) {
   return &_mi_tld()->stats;
 }
 
-static void mi_stats_merge_from(mi_stats_t* stats) {
-  mi_subproc_t* subproc = _mi_subproc();
-  if (stats != &subproc->stats) {
-    mi_stats_add(&subproc->stats, stats);
-    _mi_memzero(stats, sizeof(mi_stats_t));
-  }
-}
-
 void mi_stats_reset(void) mi_attr_noexcept {
   mi_stats_t* stats = mi_get_tld_stats();
   mi_subproc_t* subproc = _mi_subproc();
@@ -427,16 +419,23 @@ void mi_stats_reset(void) mi_attr_noexcept {
   if (mi_process_start == 0) { mi_process_start = _mi_clock_start(); };
 }
 
-void mi_stats_merge(void) mi_attr_noexcept {
-  mi_stats_merge_from( mi_get_tld_stats() );
+void _mi_stats_merge_from(mi_stats_t* to, mi_stats_t* from) {
+  if (to != from) {
+    mi_stats_add(to, from);
+    _mi_memzero(from, sizeof(mi_stats_t));
+  }
 }
 
 void _mi_stats_done(mi_stats_t* stats) {  // called from `mi_thread_done`
-  mi_stats_merge_from(stats);
+  _mi_stats_merge_from(&_mi_subproc()->stats, stats);
+}
+
+void mi_stats_merge(void) mi_attr_noexcept {
+  _mi_stats_done( mi_get_tld_stats() );
 }
 
 void mi_stats_print_out(mi_output_fun* out, void* arg) mi_attr_noexcept {
-  mi_stats_merge_from(mi_get_tld_stats());
+  mi_stats_merge();
   _mi_stats_print(&_mi_subproc()->stats, out, arg);
 }