From 309449d9d7c6936e3cbefe16fc2fc0ad06ee366f Mon Sep 17 00:00:00 2001
From: Carl Smith <carl.smith@alliedtelesis.co.nz>
Date: Wed, 18 Sep 2024 07:54:14 +1200
Subject: [PATCH] Improve memuse by providing total in use and in the malloc
 arena

---
 apteryx.c  | 23 +++++++++++++++--------
 apteryx.h  |  2 ++
 apteryxc.c | 10 ++++++++--
 apteryxd.c | 16 +++++++++++++++-
 4 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/apteryx.c b/apteryx.c
index 0fb9163..b76c87a 100644
--- a/apteryx.c
+++ b/apteryx.c
@@ -2262,15 +2262,22 @@ apteryx_memuse (const char *path)
     DEBUG ("MEMUSE: %s\n", path ?: "NULL");
 
     /* Check path */
-    path = validate_path (path, &url);
-    /* if path is empty, or path ends in '/' but is not the root db path (ie "/") */
-    if (!path ||
-        ((path[strlen(path)-1] == '/') && strlen(path) > 1))
+    if (path[0] == '.')
     {
-        ERROR ("MEMUSE: invalid path (%s)!\n", path ?: "NULL");
-        free (url);
-        assert (!apteryx_debug || path);
-        return 0;
+        url = strdup(default_url);
+    }
+    else
+    {
+        path = validate_path (path, &url);
+        /* if path is empty, or path ends in '/' but is not the root db path (ie "/") */
+        if (!path ||
+            ((path[strlen(path)-1] == '/') && strlen(path) > 1))
+        {
+            ERROR ("MEMUSE: invalid path (%s)!\n", path ?: "NULL");
+            free (url);
+            assert (!apteryx_debug || path);
+            return 0;
+        }
     }
 
     /* IPC */
diff --git a/apteryx.h b/apteryx.h
index 832aac1..800b909 100644
--- a/apteryx.h
+++ b/apteryx.h
@@ -229,6 +229,8 @@ uint64_t apteryx_timestamp (const char *path);
 /**
  * Get the memory usage in bytes of a given path
  * @param path path to get the memory usage for
+ *             '.'  = total memory in use (mi.uordblks + mi.hblkhd)
+ *             '..' = total memory allocated (mi.arena + mi.hblkhd)
  * @return 0 if the path doesn't exist, memory usage in bytes otherwise
  */
 uint64_t apteryx_memuse (const char *path);
diff --git a/apteryxc.c b/apteryxc.c
index c9dcd7b..6688e70 100644
--- a/apteryxc.c
+++ b/apteryxc.c
@@ -463,8 +463,7 @@ main (int argc, char **argv)
     {
         if (!path || param)
         {
-            usage ();
-            return 0;
+            path = "/";
         }
         apteryx_init (apteryx_debug);
         if (path[strlen(path) - 1] != '/')
@@ -479,6 +478,13 @@ main (int argc, char **argv)
                 printf ("%10"PRIu64" %s\n", size, (char *) _iter->data);
         }
         g_list_free_full (paths, free);
+        if (g_strcmp0 (path, "/") == 0)
+        {
+            uint64_t total = apteryx_memuse ("..");
+            uint64_t used = apteryx_memuse (".");
+            uint64_t db = apteryx_memuse ("/");
+            printf ("%"PRIu64"/%"PRIu64" total bytes used/arena (%"PRIu64" in database)\n", used, total, db);
+        }
         apteryx_shutdown ();
         g_free (path);
         break;
diff --git a/apteryxd.c b/apteryxd.c
index 5dd6131..475567a 100644
--- a/apteryxd.c
+++ b/apteryxd.c
@@ -26,6 +26,7 @@
 #include <sys/poll.h>
 #include <sys/wait.h>
 #include <arpa/inet.h>
+#include <malloc.h>
 #include "apteryx.h"
 #include "internal.h"
 
@@ -2436,7 +2437,20 @@ handle_memuse (rpc_message msg)
     INC_COUNTER (counters.memuse);
 
     /* Lookup value */
-    value = db_memuse (path);
+    if (path[0] == '.' && path[1] == '\0')
+    {
+        /* Total memory in use */
+        struct mallinfo2 mi = mallinfo2 ();
+        value = (unsigned int) (mi.uordblks) + (unsigned int) (mi.hblkhd);
+    }
+    else if (path[0] == '.' && path[1] == '.' && path[2] == '\0')
+    {
+        /* Total memory allocated */
+        struct mallinfo2 mi = mallinfo2 ();
+        value = (unsigned int) (mi.arena) + (unsigned int) (mi.hblkhd);
+    }
+    else
+        value = db_memuse (path);
 
     /* Send result */
     DEBUG ("     = %"PRIu64"\n", value);