From 3a7f4ef0ed9fb2583387c0fbcc7fc28ab403b3bc Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Thu, 7 Dec 2023 13:56:03 +0000 Subject: [PATCH] [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; } -- 2.39.5