From b8bcfc079837b4650e142f3ae5b4205945d79ce1 Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mattpark@microsoft.com>
Date: Wed, 15 May 2019 17:21:43 +0100
Subject: [PATCH] Made the statistics print atexit

Fixed some statistics and made them automatically print atexit.
---
 src/ds/helpers.h      | 12 +++++++++++-
 src/mem/allocstats.h  | 13 +++++++++----
 src/mem/largealloc.h  |  4 ++++
 src/mem/threadalloc.h | 22 +++++++++++++++++++++-
 4 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/src/ds/helpers.h b/src/ds/helpers.h
index 819d73bef..e94227db1 100644
--- a/src/ds/helpers.h
+++ b/src/ds/helpers.h
@@ -13,12 +13,20 @@ namespace snmalloc
   class Singleton
   {
   public:
-    inline static Object& get()
+    /**
+     * If argument is non-null, then it is assigned the value
+     * true, if this is the first call to get.
+     * At most one call will be first.
+     */
+    inline static Object& get(bool* first = nullptr)
     {
       static std::atomic_flag flag;
       static std::atomic<bool> initialised;
       static Object obj;
 
+      // If defined should be initially false;
+      assert(first == nullptr || *first == false);
+
       if (!initialised.load(std::memory_order_acquire))
       {
         FlagLock lock(flag);
@@ -26,6 +34,8 @@ namespace snmalloc
         {
           obj = init();
           initialised.store(true, std::memory_order_release);
+          if (first != nullptr)
+            *first = true;
         }
       }
       return obj;
diff --git a/src/mem/allocstats.h b/src/mem/allocstats.h
index 710dadb62..b047a3a0a 100644
--- a/src/mem/allocstats.h
+++ b/src/mem/allocstats.h
@@ -21,10 +21,12 @@ namespace snmalloc
     {
       size_t current = 0;
       size_t max = 0;
+      size_t used = 0;
 
       void inc()
       {
         current++;
+        used++;
         if (current > max)
           max++;
       }
@@ -48,11 +50,12 @@ namespace snmalloc
       {
         current += that.current;
         max += that.max;
+        used += that.used;
       }
 #ifdef USE_SNMALLOC_STATS
       void print(CSVStream& csv, size_t multiplier = 1)
       {
-        csv << current * multiplier << max * multiplier;
+        csv << current * multiplier << max * multiplier << used * multiplier;
       }
 #endif
     };
@@ -327,10 +330,12 @@ namespace snmalloc
             << "AllocatorID"
             << "Size group"
             << "Size"
-            << "Current bytes"
-            << "Max bytes"
+            << "Current count"
+            << "Max count"
+            << "Total Allocs"
             << "Current Slab bytes"
             << "Max Slab bytes"
+            << "Total slab allocs"
             << "Average Slab Usage"
             << "Average wasted space" << csv.endl;
 
@@ -353,7 +358,7 @@ namespace snmalloc
         csv << "BucketedStats" << dumpid << allocatorid << i
             << sizeclass_to_size(i);
 
-        sizeclass[i].print(csv, sizeclass_to_size(i), SLAB_SIZE);
+        sizeclass[i].print(csv, sizeclass_to_size(i));
       }
 
       for (uint8_t i = 0; i < LARGE_N; i++)
diff --git a/src/mem/largealloc.h b/src/mem/largealloc.h
index 61053b16a..929cbd7b0 100644
--- a/src/mem/largealloc.h
+++ b/src/mem/largealloc.h
@@ -345,11 +345,14 @@ namespace snmalloc
         p = reserved_start;
         reserved_start = pointer_offset(p, rsize);
 
+        stats.superslab_fresh();
         // All memory is zeroed since it comes from reserved space.
         memory_provider.template notify_using<NoZero>(p, size);
       }
       else
       {
+        stats.superslab_pop();
+
         if constexpr (decommit_strategy == DecommitSuperLazy)
         {
           if (static_cast<Baseslab*>(p)->get_kind() == Decommitted)
@@ -400,6 +403,7 @@ namespace snmalloc
 
     void dealloc(void* p, size_t large_class)
     {
+      stats.superslab_push();
       memory_provider.large_stack[large_class].push(static_cast<Largeslab*>(p));
       memory_provider.lazy_decommit_if_needed();
     }
diff --git a/src/mem/threadalloc.h b/src/mem/threadalloc.h
index 717a6a230..59a05e3ff 100644
--- a/src/mem/threadalloc.h
+++ b/src/mem/threadalloc.h
@@ -230,6 +230,15 @@ namespace snmalloc
       return per_thread;
     }
 
+#  ifdef USE_SNMALLOC_STATS
+    static void print_stats()
+    {
+      Stats s;
+      current_alloc_pool()->aggregate_stats(s);
+      s.print<Alloc>(std::cout);
+    }
+#  endif
+
     /**
      * Private initialiser for the per thread allocator
      */
@@ -248,9 +257,20 @@ namespace snmalloc
         // allocator.
         per_thread = current_alloc_pool()->acquire();
 
-        tls_key_t key = Singleton<tls_key_t, tls_key_create>::get();
+        bool first = false;
+        tls_key_t key = Singleton<tls_key_t, tls_key_create>::get(&first);
         // Associate the new allocator with the destructor.
         tls_set_value(key, &per_thread);
+
+#  ifdef USE_SNMALLOC_STATS
+        // Allocator is up and running now, safe to call atexit.
+        if (first)
+        {
+          atexit(print_stats);
+        }
+#  else
+        UNUSED(first);
+#  endif
       }
       return per_thread;
     }