4 files changed, 104 insertions, 0 deletions

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_ */