From eba3948ce08b642945ca02942d9c391a5a13d8f4 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Thu, 7 Nov 2019 14:44:45 +0000
Subject: [PATCH] [Rework] Another phase of finish actions rework

---
 src/libserver/worker_util.c | 99 ++++++++++++++++++++++++++++---------
 1 file changed, 75 insertions(+), 24 deletions(-)

diff --git a/src/libserver/worker_util.c b/src/libserver/worker_util.c
index 63342a72a..e3927063d 100644
--- a/src/libserver/worker_util.c
+++ b/src/libserver/worker_util.c
@@ -184,29 +184,33 @@ rspamd_worker_terminate_handlers (struct rspamd_worker *w)
 
 		return;
 	}
-	else if (w->nconns > 0) {
-		/*
-		 * Wait until all connections are terminated
-		 */
-		w->state = rspamd_worker_wait_connections;
-	}
 	else {
-		/*
-		 * Start finish scripts
-		 */
-		w->state = rspamd_worker_wait_final_scripts;
-		msg_info ("performing finishing actions");
-
-		if ((w->flags & RSPAMD_WORKER_SCANNER) &&
-			rspamd_worker_call_finish_handlers (w)) {
-			w->state = rspamd_worker_wait_final_scripts;
+		if (w->nconns > 0) {
+			/*
+			 * Wait until all connections are terminated
+			 */
+			w->state = rspamd_worker_wait_connections;
 		}
 		else {
 			/*
-			 * We are done now
+			 * Start finish scripts
 			 */
-			ev_break (actx->event_loop, EVBREAK_ALL);
-			w->state = rspamd_worker_wanna_die;
+			if (w->state != rspamd_worker_wait_final_scripts) {
+				w->state = rspamd_worker_wait_final_scripts;
+				msg_info ("performing finishing actions");
+
+				if ((w->flags & RSPAMD_WORKER_SCANNER) &&
+					rspamd_worker_call_finish_handlers (w)) {
+					w->state = rspamd_worker_wait_final_scripts;
+				}
+				else {
+					/*
+					 * We are done now
+					 */
+					ev_break (actx->event_loop, EVBREAK_ALL);
+					w->state = rspamd_worker_wanna_die;
+				}
+			}
 		}
 	}
 }
@@ -214,12 +218,34 @@ rspamd_worker_terminate_handlers (struct rspamd_worker *w)
 static void
 rspamd_worker_on_delayed_shutdown (EV_P_ ev_timer *w, int revents)
 {
+	struct rspamd_worker *worker = (struct rspamd_worker *)w->data;
+
+	worker->state = rspamd_worker_wanna_die;
+	ev_timer_stop (EV_A_ w);
 	ev_break (loop, EVBREAK_ALL);
 #ifdef WITH_GPERF_TOOLS
 	ProfilerStop ();
 #endif
 }
 
+static void
+rspamd_worker_shutdown_check (EV_P_ ev_timer *w, int revents)
+{
+	struct rspamd_worker *worker = (struct rspamd_worker *)w->data;
+
+	if (worker->state != rspamd_worker_wanna_die) {
+		ev_timer_again (EV_A_ w);
+		rspamd_worker_terminate_handlers (worker);
+
+		if (worker->state == rspamd_worker_wanna_die) {
+			ev_timer_stop (EV_A_ w);
+		}
+	}
+	else {
+		ev_timer_stop (EV_A_ w);
+	}
+}
+
 /*
  * Config reload is designed by sending sigusr2 to active workers and pending shutdown of them
  */
@@ -228,7 +254,7 @@ rspamd_worker_usr2_handler (struct rspamd_worker_signal_handler *sigh, void *arg
 {
 	/* Do not accept new connections, preparing to end worker's process */
 	if (sigh->worker->state == rspamd_worker_state_running) {
-		static ev_timer shutdown_ev;
+		static ev_timer shutdown_ev, shutdown_check_ev;
 		ev_tstamp shutdown_ts;
 
 		shutdown_ts = MAX (SOFT_SHUTDOWN_TIME,
@@ -243,9 +269,18 @@ rspamd_worker_usr2_handler (struct rspamd_worker_signal_handler *sigh, void *arg
 				G_STRFUNC,
 				"worker's shutdown is pending in %.2f sec",
 				shutdown_ts);
+
+		/* Soft shutdown timer */
+		shutdown_ev.data = sigh->worker;
 		ev_timer_init (&shutdown_ev, rspamd_worker_on_delayed_shutdown,
 				shutdown_ts, 0.0);
 		ev_timer_start (sigh->event_loop, &shutdown_ev);
+
+		/* This timer checks if we are ready to die and is called frequently */
+		shutdown_check_ev.data = sigh->worker;
+		ev_timer_init (&shutdown_check_ev, rspamd_worker_shutdown_check,
+				0.5, 0.5);
+		ev_timer_start (sigh->event_loop, &shutdown_check_ev);
 		rspamd_worker_stop_accept (sigh->worker);
 	}
 
@@ -272,7 +307,11 @@ static gboolean
 rspamd_worker_term_handler (struct rspamd_worker_signal_handler *sigh, void *arg)
 {
 	if (sigh->worker->state == rspamd_worker_state_running) {
-		static ev_timer shutdown_ev;
+		static ev_timer shutdown_ev, shutdown_check_ev;
+		ev_tstamp shutdown_ts;
+
+		shutdown_ts = MAX (SOFT_SHUTDOWN_TIME,
+				sigh->worker->srv->cfg->task_timeout * 2.0);
 
 		rspamd_worker_ignore_signal (sigh);
 		sigh->worker->state = rspamd_worker_state_terminating;
@@ -283,11 +322,23 @@ rspamd_worker_term_handler (struct rspamd_worker_signal_handler *sigh, void *arg
 				"terminating after receiving signal %s",
 				g_strsignal (sigh->signo));
 
-		rspamd_worker_terminate_handlers (sigh->worker);
-		ev_timer_init (&shutdown_ev, rspamd_worker_on_delayed_shutdown,
-				0.0, 0.0);
-		ev_timer_start (sigh->event_loop, &shutdown_ev);
 		rspamd_worker_stop_accept (sigh->worker);
+		rspamd_worker_terminate_handlers (sigh->worker);
+
+		/* Check if we are ready to die */
+		if (sigh->worker->state != rspamd_worker_wanna_die) {
+			/* This timer is called when we have no choices but to die */
+			shutdown_ev.data = sigh->worker;
+			ev_timer_init (&shutdown_ev, rspamd_worker_on_delayed_shutdown,
+					shutdown_ts, 0.0);
+			ev_timer_start (sigh->event_loop, &shutdown_ev);
+
+			/* This timer checks if we are ready to die and is called frequently */
+			shutdown_check_ev.data = sigh->worker;
+			ev_timer_init (&shutdown_check_ev, rspamd_worker_shutdown_check,
+					0.5, 0.5);
+			ev_timer_start (sigh->event_loop, &shutdown_check_ev);
+		}
 	}
 
 	/* Stop reacting on signals */
-- 
2.39.5