[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.

diff --git a/src/client/rspamc.c b/src/client/rspamc.c
index d71a3e089e6aa1f68c3321cc06bf4160fa9d6f08..0c8a270df763353918d7f168222005205de51186 100644 (file)
--- 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 f5c07ba66791ec8d45f334972b5f60b13469f6c7..a3f7f916767d752222e03ab4fe86e81ded822500 100644 (file)
--- 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 872b4dd6d3740c7f673ae969123e89c310586fff..fe3e9fb4f86bf690e9c64c5d770a64b41d5ac7ae 100644 (file)
--- 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 f02665afd33d22b39bad1fee674bae3a06a81fd0..dc3198609892a8c87a2f512e97549612127a46a3 100644 (file)
--- 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 be94f00c2692e2fea07fb17f7e1b9f9e757939d9..1779310bb7340f2b867cdf3af4d0a2b612095012 100644 (file)
--- 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_ */