[Feature] Add termination callbacks for workers

4 files changed, 33 insertions, 7 deletions

diff --git a/src/controller.c b/src/controller.c
index 99c81b95b..1c533cc0a 100644
--- a/src/controller.c
+++ b/src/controller.c
@@ -2473,13 +2473,13 @@ rspamd_controller_store_saved_stats (struct rspamd_controller_worker_ctx *ctx)
 	fd = open (ctx->saved_stats_path, O_WRONLY|O_CREAT|O_TRUNC, 00644);
 
 	if (fd == -1) {
-		msg_err_ctx ("cannot load controller stats from %s: %s",
+		msg_err_ctx ("cannot open for writing controller stats from %s: %s",
 				ctx->saved_stats_path, strerror (errno));
 		return;
 	}
 
 	if (rspamd_file_lock (fd, FALSE) == -1) {
-		msg_err_ctx ("cannot load controller stats from %s: %s",
+		msg_err_ctx ("cannot lock controller stats in %s: %s",
 				ctx->saved_stats_path, strerror (errno));
 		close (fd);
 
@@ -2660,6 +2660,19 @@ init_controller_worker (struct rspamd_config *cfg)
 	return ctx;
 }
 
+static void
+rspamd_controller_on_terminate (struct rspamd_worker *worker)
+{
+	struct rspamd_controller_worker_ctx *ctx = worker->ctx;
+
+	rspamd_controller_store_saved_stats (ctx);
+
+	if (ctx->rrd) {
+		event_del (ctx->rrd_event);
+		rspamd_rrd_close (ctx->rrd);
+	}
+}
+
 /*
  * Start worker process
  */
@@ -2697,6 +2710,8 @@ start_controller_worker (struct rspamd_worker *worker)
 				DEFAULT_STATS_PATH);
 	}
 
+	g_ptr_array_add (worker->finish_actions,
+			(gpointer)rspamd_controller_on_terminate);
 	rspamd_controller_load_saved_stats (ctx);
 
 	/* RRD collector */
@@ -2832,11 +2847,6 @@ start_controller_worker (struct rspamd_worker *worker)
 	rspamd_stat_close ();
 	rspamd_http_router_free (ctx->http);
 	rspamd_log_close (worker->srv->logger);
-	rspamd_controller_store_saved_stats (ctx);
-
-	if (ctx->rrd) {
-		rspamd_rrd_close (ctx->rrd);
-	}
 
 	if (ctx->cached_password.len > 0) {
 		m = (gpointer)ctx->cached_password.begin;
diff --git a/src/libserver/worker_util.c b/src/libserver/worker_util.c
index b7aa5035f..2a0a76d04 100644
--- a/src/libserver/worker_util.c
+++ b/src/libserver/worker_util.c
@@ -89,6 +89,17 @@ rspamd_get_worker_by_type (struct rspamd_config *cfg, GQuark type)
 
 sig_atomic_t wanna_die = 0;
 
+static void
+rspamd_worker_terminate_handlers (struct rspamd_worker *w)
+{
+	guint i;
+	void (*cb)(struct rspamd_worker *);
+
+	for (i = 0; i < w->finish_actions->len; i ++) {
+		cb = g_ptr_array_index (w->finish_actions, i);
+		cb (w);
+	}
+}
 /*
  * Config reload is designed by sending sigusr2 to active workers and pending shutdown of them
  */
@@ -102,6 +113,7 @@ rspamd_worker_usr2_handler (struct rspamd_worker_signal_handler *sigh, void *arg
 		tv.tv_sec = SOFT_SHUTDOWN_TIME;
 		tv.tv_usec = 0;
 		wanna_die = 1;
+		rspamd_worker_terminate_handlers (sigh->worker);
 		rspamd_default_log_function (G_LOG_LEVEL_INFO,
 				sigh->worker->srv->server_pool->tag.tagname,
 				sigh->worker->srv->server_pool->tag.uid,
@@ -134,6 +146,7 @@ rspamd_worker_term_handler (struct rspamd_worker_signal_handler *sigh, void *arg
 				G_STRFUNC,
 				"terminating after receiving signal %s",
 				g_strsignal (sigh->signo));
+		rspamd_worker_terminate_handlers (sigh->worker);
 		wanna_die = 1;
 		tv.tv_sec = 0;
 		tv.tv_usec = 0;
@@ -525,6 +538,7 @@ rspamd_fork_worker (struct rspamd_main *rspamd_main,
 	memcpy (wrk->cf, cf, sizeof (struct rspamd_worker_conf));
 	wrk->index = index;
 	wrk->ctx = cf->ctx;
+	wrk->finish_actions = g_ptr_array_new ();
 
 	wrk->pid = fork ();
 
diff --git a/src/rspamd.c b/src/rspamd.c
index 056de6e12..3ddee3ff4 100644
--- a/src/rspamd.c
+++ b/src/rspamd.c
@@ -692,6 +692,7 @@ wait_for_workers (gpointer key, gpointer value, gpointer unused)
 			g_quark_to_string (w->type), w->pid,
 			WTERMSIG (res) == SIGKILL ? "hardly" : "softly");
 	event_del (&w->srv_ev);
+	g_ptr_array_free (w->finish_actions, TRUE);
 	g_free (w->cf);
 	g_free (w);
 
diff --git a/src/rspamd.h b/src/rspamd.h
index ffebfe387..999061a99 100644
--- a/src/rspamd.h
+++ b/src/rspamd.h
@@ -62,6 +62,7 @@ struct rspamd_worker {
 	                                     main process. [0] - main, [1] - worker			*/
 	struct event srv_ev;            /**< used by main for read workers' requests		*/
 	gpointer control_data;          /**< used by control protocol to handle commands	*/
+	GPtrArray *finish_actions;      /**< called when worker is terminated				*/
 };
 
 struct rspamd_abstract_worker_ctx {