[Project] Final things on redis cache rework

3 files changed, 89 insertions, 5 deletions

diff --git a/lualib/lua_bayes_redis.lua b/lualib/lua_bayes_redis.lua
index 3988b937a..576f88b8a 100644
--- a/lualib/lua_bayes_redis.lua
+++ b/lualib/lua_bayes_redis.lua
@@ -177,11 +177,15 @@ local function gen_cache_check_functor(redis_params, check_script_id, conf)
   return function(task, cache_id, callback)
 
     local function classify_redis_cb(err, data)
-      lua_util.debugm(N, task, 'check cache redis cb: %s, %s', err, data)
+      lua_util.debugm(N, task, 'check cache redis cb: %s, %s (%s)', err, data, type(data))
       if err then
         callback(task, false, err)
       else
-        callback(task, true, tonumber(data))
+        if type(data) == 'number' then
+          callback(task, true, data)
+        else
+          callback(task, false, 'not found')
+        end
       end
     end
 
@@ -217,12 +221,19 @@ exports.lua_bayes_init_cache = function(classifier_ucl, statfile_ucl)
 
   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
+    cache_max_elt = 10000, -- Maximum number of elements in the cache key
+    cache_max_keys = 5, -- Maximum number of keys in the cache
+    cache_per_user_mult = 0.1, -- Multiplier for per user cache size
+    cache_elt_len = 32, -- Length of the element in the cache (will trim id to that value)
   }
 
   local conf = lua_util.override_defaults(default_conf, classifier_ucl)
+  -- Clean all not known configurations
+  for k, _ in pairs(conf) do
+    if default_conf[k] == nil then
+      conf[k] = nil
+    end
+  end
 
   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)
diff --git a/lualib/redis_scripts/bayes_cache_check.lua b/lualib/redis_scripts/bayes_cache_check.lua
index e69de29bb..3bea57d5f 100644
--- a/lualib/redis_scripts/bayes_cache_check.lua
+++ b/lualib/redis_scripts/bayes_cache_check.lua
@@ -0,0 +1,20 @@
+-- Lua script to perform cache checking for bayes classification
+-- This script accepts the following parameters:
+-- key1 - cache id
+-- key2 - configuration table in message pack
+
+local cache_id = KEYS[1]
+local conf = cmsgpack.unpack(KEYS[2])
+cache_id = string.sub(cache_id, 1, conf.cache_prefix_len)
+
+-- Try each prefix that is in Redis
+for i = 0, conf.cache_max_keys do
+  local prefix = conf.cache_prefix .. string.rep("X", i)
+  local have = redis.call('HGET', prefix, cache_id)
+
+  if have then
+    return tonumber(have)
+  end
+end
+
+return nil
diff --git a/lualib/redis_scripts/bayes_cache_learn.lua b/lualib/redis_scripts/bayes_cache_learn.lua
index e69de29bb..5335ce832 100644
--- a/lualib/redis_scripts/bayes_cache_learn.lua
+++ b/lualib/redis_scripts/bayes_cache_learn.lua
@@ -0,0 +1,53 @@
+-- Lua script to perform cache checking for bayes classification
+-- This script accepts the following parameters:
+-- key1 - cache id
+-- key3 - is spam
+-- key3 - configuration table in message pack
+
+local cache_id = KEYS[1]
+local is_spam = KEYS[2]
+local conf = cmsgpack.unpack(KEYS[3])
+cache_id = string.sub(cache_id, 1, conf.cache_prefix_len)
+
+-- Try each prefix that is in Redis (as some other instance might have set it)
+for i = 0, conf.cache_max_keys do
+  local prefix = conf.cache_prefix .. string.rep("X", i)
+  local have = redis.call('HGET', prefix, cache_id)
+
+  if have then
+    -- Already in cache
+    return false
+  end
+end
+
+local added = false
+for i = 0, conf.cache_max_keys do
+  if not added then
+    local prefix = conf.cache_prefix .. string.rep("X", i)
+    local count = redis.call('HLEN', prefix)
+
+    if count < conf.cache_max_elt then
+      -- We can add it to this prefix
+      redis.call('HSET', prefix, cache_id, is_spam)
+      added = true
+    end
+  end
+end
+
+if not added then
+  -- Need to expire some keys
+  for i = 0, conf.cache_max_keys do
+    local prefix = conf.cache_prefix .. string.rep("X", i)
+    local exists = redis.call('EXISTS', prefix)
+
+    if exists then
+      redis.call('DEL', prefix)
+      redis.call('HSET', prefix, cache_id, is_spam)
+
+      -- Do not expire anything else
+      return true
+    end
+  end
+end
+
+return true
\ No newline at end of file