commit 3a7f4ef: [Minor] Add initial stat callback

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2023-12-07 13:56:03 +0000
URL: https://github.com/rspamd/rspamd/commit/3a7f4ef0ed9fb2583387c0fbcc7fc28ab403b3bc

[Minor] Add initial stat callback

---
 lualib/lua_bayes_redis.lua             | 22 +++++++++++--
 src/libstat/backends/redis_backend.cxx | 60 +++++++++++++++-------------------
 2 files changed, 47 insertions(+), 35 deletions(-)

diff --git a/lualib/lua_bayes_redis.lua b/lualib/lua_bayes_redis.lua
index 25c56d58b..575beff4b 100644
--- a/lualib/lua_bayes_redis.lua
+++ b/lualib/lua_bayes_redis.lua
@@ -21,7 +21,7 @@ local lua_redis = require "lua_redis"
 local logger = require "rspamd_logger"
 local lua_util = require "lua_util"
 
-local N = "stat_redis"
+local N = "bayes"
 
 local function gen_classify_functor(redis_params, classify_script_id)
   return function(task, expanded_key, id, is_spam, stat_tokens, callback)
@@ -52,7 +52,7 @@ end
 --- @param classifier_ucl ucl of the classifier config
 --- @param statfile_ucl ucl of the statfile config
 --- @return a pair of (classify_functor, learn_functor) or `nil` in case of error
-exports.lua_bayes_init_classifier = function(classifier_ucl, statfile_ucl)
+exports.lua_bayes_init_classifier = function(classifier_ucl, statfile_ucl, symbol, stat_periodic_cb)
   local redis_params
 
   if classifier_ucl.backend then
@@ -78,6 +78,24 @@ exports.lua_bayes_init_classifier = function(classifier_ucl, statfile_ucl)
 
   local classify_script_id = lua_redis.load_redis_script_from_file("bayes_classify.lua", redis_params)
   local learn_script_id = lua_redis.load_redis_script_from_file("bayes_learn.lua", redis_params)
+  local stat_script_id = lua_redis.load_redis_script_from_file("bayes_stat.lua", redis_params)
+  local max_users = classifier_ucl.max_users or 1000
+
+  rspamd_config:add_on_load(function(_, ev_base, _)
+
+    rspamd_config:add_periodic(ev_base, 0.0, function(cfg, _)
+
+      local function stat_redis_cb(err, data)
+        -- TODO: write this function
+
+      end
+
+      lua_redis.exec_redis_script(stat_script_id,
+          { ev_base = ev_base, cfg = cfg, is_write = false },
+          stat_redis_cb, { symbol, max_users })
+      return 30.0 -- TODO: make configurable
+    end)
+  end)
 
   return gen_classify_functor(redis_params, classify_script_id), gen_learn_functor(redis_params, learn_script_id)
 end
diff --git a/src/libstat/backends/redis_backend.cxx b/src/libstat/backends/redis_backend.cxx
index 1ff67bcb4..fe38b52bc 100644
--- a/src/libstat/backends/redis_backend.cxx
+++ b/src/libstat/backends/redis_backend.cxx
@@ -43,7 +43,6 @@ INIT_LOG_MODULE(stat_redis)
 struct redis_stat_ctx {
 	lua_State *L;
 	struct rspamd_statfile_config *stcf;
-	struct rspamd_stat_async_elt *stat_elt;
 	const char *redis_object = REDIS_DEFAULT_OBJECT;
 	bool enable_users = false;
 	bool store_tokens = false;
@@ -145,27 +144,6 @@ public:
 	}
 };
 
-/* Used to get statistics from redis */
-struct rspamd_redis_stat_cbdata;
-
-struct rspamd_redis_stat_elt {
-	struct redis_stat_ctx *ctx;
-	struct rspamd_stat_async_elt *async;
-	struct ev_loop *event_loop;
-	ucl_object_t *stat;
-	struct rspamd_redis_stat_cbdata *cbdata;
-};
-
-struct rspamd_redis_stat_cbdata {
-	struct rspamd_redis_stat_elt *elt;
-	redisAsyncContext *redis;
-	ucl_object_t *cur;
-	GPtrArray *cur_keys;
-	struct upstream *selected;
-	guint inflight;
-	gboolean wanna_die;
-};
-
 #define GET_TASK_ELT(task, elt) (task == nullptr ? nullptr : (task)->elt)
 
 static const gchar *M = "redis statistics";
@@ -1132,6 +1110,21 @@ rspamd_redis_async_stat_fin(struct rspamd_stat_async_elt *elt, gpointer d)
 
 #endif
 
+static int
+rspamd_redis_stat_cb(lua_State *L)
+{
+	const auto *cookie = lua_tostring(L, lua_upvalueindex(1));
+	auto *cfg = lua_check_config(L, 1);
+	auto *backend = REDIS_CTX(rspamd_mempool_get_variable(cfg->cfg_pool, cookie));
+
+	if (backend == nullptr) {
+		msg_err("internal error: cookie %s is not found", cookie);
+
+		return 0;
+	}
+
+	return 0;
+}
 
 static bool
 rspamd_redis_parse_classifier_opts(struct redis_stat_ctx *backend,
@@ -1239,8 +1232,18 @@ rspamd_redis_parse_classifier_opts(struct redis_stat_ctx *backend,
 	/* Push arguments */
 	ucl_object_push_lua(L, classifier_obj, false);
 	ucl_object_push_lua(L, statfile_obj, false);
+	lua_pushstring(L, backend->stcf->symbol);
 
-	if (lua_pcall(L, 2, 2, err_idx) != 0) {
+	/* Store backend in random cookie */
+	char *cookie = (char *) rspamd_mempool_alloc(cfg->cfg_pool, 16);
+	rspamd_random_hex(cookie, 16);
+	cookie[15] = '\0';
+	rspamd_mempool_set_variable(cfg->cfg_pool, cookie, backend, nullptr);
+	/* Callback */
+	lua_pushstring(L, cookie);
+	lua_pushcclosure(L, &rspamd_redis_stat_cb, 1);
+
+	if (lua_pcall(L, 4, 2, err_idx) != 0) {
 		msg_err("call to lua_bayes_init_classifier "
 				"script failed: %s",
 				lua_tostring(L, -1));
@@ -1289,9 +1292,6 @@ rspamd_redis_init(struct rspamd_stat_ctx *ctx,
 	st->stcf->clcf->flags |= RSPAMD_FLAG_CLASSIFIER_INCREMENTING_BACKEND;
 	backend->stcf = st->stcf;
 
-	auto *st_elt = g_new0(struct rspamd_redis_stat_elt, 1);
-	st_elt->event_loop = ctx->event_loop;
-	st_elt->ctx = backend;
 #if 0
 	backend->stat_elt = rspamd_stat_ctx_register_async(
 		rspamd_redis_async_stat_cb,
@@ -1629,13 +1629,7 @@ rspamd_redis_get_stat(gpointer runtime,
 	struct rspamd_redis_stat_elt *st;
 	redisAsyncContext *redis;
 
-	if (rt->ctx->stat_elt) {
-		st = (struct rspamd_redis_stat_elt *) rt->ctx->stat_elt->ud;
-
-		if (st->stat) {
-			return ucl_object_ref(st->stat);
-		}
-	}
+	/* TODO: write extraction */
 
 	return nullptr;
 }