[Fix] Fix on_finish scripts and async handlers

3 files changed, 44 insertions, 11 deletions

diff --git a/src/controller.c b/src/controller.c
index 65b40ca5a..6375c844b 100644
--- a/src/controller.c
+++ b/src/controller.c
@@ -2930,7 +2930,7 @@ lua_csession_send_string (lua_State *L)
 	return 0;
 }
 
-static void
+static gboolean
 rspamd_controller_on_terminate (struct rspamd_worker *worker)
 {
 	struct rspamd_controller_worker_ctx *ctx = worker->ctx;
@@ -2942,6 +2942,8 @@ rspamd_controller_on_terminate (struct rspamd_worker *worker)
 		event_del (ctx->rrd_event);
 		rspamd_rrd_close (ctx->rrd);
 	}
+
+	return FALSE;
 }
 
 /*
diff --git a/src/libserver/worker_util.c b/src/libserver/worker_util.c
index 6261e1991..4c1bb4a1b 100644
--- a/src/libserver/worker_util.c
+++ b/src/libserver/worker_util.c
@@ -87,16 +87,21 @@ rspamd_get_worker_by_type (struct rspamd_config *cfg, GQuark type)
 	return NULL;
 }
 
-static void
+static gboolean
 rspamd_worker_terminate_handlers (struct rspamd_worker *w)
 {
 	guint i;
-	void (*cb)(struct rspamd_worker *);
+	gboolean (*cb)(struct rspamd_worker *);
+	gboolean ret = FALSE;
 
 	for (i = 0; i < w->finish_actions->len; i ++) {
 		cb = g_ptr_array_index (w->finish_actions, i);
-		cb (w);
+		if (cb (w)) {
+			ret = TRUE;
+		}
 	}
+
+	return ret;
 }
 /*
  * Config reload is designed by sending sigusr2 to active workers and pending shutdown of them
@@ -144,11 +149,16 @@ 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);
-		sigh->worker->wanna_die = 1;
-		tv.tv_sec = 0;
+
 		tv.tv_usec = 0;
+		if (rspamd_worker_terminate_handlers (sigh->worker)) {
+			tv.tv_sec =  SOFT_SHUTDOWN_TIME;
+		}
+		else {
+			tv.tv_sec = 0;
+		}
 
+		sigh->worker->wanna_die = 1;
 		event_base_loopexit (sigh->base, &tv);
 #ifdef WITH_GPERF_TOOLS
 		ProfilerStop ();
diff --git a/src/worker.c b/src/worker.c
index fcf4e2c73..fcb9d9fe3 100644
--- a/src/worker.c
+++ b/src/worker.c
@@ -73,7 +73,19 @@ worker_t normal_worker = {
         G_STRFUNC, \
         __VA_ARGS__)
 
-static void
+static gboolean
+rspamd_worker_finalize (gpointer user_data)
+{
+	struct rspamd_task *task = user_data;
+	struct timeval tv = {.tv_sec = 0, .tv_usec = 0};
+
+	msg_info_task ("finishing actions has been processed, terminating");
+	event_base_loopexit (task->ev_base, &tv);
+
+	return TRUE;
+}
+
+static gboolean
 rspamd_worker_call_finish_handlers (struct rspamd_worker *worker)
 {
 	struct rspamd_task *task;
@@ -88,7 +100,7 @@ rspamd_worker_call_finish_handlers (struct rspamd_worker *worker)
 		task->resolver = ctx->resolver;
 		task->ev_base = ctx->ev_base;
 		task->s = rspamd_session_create (task->task_pool,
-				NULL,
+				rspamd_worker_finalize,
 				NULL,
 				(event_finalizer_t) rspamd_task_free,
 				task);
@@ -96,8 +108,13 @@ rspamd_worker_call_finish_handlers (struct rspamd_worker *worker)
 		DL_FOREACH (cfg->finish_callbacks, sc) {
 			lua_call_finish_script (cfg->lua_state, sc, task);
 		}
+
+		if (rspamd_session_pending (task->s)) {
+			return TRUE;
+		}
 	}
 
+	return FALSE;
 }
 
 /*
@@ -558,13 +575,17 @@ init_worker (struct rspamd_config *cfg)
 	return ctx;
 }
 
-static void
+static gboolean
 rspamd_worker_on_terminate (struct rspamd_worker *worker)
 {
 	if (worker->nconns == 0) {
 		msg_info ("performing finishing actions");
-		rspamd_worker_call_finish_handlers (worker);
+		if (rspamd_worker_call_finish_handlers (worker)) {
+			return TRUE;
+		}
 	}
+
+	return FALSE;
 }
 
 /*