[Project] Track more memory allocations in a task

9 files changed, 23 insertions, 3 deletions

diff --git a/src/libmime/email_addr.c b/src/libmime/email_addr.c
index 24b1d0111..4e09de6fb 100644
--- a/src/libmime/email_addr.c
+++ b/src/libmime/email_addr.c
@@ -110,6 +110,7 @@ rspamd_email_address_add (rspamd_mempool_t *pool,
 	guint nlen;
 
 	elt = g_malloc0 (sizeof (*elt));
+	rspamd_mempool_notify_alloc (pool, sizeof (*elt));
 
 	if (addr != NULL) {
 		memcpy (elt, addr, sizeof (*addr));
@@ -132,6 +133,7 @@ rspamd_email_address_add (rspamd_mempool_t *pool,
 		/* We need to unquote addr */
 		nlen = elt->domain_len + elt->user_len + 2;
 		elt->addr = g_malloc (nlen + 1);
+		rspamd_mempool_notify_alloc (pool, nlen + 1);
 		elt->addr_len = rspamd_snprintf ((char *)elt->addr, nlen, "%*s@%*s",
 				(gint)elt->user_len, elt->user,
 				(gint)elt->domain_len, elt->domain);
@@ -143,6 +145,7 @@ rspamd_email_address_add (rspamd_mempool_t *pool,
 		elt->name = rspamd_mime_header_decode (pool, name->str, name->len, NULL);
 	}
 
+	rspamd_mempool_notify_alloc (pool, name->len);
 	g_ptr_array_add (ar, elt);
 }
 
@@ -481,6 +484,7 @@ rspamd_email_address_from_mime (rspamd_mempool_t *pool,
 		break;
 	}
 
+	rspamd_mempool_notify_alloc (pool, cpy->len);
 	g_string_free (ns, TRUE);
 
 	return res;
diff --git a/src/libmime/message.c b/src/libmime/message.c
index e99fc4af7..95a1ab708 100644
--- a/src/libmime/message.c
+++ b/src/libmime/message.c
@@ -549,6 +549,8 @@ rspamd_normalize_text_part (struct rspamd_task *task,
 	rspamd_mempool_add_destructor (task->task_pool,
 			(rspamd_mempool_destruct_t) free_byte_array_callback,
 			part->utf_stripped_content);
+	rspamd_mempool_notify_alloc (task->task_pool,
+			part->utf_stripped_content->len);
 	rspamd_mempool_add_destructor (task->task_pool,
 			(rspamd_mempool_destruct_t) rspamd_ptr_array_free_hard,
 			part->newlines);
diff --git a/src/libmime/mime_encoding.c b/src/libmime/mime_encoding.c
index d98f96b85..0ba0e0edd 100644
--- a/src/libmime/mime_encoding.c
+++ b/src/libmime/mime_encoding.c
@@ -672,6 +672,8 @@ rspamd_mime_text_part_maybe_convert (struct rspamd_task *task,
 	part_content = g_byte_array_sized_new (text_part->parsed.len);
 	memcpy (part_content->data, text_part->parsed.begin, text_part->parsed.len);
 	part_content->len = text_part->parsed.len;
+	rspamd_mempool_notify_alloc (task->task_pool,
+			part_content->len);
 	rspamd_mempool_add_destructor (task->task_pool,
 			(rspamd_mempool_destruct_t)g_byte_array_unref, part_content);
 
diff --git a/src/libmime/mime_headers.c b/src/libmime/mime_headers.c
index e0c91c478..aed7575eb 100644
--- a/src/libmime/mime_headers.c
+++ b/src/libmime/mime_headers.c
@@ -787,6 +787,7 @@ rspamd_mime_header_decode (rspamd_mempool_t *pool, const gchar *in,
 	g_byte_array_free (token, TRUE);
 	g_byte_array_free (decoded, TRUE);
 	rspamd_mime_header_sanity_check (out);
+	rspamd_mempool_notify_alloc (pool, out->len);
 	ret = g_string_free (out, FALSE);
 	rspamd_mempool_add_destructor (pool, g_free, ret);
 
diff --git a/src/libmime/mime_parser.c b/src/libmime/mime_parser.c
index 18c54d13b..d989a8e2e 100644
--- a/src/libmime/mime_parser.c
+++ b/src/libmime/mime_parser.c
@@ -539,6 +539,7 @@ rspamd_mime_parse_normal_part (struct rspamd_task *task,
 			memcpy (parsed->str, part->raw_data.begin, parsed->len);
 			part->parsed_data.begin = parsed->str;
 			part->parsed_data.len = parsed->len;
+			rspamd_mempool_notify_alloc (task->task_pool, parsed->len);
 			rspamd_mempool_add_destructor (task->task_pool,
 					(rspamd_mempool_destruct_t)rspamd_fstring_free, parsed);
 		}
@@ -555,6 +556,7 @@ rspamd_mime_parse_normal_part (struct rspamd_task *task,
 			parsed->len = r;
 			part->parsed_data.begin = parsed->str;
 			part->parsed_data.len = parsed->len;
+			rspamd_mempool_notify_alloc (task->task_pool, parsed->len);
 			rspamd_mempool_add_destructor (task->task_pool,
 					(rspamd_mempool_destruct_t)rspamd_fstring_free, parsed);
 		}
@@ -566,6 +568,7 @@ rspamd_mime_parse_normal_part (struct rspamd_task *task,
 			parsed->len = part->raw_data.len;
 			part->parsed_data.begin = parsed->str;
 			part->parsed_data.len = parsed->len;
+			rspamd_mempool_notify_alloc (task->task_pool, parsed->len);
 			rspamd_mempool_add_destructor (task->task_pool,
 					(rspamd_mempool_destruct_t)rspamd_fstring_free, parsed);
 		}
@@ -577,6 +580,7 @@ rspamd_mime_parse_normal_part (struct rspamd_task *task,
 				parsed->str, &parsed->len);
 		part->parsed_data.begin = parsed->str;
 		part->parsed_data.len = parsed->len;
+		rspamd_mempool_notify_alloc (task->task_pool, parsed->len);
 		rspamd_mempool_add_destructor (task->task_pool,
 				(rspamd_mempool_destruct_t)rspamd_fstring_free, parsed);
 		break;
@@ -584,6 +588,7 @@ rspamd_mime_parse_normal_part (struct rspamd_task *task,
 		parsed = rspamd_fstring_sized_new (part->raw_data.len / 4 * 3 + 12);
 		r = rspamd_decode_uue_buf (part->raw_data.begin, part->raw_data.len,
 				parsed->str, parsed->allocated);
+		rspamd_mempool_notify_alloc (task->task_pool, parsed->len);
 		rspamd_mempool_add_destructor (task->task_pool,
 				(rspamd_mempool_destruct_t)rspamd_fstring_free, parsed);
 		if (r != -1) {
@@ -597,6 +602,7 @@ rspamd_mime_parse_normal_part (struct rspamd_task *task,
 			part->cte = RSPAMD_CTE_8BIT;
 			parsed->len = MIN (part->raw_data.len, parsed->allocated);
 			memcpy (parsed->str, part->raw_data.begin, parsed->len);
+			rspamd_mempool_notify_alloc (task->task_pool, parsed->len);
 			part->parsed_data.begin = parsed->str;
 			part->parsed_data.len = parsed->len;
 		}
diff --git a/src/libserver/html.c b/src/libserver/html.c
index 5bb3ca05b..d9cddb468 100644
--- a/src/libserver/html.c
+++ b/src/libserver/html.c
@@ -1801,6 +1801,7 @@ rspamd_html_process_img_tag (rspamd_mempool_t *pool, struct html_tag *tag,
 
 	if (hc->images == NULL) {
 		hc->images = g_ptr_array_sized_new (4);
+		rspamd_mempool_notify_alloc (pool, 4 * sizeof (gpointer) + sizeof (GPtrArray));
 		rspamd_mempool_add_destructor (pool, rspamd_ptr_array_free_hard,
 				hc->images);
 	}
@@ -2370,6 +2371,7 @@ rspamd_html_process_block_tag (rspamd_mempool_t *pool, struct html_tag *tag,
 
 	if (hc->blocks == NULL) {
 		hc->blocks = g_ptr_array_sized_new (64);
+		rspamd_mempool_notify_alloc (pool, 64 * sizeof (gpointer) + sizeof (GPtrArray));
 		rspamd_mempool_add_destructor (pool, rspamd_ptr_array_free_hard,
 				hc->blocks);
 	}
diff --git a/src/libserver/task.c b/src/libserver/task.c
index ccd959920..634ebb392 100644
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -1158,6 +1158,7 @@ rspamd_task_log_metric_res (struct rspamd_task *task,
 			rspamd_mempool_add_destructor (task->task_pool,
 					(rspamd_mempool_destruct_t)rspamd_fstring_free,
 					symbuf);
+			rspamd_mempool_notify_alloc (task->task_pool, symbuf->len);
 			res.begin = symbuf->str;
 			res.len = symbuf->len;
 			break;
@@ -1203,6 +1204,7 @@ rspamd_task_log_metric_res (struct rspamd_task *task,
 			rspamd_mempool_add_destructor (task->task_pool,
 					(rspamd_mempool_destruct_t) rspamd_fstring_free,
 					symbuf);
+			rspamd_mempool_notify_alloc (task->task_pool, symbuf->len);
 			res.begin = symbuf->str;
 			res.len = symbuf->len;
 			break;
diff --git a/src/libstat/stat_process.c b/src/libstat/stat_process.c
index 603e0ac96..6bd17f38f 100644
--- a/src/libstat/stat_process.c
+++ b/src/libstat/stat_process.c
@@ -142,6 +142,7 @@ rspamd_stat_process_tokenize (struct rspamd_stat_ctx *st_ctx,
 	task->tokens = g_ptr_array_sized_new (reserved_len);
 	rspamd_mempool_add_destructor (task->task_pool,
 			rspamd_ptr_array_free_hard, task->tokens);
+	rspamd_mempool_notify_alloc (task->task_pool, reserved_len * sizeof (gpointer));
 	pdiff = rspamd_mempool_get_variable (task->task_pool, "parts_distance");
 
 	PTR_ARRAY_FOREACH (MESSAGE_FIELD (task, text_parts), i, part) {
diff --git a/src/libutil/mem_pool.c b/src/libutil/mem_pool.c
index f2272fa7c..df8c3cb07 100644
--- a/src/libutil/mem_pool.c
+++ b/src/libutil/mem_pool.c
@@ -445,10 +445,10 @@ RSPAMD_ATTR_ALLOC_SIZE(2) RSPAMD_ATTR_ALLOC_ALIGN(MIN_MEM_ALIGNMENT) RSPAMD_ATTR
 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 (pool && G_UNLIKELY (pool->priv->flags & RSPAMD_MEMPOOL_DEBUG)) {
+		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) {