More changes to classification stages

2 files changed, 33 insertions, 29 deletions

diff --git a/src/libstat/stat_internal.h b/src/libstat/stat_internal.h
index c049486fe..640196788 100644
--- a/src/libstat/stat_internal.h
+++ b/src/libstat/stat_internal.h
@@ -53,6 +53,7 @@ struct rspamd_statfile_runtime {
 
 struct rspamd_classifier_runtime {
 	struct rspamd_classifier_config *clcf;
+	struct classifier_ctx *clctx;
 	struct rspamd_stat_classifier *cl;
 	struct rspamd_stat_backend *backend;
 	struct rspamd_tokenizer_runtime *tok;
diff --git a/src/libstat/stat_process.c b/src/libstat/stat_process.c
index 89df10a07..db1ed9276 100644
--- a/src/libstat/stat_process.c
+++ b/src/libstat/stat_process.c
@@ -546,7 +546,6 @@ rspamd_stat_classify (struct rspamd_task *task, lua_State *L, guint stage,
 	struct rspamd_stat_ctx *st_ctx;
 	struct rspamd_statfile_runtime *st_run;
 	struct rspamd_classifier_runtime *cl_run;
-	struct classifier_ctx *cl_ctx;
 	GList *cl_runtimes;
 	GList *cur, *curst;
 	gboolean ret = RSPAMD_STAT_PROCESS_OK;
@@ -564,19 +563,37 @@ rspamd_stat_classify (struct rspamd_task *task, lua_State *L, guint stage,
 			}
 
 			task->cl_runtimes = cl_runtimes;
+			cur = cl_runtimes;
+
+			/* Finalize backend so it can load tokens delayed if needed */
+			while (cur) {
+				cl_run = (struct rspamd_classifier_runtime *) cur->data;
+				curst = cl_run->st_runtime;
+
+				while (curst) {
+					st_run = curst->data;
+					cl_run->backend->finalize_process (task,
+							st_run->backend_runtime,
+							cl_run->backend->ctx);
+					curst = g_list_next (curst);
+				}
+			}
 		}
-
+	}
+	else if (stage == RSPAMD_TASK_STAGE_CLASSIFIERS) {
 		cur = cl_runtimes;
 
+		/* The first stage of classification */
 		while (cur) {
 			cl_run = (struct rspamd_classifier_runtime *) cur->data;
 			cl_run->stage = RSPAMD_STAT_STAGE_PRE;
 
 			if (cl_run->cl) {
-				cl_ctx = cl_run->cl->init_func (task->task_pool, cl_run->clcf);
+				cl_run->clctx = cl_run->cl->init_func (task->task_pool,
+						cl_run->clcf);
 
-				if (cl_ctx != NULL) {
-					cl_run->cl->classify_func (cl_ctx, cl_run->tok->tokens,
+				if (cl_run->clctx != NULL) {
+					cl_run->cl->classify_func (cl_run->clctx, cl_run->tok->tokens,
 							cl_run, task);
 				}
 			}
@@ -584,8 +601,9 @@ rspamd_stat_classify (struct rspamd_task *task, lua_State *L, guint stage,
 			cur = g_list_next (cur);
 		}
 	}
-	else if (stage == RSPAMD_TASK_STAGE_CLASSIFIERS) {
+	else if (stage == RSPAMD_TASK_STAGE_CLASSIFIERS_POST) {
 		cur = cl_runtimes;
+		/* The second stage of classification */
 		while (cur) {
 			cl_run = (struct rspamd_classifier_runtime *) cur->data;
 			cl_run->stage = RSPAMD_STAT_STAGE_POST;
@@ -595,21 +613,6 @@ rspamd_stat_classify (struct rspamd_task *task, lua_State *L, guint stage,
 				continue;
 			}
 
-			if (cl_run->cl) {
-				if (cl_ctx != NULL) {
-					if (cl_run->cl->classify_func (cl_ctx, cl_run->tok->tokens,
-							cl_run, task)) {
-						ret = RSPAMD_STAT_PROCESS_OK;
-					}
-				}
-			}
-
-			cur = g_list_next (cur);
-		}
-	}
-	else if (stage == RSPAMD_TASK_STAGE_CLASSIFIERS_POST) {
-		cur = cl_runtimes;
-		while (cur) {
 			cl_run = (struct rspamd_classifier_runtime *) cur->data;
 			cl_run->stage = RSPAMD_STAT_STAGE_POST;
 
@@ -618,14 +621,14 @@ rspamd_stat_classify (struct rspamd_task *task, lua_State *L, guint stage,
 				continue;
 			}
 
-			curst = cl_run->st_runtime;
-
-			while (curst) {
-				st_run = curst->data;
-				cl_run->backend->finalize_process (task,
-						st_run->backend_runtime,
-						cl_run->backend->ctx);
-				curst = g_list_next (curst);
+			if (cl_run->cl) {
+				if (cl_run->clctx != NULL) {
+					if (cl_run->cl->classify_func (cl_run->clctx,
+							cl_run->tok->tokens,
+							cl_run, task)) {
+						ret = RSPAMD_STAT_PROCESS_OK;
+					}
+				}
 			}
 
 			cur = g_list_next (cur);