[Project] Add cache config

1 files changed, 21 insertions, 6 deletions

diff --git a/lualib/lua_bayes_redis.lua b/lualib/lua_bayes_redis.lua
index 6f6c27403..3988b937a 100644
--- a/lualib/lua_bayes_redis.lua
+++ b/lualib/lua_bayes_redis.lua
@@ -20,6 +20,7 @@ local exports = {}
 local lua_redis = require "lua_redis"
 local logger = require "rspamd_logger"
 local lua_util = require "lua_util"
+local ucl = require "ucl"
 
 local N = "bayes"
 
@@ -171,11 +172,12 @@ exports.lua_bayes_init_statfile = function(classifier_ucl, statfile_ucl, symbol,
   return gen_classify_functor(redis_params, classify_script_id), gen_learn_functor(redis_params, learn_script_id)
 end
 
-local function gen_cache_check_functor(redis_params, check_script_id)
+local function gen_cache_check_functor(redis_params, check_script_id, conf)
+  local packed_conf = ucl.to_format(conf, 'msgpack')
   return function(task, cache_id, callback)
 
     local function classify_redis_cb(err, data)
-      lua_util.debugm(N, task, 'classify redis cb: %s, %s', err, data)
+      lua_util.debugm(N, task, 'check cache redis cb: %s, %s', err, data)
       if err then
         callback(task, false, err)
       else
@@ -183,22 +185,25 @@ local function gen_cache_check_functor(redis_params, check_script_id)
       end
     end
 
+    lua_util.debugm(N, task, 'checking cache: %s', cache_id)
     lua_redis.exec_redis_script(check_script_id,
         { task = task, is_write = false, key = cache_id },
-        classify_redis_cb, { cache_id })
+        classify_redis_cb, { cache_id, packed_conf })
   end
 end
 
-local function gen_cache_learn_functor(redis_params, learn_script_id)
+local function gen_cache_learn_functor(redis_params, learn_script_id, conf)
+  local packed_conf = ucl.to_format(conf, 'msgpack')
   return function(task, cache_id, is_spam)
     local function learn_redis_cb(err, data)
       lua_util.debugm(N, task, 'learn_cache redis cb: %s, %s', err, data)
     end
 
+    lua_util.debugm(N, task, 'try to learn cache: %s', cache_id)
     lua_redis.exec_redis_script(learn_script_id,
         { task = task, is_write = true, key = cache_id },
         learn_redis_cb,
-        { cache_id, is_spam and "1" or "0" })
+        { cache_id, is_spam and "1" or "0", packed_conf })
 
   end
 end
@@ -210,10 +215,20 @@ exports.lua_bayes_init_cache = function(classifier_ucl, statfile_ucl)
     return nil
   end
 
+  local default_conf = {
+    cache_prefix = "learned_ids",
+    max_elt = 10000, -- Maximum number of elements in the cache key
+    max_keys = 10, -- Maximum number of keys in the cache
+    per_user_mult = 0.1, -- Multiplier for per user cache size
+  }
+
+  local conf = lua_util.override_defaults(default_conf, classifier_ucl)
+
   local check_script_id = lua_redis.load_redis_script_from_file("bayes_cache_check.lua", redis_params)
   local learn_script_id = lua_redis.load_redis_script_from_file("bayes_cache_learn.lua", redis_params)
 
-  return gen_cache_check_functor(redis_params, check_script_id), gen_cache_learn_functor(redis_params, learn_script_id)
+  return gen_cache_check_functor(redis_params, check_script_id, conf), gen_cache_learn_functor(redis_params,
+      learn_script_id, conf)
 end
 
 return exports