From bfd6f0dba6a5147e7718c816d064260762d37014 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@rspamd.com>
Date: Fri, 19 Aug 2022 21:13:03 +0100
Subject: [PATCH] [Minor] Add utility function for task timeouts

---
 src/controller.c            | 11 ++---------
 src/libserver/worker_util.c | 29 +++++++++++++++++++++++++++++
 src/libserver/worker_util.h |  9 +++++++++
 src/rspamd_proxy.c          |  5 ++++-
 src/worker.c                | 11 ++---------
 5 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/src/controller.c b/src/controller.c
index 765aa1623..adebce251 100644
--- a/src/controller.c
+++ b/src/controller.c
@@ -2189,7 +2189,7 @@ rspamd_controller_handle_scan (struct rspamd_http_connection_entry *conn_ent,
 		goto end;
 	}
 
-	if (ctx->task_timeout > 0.0) {
+	if (!isnan(ctx->task_timeout) && ctx->task_timeout > 0.0) {
 		task->timeout_ev.data = task;
 		ev_timer_init (&task->timeout_ev, rspamd_task_timeout,
 				ctx->task_timeout, ctx->task_timeout);
@@ -4080,14 +4080,7 @@ start_controller_worker (struct rspamd_worker *worker)
 			rspamd_ftok_icase_equal, rspamd_fstring_mapped_ftok_free,
 			rspamd_plugin_cbdata_dtor);
 
-	if (isnan (ctx->task_timeout)) {
-		if (isnan (ctx->cfg->task_timeout)) {
-			ctx->task_timeout = 0;
-		}
-		else {
-			ctx->task_timeout = ctx->cfg->task_timeout;
-		}
-	}
+	ctx->task_timeout = rspamd_worker_check_and_adjust_timeout(ctx->cfg, ctx->task_timeout);
 
 	if (ctx->secure_ip != NULL) {
 		rspamd_config_radix_from_ucl (ctx->cfg, ctx->secure_ip,
diff --git a/src/libserver/worker_util.c b/src/libserver/worker_util.c
index 6392aa8a9..8c716fb94 100644
--- a/src/libserver/worker_util.c
+++ b/src/libserver/worker_util.c
@@ -57,6 +57,8 @@
 
 #ifdef HAVE_SYS_WAIT_H
 #include <sys/wait.h>
+#include <math.h>
+
 #endif
 
 #include "contrib/libev/ev.h"
@@ -2221,3 +2223,30 @@ rspamd_worker_init_controller (struct rspamd_worker *worker,
 				ctx->resolver, worker, RSPAMD_MAP_WATCH_SCANNER);
 	}
 }
+
+gdouble
+rspamd_worker_check_and_adjust_timeout (struct rspamd_config *cfg, gdouble timeout)
+{
+	if (isnan (timeout)) {
+		/* Use implicit timeout from cfg->task_timeout */
+		timeout = cfg->task_timeout;
+	}
+
+	if (isnan (timeout)) {
+		return timeout;
+	}
+
+	struct rspamd_symcache_timeout_result *tres = rspamd_symcache_get_max_timeout (cfg->cache);
+	g_assert (tres != 0);
+
+	if (tres->max_timeout > timeout) {
+		msg_info_config("configured task_timeout %.2f is less than maximum symbols cache timeout %.2f, so"
+						"some symbols could be terminated early", timeout, tres->max_timeout);
+		/* TODO: list timeouts for top symbols */
+	}
+
+	rspamd_symcache_timeout_result_free (tres);
+
+	/* TODO: maybe adjust timeout */
+	return timeout;
+}
diff --git a/src/libserver/worker_util.h b/src/libserver/worker_util.h
index ed7ada3d3..677f8fcbf 100644
--- a/src/libserver/worker_util.h
+++ b/src/libserver/worker_util.h
@@ -173,6 +173,15 @@ void rspamd_hard_terminate (struct rspamd_main *rspamd_main) G_GNUC_NORETURN;
  */
 gboolean rspamd_worker_is_scanner (struct rspamd_worker *w);
 
+/**
+ * Checks
+ * @param cfg
+ * @param timeout
+ * @return
+ */
+gdouble rspamd_worker_check_and_adjust_timeout (struct rspamd_config *cfg,
+		gdouble timeout);
+
 /**
  * Returns TRUE if a specific worker is a primary controller
  * @param w
diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c
index 3bc5062c8..3bf83b0e4 100644
--- a/src/rspamd_proxy.c
+++ b/src/rspamd_proxy.c
@@ -157,6 +157,7 @@ struct rspamd_proxy_ctx {
 	struct rspamd_milter_context milter_ctx;
 	/* Language detector */
 	struct rspamd_lang_detector *lang_det;
+	gdouble task_timeout;
 };
 
 enum rspamd_backend_flags {
@@ -1886,7 +1887,7 @@ rspamd_proxy_self_scan (struct rspamd_proxy_session *session)
 
 	}
 	else if (session->ctx->has_self_scan) {
-		if (session->ctx->cfg->task_timeout > 0) {
+		if (!isnan(session->ctx->task_timeout) && session->ctx->task_timeout > 0) {
 			task->timeout_ev.data = task;
 			ev_timer_init (&task->timeout_ev, rspamd_task_timeout,
 					session->ctx->cfg->task_timeout,
@@ -2374,6 +2375,8 @@ start_rspamd_proxy (struct rspamd_worker *worker)
 		/* Additional initialisation needed */
 		rspamd_worker_init_scanner (worker, ctx->event_loop, ctx->resolver,
 				&ctx->lang_det);
+		/* Always yse cfg->task_timeout */
+		ctx->task_timeout = rspamd_worker_check_and_adjust_timeout(ctx->cfg, NAN);
 
 		if (worker->index == 0) {
 			/*
diff --git a/src/worker.c b/src/worker.c
index 87b63bbdf..26cdc2390 100644
--- a/src/worker.c
+++ b/src/worker.c
@@ -186,7 +186,7 @@ rspamd_worker_body_handler (struct rspamd_http_connection *conn,
 	}
 
 	/* Set global timeout for the task */
-	if (ctx->task_timeout > 0.0) {
+	if (!isnan(ctx->task_timeout) && ctx->task_timeout > 0.0) {
 		task->timeout_ev.data = task;
 		ev_timer_init (&task->timeout_ev, rspamd_task_timeout,
 				ctx->task_timeout,
@@ -493,14 +493,7 @@ start_worker (struct rspamd_worker *worker)
 	rspamd_symcache_start_refresh (worker->srv->cfg->cache, ctx->event_loop,
 			worker);
 
-	if (isnan (ctx->task_timeout)) {
-		if (isnan (ctx->cfg->task_timeout)) {
-			ctx->task_timeout = 0;
-		}
-		else {
-			ctx->task_timeout = ctx->cfg->task_timeout;
-		}
-	}
+	ctx->task_timeout = rspamd_worker_check_and_adjust_timeout(ctx->cfg, ctx->task_timeout);
 
 	ctx->resolver = rspamd_dns_resolver_init (worker->srv->logger,
 			ctx->event_loop,
-- 
2.39.5