[Feature] Support metrics command in normal/proxy workers

1 files changed, 12 insertions, 1 deletions

diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c
index 5bf1fb8f2..4f08e81b9 100644
--- a/src/rspamd_proxy.c
+++ b/src/rspamd_proxy.c
@@ -160,6 +160,7 @@ struct rspamd_proxy_ctx {
 	/* Language detector */
 	struct rspamd_lang_detector *lang_det;
 	double task_timeout;
+	struct rspamd_main *srv;
 };
 
 enum rspamd_backend_flags {
@@ -1734,6 +1735,8 @@ rspamd_proxy_scan_self_reply(struct rspamd_task *task)
 	int out_type = UCL_EMIT_JSON_COMPACT;
 	const char *ctype = "application/json";
 	const rspamd_ftok_t *accept_hdr = rspamd_task_get_request_header(task, "Accept");
+	rspamd_fstring_t *output;
+	struct rspamd_stat stat_copy;
 
 	if (accept_hdr && rspamd_substring_search(accept_hdr->begin, accept_hdr->len,
 											  "application/msgpack", sizeof("application/msgpack") - 1) != -1) {
@@ -1759,6 +1762,13 @@ rspamd_proxy_scan_self_reply(struct rspamd_task *task)
 		rspamd_http_message_set_body(msg, "pong" CRLF, 6);
 		ctype = "text/plain";
 		break;
+	case CMD_METRICS:
+		memcpy(&stat_copy, session->ctx->srv->stat, sizeof(stat_copy));
+		output = rspamd_metrics_to_prometheus_string(
+			rspamd_worker_metrics_object(task->cfg, &stat_copy, ev_time() - session->ctx->srv->start_time));
+		rspamd_http_message_set_body_from_fstring_steal(msg, output);
+		ctype = "application/openmetrics-text; version=1.0.0; charset=utf-8";
+		break;
 	default:
 		msg_err_task("BROKEN");
 		break;
@@ -1903,7 +1913,7 @@ rspamd_proxy_self_scan(struct rspamd_proxy_session *session)
 		task->flags |= RSPAMD_TASK_FLAG_SKIP;
 	}
 	else {
-		if (task->cmd == CMD_PING) {
+		if (task->cmd == CMD_PING || task->cmd == CMD_METRICS) {
 			task->flags |= RSPAMD_TASK_FLAG_SKIP;
 		}
 		else {
@@ -2412,6 +2422,7 @@ start_rspamd_proxy(struct rspamd_worker *worker)
 
 	g_assert(rspamd_worker_check_context(worker->ctx, rspamd_rspamd_proxy_magic));
 	ctx->cfg = worker->srv->cfg;
+	ctx->srv = worker->srv;
 	ctx->event_loop = rspamd_prepare_worker(worker, "rspamd_proxy",
 											proxy_accept_socket);