From af417271a47e88d47cc3dc25f93f56d74864c3dd Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Mon, 23 Dec 2019 19:08:05 +0000
Subject: [PATCH] [Project] Add rspamd_mempool_notify_alloc function

---
 src/libutil/mem_pool.c | 35 ++++++++++++++++++++++-------------
 src/libutil/mem_pool.h | 10 ++++++++++
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/src/libutil/mem_pool.c b/src/libutil/mem_pool.c
index 5e4d1fe1a..f2272fa7c 100644
--- a/src/libutil/mem_pool.c
+++ b/src/libutil/mem_pool.c
@@ -441,6 +441,27 @@ memory_pool_alloc_common (rspamd_mempool_t * pool, gsize size,
 						  const gchar *loc)
 RSPAMD_ATTR_ALLOC_SIZE(2) RSPAMD_ATTR_ALLOC_ALIGN(MIN_MEM_ALIGNMENT) RSPAMD_ATTR_RETURNS_NONNUL;
 
+
+void
+rspamd_mempool_notify_alloc_ (rspamd_mempool_t *pool, gsize size, const gchar *loc)
+{
+	GHashTable *debug_tbl = *(GHashTable **)(((guchar *)pool + sizeof (*pool)));
+	gpointer ptr;
+
+	if (G_UNLIKELY (pool->priv->flags & RSPAMD_MEMPOOL_DEBUG)) {
+		ptr = g_hash_table_lookup (debug_tbl, loc);
+
+		if (ptr) {
+			ptr = GSIZE_TO_POINTER (GPOINTER_TO_SIZE (ptr) + size);
+		}
+		else {
+			ptr = GSIZE_TO_POINTER (size);
+		}
+
+		g_hash_table_insert (debug_tbl, (gpointer) loc, ptr);
+	}
+}
+
 static void *
 memory_pool_alloc_common (rspamd_mempool_t * pool, gsize size,
 		enum rspamd_mempool_chain_type pool_type, const gchar *loc)
@@ -454,19 +475,7 @@ memory_pool_alloc_common (rspamd_mempool_t * pool, gsize size,
 		pool->priv->used_memory += size;
 
 		if (G_UNLIKELY (pool->priv->flags & RSPAMD_MEMPOOL_DEBUG)) {
-			GHashTable *debug_tbl = *(GHashTable **)(((guchar *)pool + sizeof (*pool)));
-			gpointer ptr;
-
-			ptr = g_hash_table_lookup (debug_tbl, loc);
-
-			if (ptr) {
-				ptr = GSIZE_TO_POINTER (GPOINTER_TO_SIZE (ptr) + size);
-			}
-			else {
-				ptr = GSIZE_TO_POINTER (size);
-			}
-
-			g_hash_table_insert (debug_tbl, (gpointer)loc, ptr);
+			rspamd_mempool_notify_alloc_ (pool, size, loc);
 		}
 
 		if (always_malloc && pool_type != RSPAMD_MEMPOOL_SHARED) {
diff --git a/src/libutil/mem_pool.h b/src/libutil/mem_pool.h
index 773be33d3..1554edbd2 100644
--- a/src/libutil/mem_pool.h
+++ b/src/libutil/mem_pool.h
@@ -154,6 +154,16 @@ RSPAMD_ATTR_ALLOC_SIZE(2) RSPAMD_ATTR_ALLOC_ALIGN(MIN_MEM_ALIGNMENT) RSPAMD_ATTR
 #define rspamd_mempool_alloc(pool, size) \
 	rspamd_mempool_alloc_((pool), (size), G_STRLOC)
 
+/**
+ * Notify external memory usage for memory pool
+ * @param pool
+ * @param size
+ * @param loc
+ */
+void rspamd_mempool_notify_alloc_ (rspamd_mempool_t *pool, gsize size, const gchar *loc);
+#define rspamd_mempool_notify_alloc(pool, size) \
+	rspamd_mempool_notify_alloc_((pool), (size), G_STRLOC)
+
 /**
  * Get memory and set it to zero
  * @param pool memory pool object
-- 
2.39.5