lib: count total memory allocation per MTYPE

If malloc_usable_size() or malloc_size() are available, we can count
total usage of a particular MTYPE.  (Without the functions, we don't
know how much to subtract on free.)

Signed-off-by: David Lamparter <equinox@diac24.net>

3 files changed, 42 insertions, 7 deletions

diff --git a/lib/memory.c b/lib/memory.c
index e279b17d1..87cba69fc 100644
--- a/lib/memory.c
+++ b/lib/memory.c
@@ -17,6 +17,12 @@
 #include <zebra.h>
 
 #include <stdlib.h>
+#ifdef HAVE_MALLOC_H
+#include <malloc.h>
+#endif
+#ifdef HAVE_MALLOC_MALLOC_H
+#include <malloc/malloc.h>
+#endif
 
 #include "memory.h"
 #include "log.h"
@@ -28,7 +34,7 @@ DEFINE_MGROUP(LIB, "libfrr")
 DEFINE_MTYPE(LIB, TMP, "Temporary memory")
 DEFINE_MTYPE(LIB, PREFIX_FLOWSPEC, "Prefix Flowspec")
 
-static inline void mt_count_alloc(struct memtype *mt, size_t size)
+static inline void mt_count_alloc(struct memtype *mt, size_t size, void *ptr)
 {
 	size_t oldsize;
 
@@ -41,12 +47,24 @@ static inline void mt_count_alloc(struct memtype *mt, size_t size)
 	if (oldsize != 0 && oldsize != size && oldsize != SIZE_VAR)
 		atomic_store_explicit(&mt->size, SIZE_VAR,
 				      memory_order_relaxed);
+
+#ifdef HAVE_MALLOC_USABLE_SIZE
+	size_t mallocsz = malloc_usable_size(ptr);
+
+	atomic_fetch_add_explicit(&mt->total, mallocsz, memory_order_relaxed);
+#endif
 }
 
-static inline void mt_count_free(struct memtype *mt)
+static inline void mt_count_free(struct memtype *mt, void *ptr)
 {
 	assert(mt->n_alloc);
 	atomic_fetch_sub_explicit(&mt->n_alloc, 1, memory_order_relaxed);
+
+#ifdef HAVE_MALLOC_USABLE_SIZE
+	size_t mallocsz = malloc_usable_size(ptr);
+
+	atomic_fetch_sub_explicit(&mt->total, mallocsz, memory_order_relaxed);
+#endif
 }
 
 static inline void *mt_checkalloc(struct memtype *mt, void *ptr, size_t size)
@@ -58,7 +76,7 @@ static inline void *mt_checkalloc(struct memtype *mt, void *ptr, size_t size)
 		}
 		return NULL;
 	}
-	mt_count_alloc(mt, size);
+	mt_count_alloc(mt, size, ptr);
 	return ptr;
 }
 
@@ -75,7 +93,7 @@ void *qcalloc(struct memtype *mt, size_t size)
 void *qrealloc(struct memtype *mt, void *ptr, size_t size)
 {
 	if (ptr)
-		mt_count_free(mt);
+		mt_count_free(mt, ptr);
 	return mt_checkalloc(mt, ptr ? realloc(ptr, size) : malloc(size), size);
 }
 
@@ -87,7 +105,7 @@ void *qstrdup(struct memtype *mt, const char *str)
 void qfree(struct memtype *mt, void *ptr)
 {
 	if (ptr)
-		mt_count_free(mt);
+		mt_count_free(mt, ptr);
 	free(ptr);
 }
 
diff --git a/lib/memory.h b/lib/memory.h
index 1fbbbe423..c39f34e3a 100644
--- a/lib/memory.h
+++ b/lib/memory.h
@@ -24,12 +24,20 @@
 
 #define array_size(ar) (sizeof(ar) / sizeof(ar[0]))
 
+#if defined(HAVE_MALLOC_SIZE) && !defined(HAVE_MALLOC_USABLE_SIZE)
+#define malloc_usable_size(x) malloc_size(x)
+#define HAVE_MALLOC_USABLE_SIZE
+#endif
+
 #define SIZE_VAR ~0UL
 struct memtype {
 	struct memtype *next, **ref;
 	const char *name;
 	_Atomic size_t n_alloc;
 	_Atomic size_t size;
+#ifdef HAVE_MALLOC_USABLE_SIZE
+	_Atomic size_t total;
+#endif
 };
 
 struct memgroup {
diff --git a/lib/memory_vty.c b/lib/memory_vty.c
index 388273128..9ee2e52ec 100644
--- a/lib/memory_vty.c
+++ b/lib/memory_vty.c
@@ -79,12 +79,21 @@ static int qmem_walker(void *arg, struct memgroup *mg, struct memtype *mt)
 		if (mt->n_alloc != 0) {
 			char size[32];
 			snprintf(size, sizeof(size), "%6zu", mt->size);
-			vty_out(vty, "%-30s: %10zu  %s\n", mt->name,
+
+#ifdef HAVE_MALLOC_USABLE_SIZE
+#define TSTR " %9zu"
+#define TARG , mt->total
+#else
+#define TSTR ""
+#define TARG
+#endif
+			vty_out(vty, "%-30s: %10zu  %-16s"TSTR"\n", mt->name,
 				mt->n_alloc,
 				mt->size == 0 ? ""
 					      : mt->size == SIZE_VAR
 							? "(variably sized)"
-							: size);
+							: size
+				TARG);
 		}
 	}
 	return 0;