[Feature] Add profiling support in client output

It is now possible to specify `--profile` option to Rspamd to get
execution time for all symbols that were executed for a specific
message. This still doesn't count async returns time.

5 files changed, 147 insertions, 0 deletions

diff --git a/src/client/rspamc.c b/src/client/rspamc.c
index d71a3e089..0c8a270df 100644
--- a/src/client/rspamc.c
+++ b/src/client/rspamc.c
@@ -59,6 +59,7 @@ static gboolean extended_urls = FALSE;
 static gboolean mime_output = FALSE;
 static gboolean empty_input = FALSE;
 static gboolean compressed = FALSE;
+static gboolean profile = FALSE;
 static gchar *key = NULL;
 static GList *children;
 
@@ -140,6 +141,8 @@ static GOptionEntry entries[] =
 	   "Learn the specified fuzzy symbol", NULL },
 	{ "compressed", 'z', 0, G_OPTION_ARG_NONE, &compressed,
 	   "Enable zstd compression", NULL },
+	{ "profile", '\0', 0, G_OPTION_ARG_NONE, &profile,
+	   "Profile symbols execution time", NULL },
 	{ "dictionary", 'D', 0, G_OPTION_ARG_FILENAME, &dictionary,
 	   "Use dictionary to compress data", NULL },
 	{ NULL, 0, 0, G_OPTION_ARG_NONE, NULL, NULL, NULL }
@@ -519,6 +522,9 @@ add_options (GQueue *opts)
 	if (extended_urls) {
 		ADD_CLIENT_HEADER (opts, "URL-Format", "extended");
 	}
+	if (profile) {
+		ADD_CLIENT_HEADER (opts, "Profile", "true");
+	}
 
 	hdr = http_headers;
 
@@ -634,6 +640,39 @@ rspamc_metric_output (FILE *out, const ucl_object_t *obj)
 	g_ptr_array_free (sym_ptr, TRUE);
 }
 
+static gint
+rspamc_profile_sort_func (gconstpointer a, gconstpointer b)
+{
+	ucl_object_t * const *ua = a, * const *ub = b;
+
+	return ucl_object_compare (*ua, *ub);
+}
+
+static void
+rspamc_profile_output (FILE *out, const ucl_object_t *obj)
+{
+	ucl_object_iter_t it = NULL;
+	const ucl_object_t *cur;
+	guint i;
+	GPtrArray *ar;
+
+	ar = g_ptr_array_sized_new (obj->len);
+
+	while ((cur = ucl_object_iterate (obj, &it, true)) != NULL) {
+		g_ptr_array_add (ar, (void *)cur);
+	}
+
+	g_ptr_array_sort (ar, rspamc_profile_sort_func);
+
+	for (i = 0; i < ar->len; i ++) {
+		cur = (const ucl_object_t *)g_ptr_array_index (ar, i);
+		rspamd_fprintf (out, "\t%s: %.3f usec\n",
+				ucl_object_key (cur), ucl_object_todouble (cur));
+	}
+
+	g_ptr_array_free (ar, TRUE);
+}
+
 static void
 rspamc_symbols_output (FILE *out, ucl_object_t *obj)
 {
@@ -682,6 +721,10 @@ rspamc_symbols_output (FILE *out, ucl_object_t *obj)
 			rspamd_fprintf (out, "DKIM-Signature: %s\n", ucl_object_tostring (
 					cur));
 		}
+		else if (g_ascii_strcasecmp (ucl_object_key (cur), "profile") == 0) {
+			rspamd_fprintf (out, "Profile data:\n");
+			rspamc_profile_output (out, cur);
+		}
 		else if (cur->type == UCL_OBJECT &&
 				g_ascii_strcasecmp (ucl_object_key (cur), "rmilter") != 0) {
 			/* Parse metric */
diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index f5c07ba66..a3f7f9167 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -101,6 +101,7 @@
 #define MLEN_HEADER "Message-Length"
 #define USER_AGENT_HEADER "User-Agent"
 #define MTA_TAG_HEADER "MTA-Tag"
+#define PROFILE_HEADER "Profile"
 
 
 static GQuark
@@ -404,6 +405,9 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 						debug_task ("pass all filters");
 					}
 				}
+				IF_HEADER (PROFILE_HEADER) {
+					task->flags |= RSPAMD_TASK_FLAG_PROFILE;
+				}
 				break;
 			case 's':
 			case 'S':
@@ -1032,6 +1036,32 @@ rspamd_ucl_tospamc_output (const ucl_object_t *top,
 	}
 }
 
+static void
+rspamd_protocol_output_profiling (struct rspamd_task *task,
+		ucl_object_t *top)
+{
+	GHashTable *tbl;
+	GHashTableIter it;
+	gpointer k, v;
+	ucl_object_t *prof;
+	gdouble val;
+
+	prof = ucl_object_typed_new (UCL_OBJECT);
+	tbl = rspamd_mempool_get_variable (task->task_pool, "profile");
+
+	if (tbl) {
+		g_hash_table_iter_init (&it, tbl);
+
+		while (g_hash_table_iter_next (&it, &k, &v)) {
+			val = *(gdouble *)v;
+			ucl_object_insert_key (prof, ucl_object_fromdouble (val),
+					(const char *)k, 0, false);
+		}
+	}
+
+	ucl_object_insert_key (top, prof, "profile", 0, false);
+}
+
 ucl_object_t *
 rspamd_protocol_write_ucl (struct rspamd_task *task)
 {
@@ -1082,6 +1112,10 @@ rspamd_protocol_write_ucl (struct rspamd_task *task)
 		}
 	}
 
+	if (G_UNLIKELY (RSPAMD_TASK_IS_PROFILING (task))) {
+		rspamd_protocol_output_profiling (task, top);
+	}
+
 	ucl_object_insert_key (top, ucl_object_fromstring (task->message_id),
 			"message-id", 0, false);
 
diff --git a/src/libserver/symbols_cache.c b/src/libserver/symbols_cache.c
index 872b4dd6d..fe3e9fb4f 100644
--- a/src/libserver/symbols_cache.c
+++ b/src/libserver/symbols_cache.c
@@ -1289,6 +1289,10 @@ rspamd_symbols_cache_check_symbol (struct rspamd_task *task,
 			t2 = rspamd_get_ticks ();
 			diff = (t2 - t1) * 1e6;
 
+			if (G_UNLIKELY (RSPAMD_TASK_IS_PROFILING (task))) {
+				rspamd_task_profile_set (task, item->symbol, diff);
+			}
+
 			if (total_diff) {
 				*total_diff += diff;
 			}
diff --git a/src/libserver/task.c b/src/libserver/task.c
index f02665afd..dc3198609 100644
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -1473,3 +1473,50 @@ rspamd_task_add_request_header (struct rspamd_task *task,
 		g_hash_table_replace (task->request_headers, name, ret);
 	}
 }
+
+
+void
+rspamd_task_profile_set (struct rspamd_task *task, const gchar *key,
+		gdouble value)
+{
+	GHashTable *tbl;
+	gdouble *pval;
+
+	if (key == NULL) {
+		return;
+	}
+
+	tbl = rspamd_mempool_get_variable (task->task_pool, "profile");
+
+	if (tbl == NULL) {
+		tbl = g_hash_table_new (g_str_hash, g_str_equal);
+		rspamd_mempool_set_variable (task->task_pool, "profile", tbl,
+				(rspamd_mempool_destruct_t)g_hash_table_unref);
+	}
+
+	pval = g_hash_table_lookup (tbl, key);
+
+	if (pval == NULL) {
+		pval = rspamd_mempool_alloc (task->task_pool, sizeof (*pval));
+		*pval = value;
+		g_hash_table_insert (tbl, (void *)key, pval);
+	}
+	else {
+		*pval = value;
+	}
+}
+
+gdouble*
+rspamd_task_profile_get (struct rspamd_task *task, const gchar *key)
+{
+	GHashTable *tbl;
+	gdouble *pval = NULL;
+
+	tbl = rspamd_mempool_get_variable (task->task_pool, "profile");
+
+	if (tbl != NULL) {
+		pval = g_hash_table_lookup (tbl, key);
+	}
+
+	return pval;
+}
diff --git a/src/libserver/task.h b/src/libserver/task.h
index be94f00c2..1779310bb 100644
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -103,6 +103,7 @@ enum rspamd_task_stage {
 #define RSPAMD_TASK_FLAG_EMPTY (1 << 22)
 #define RSPAMD_TASK_FLAG_LOCAL_CLIENT (1 << 23)
 #define RSPAMD_TASK_FLAG_COMPRESSED (1 << 24)
+#define RSPAMD_TASK_FLAG_PROFILE (1 << 25)
 
 #define RSPAMD_TASK_IS_SKIPPED(task) (((task)->flags & RSPAMD_TASK_FLAG_SKIP))
 #define RSPAMD_TASK_IS_JSON(task) (((task)->flags & RSPAMD_TASK_FLAG_JSON))
@@ -110,6 +111,7 @@ enum rspamd_task_stage {
 #define RSPAMD_TASK_IS_PROCESSED(task) (((task)->processed_stages & RSPAMD_TASK_STAGE_DONE))
 #define RSPAMD_TASK_IS_CLASSIFIED(task) (((task)->processed_stages & RSPAMD_TASK_STAGE_CLASSIFIERS))
 #define RSPAMD_TASK_IS_EMPTY(task) (((task)->flags & RSPAMD_TASK_FLAG_EMPTY))
+#define RSPAMD_TASK_IS_PROFILING(task) (((task)->flags & RSPAMD_TASK_FLAG_PROFILE))
 
 struct rspamd_email_address;
 enum rspamd_newlines_type;
@@ -309,4 +311,21 @@ void rspamd_task_add_request_header (struct rspamd_task *task,
  */
 void rspamd_task_write_log (struct rspamd_task *task);
 
+/**
+ * Set profiling value for a specific key
+ * @param task
+ * @param key
+ * @param value
+ */
+void rspamd_task_profile_set (struct rspamd_task *task, const gchar *key,
+		gdouble value);
+
+/**
+ * Get value for a specific profiling key
+ * @param task
+ * @param key
+ * @return
+ */
+gdouble* rspamd_task_profile_get (struct rspamd_task *task, const gchar *key);
+
 #endif /* TASK_H_ */