commit 196499c: [Minor] Add health check command for control interface

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2021-07-21 11:56:17 +0100
URL: https://github.com/rspamd/rspamd/commit/196499cda2dcf4d4ae1faaecf9cd68fda55fac81

[Minor] Add health check command for control interface

---
 src/libserver/rspamd_control.c | 28 ++++++++++++++++++++++++++++
 src/libserver/rspamd_control.h | 10 ++++++++++
 2 files changed, 38 insertions(+)

diff --git a/src/libserver/rspamd_control.c b/src/libserver/rspamd_control.c
index f1de1d29c..0021485c3 100644
--- a/src/libserver/rspamd_control.c
+++ b/src/libserver/rspamd_control.c
@@ -850,6 +850,31 @@ rspamd_control_handle_on_fork (struct rspamd_srv_command *cmd,
 	}
 }
 
+static void
+rspamd_fill_health_reply (struct rspamd_main *srv, struct rspamd_srv_reply *rep)
+{
+	GHashTableIter it;
+	gpointer k, v;
+
+	memset (&rep->reply.health, 0, sizeof (rep->reply));
+	g_hash_table_iter_init (&it, srv->workers);
+
+	while (g_hash_table_iter_next (&it, &k, &v)) {
+		struct rspamd_worker *wrk = (struct rspamd_worker *)v;
+
+		if (wrk->hb.nbeats < 0) {
+			rep->reply.health.workers_hb_lost ++;
+		}
+		else if (rspamd_worker_is_scanner (wrk)) {
+			rep->reply.health.scanners_count ++;
+		}
+
+		rep->reply.health.workers_count ++;
+	}
+
+	rep->reply.status = (g_hash_table_size (srv->workers) > 0);
+}
+
 
 static void
 rspamd_srv_handler (EV_P_ ev_io *w, int revents)
@@ -979,6 +1004,9 @@ rspamd_srv_handler (EV_P_ ev_io *w, int revents)
 				worker->hb.last_event = ev_time ();
 				rdata->rep.reply.heartbeat.status = 0;
 				break;
+			case RSPAMD_SRV_HEALTH:
+				rspamd_fill_health_reply (srv, &rdata->rep);
+				break;
 			default:
 				msg_err ("unknown command type: %d", cmd.type);
 				break;
diff --git a/src/libserver/rspamd_control.h b/src/libserver/rspamd_control.h
index 21ab1a663..1b5756022 100644
--- a/src/libserver/rspamd_control.h
+++ b/src/libserver/rspamd_control.h
@@ -48,6 +48,7 @@ enum rspamd_srv_type {
 	RSPAMD_SRV_LOG_PIPE,
 	RSPAMD_SRV_ON_FORK,
 	RSPAMD_SRV_HEARTBEAT,
+	RSPAMD_SRV_HEALTH,
 };
 
 enum rspamd_log_pipe_type {
@@ -172,6 +173,9 @@ struct rspamd_srv_command {
 			guint status;
 			/* TODO: add more fields */
 		} heartbeat;
+		struct {
+			guint status;
+		} health;
 	} cmd;
 };
 
@@ -197,6 +201,12 @@ struct rspamd_srv_reply {
 		struct {
 			gint status;
 		} heartbeat;
+		struct {
+			guint status;
+			guint workers_count;
+			guint scanners_count;
+			guint workers_hb_lost;
+		} health;
 	} reply;
 };