commit 266daff: [Minor] Neural: Extract lua scripts
Vsevolod Stakhov
vsevolod at rspamd.com
Sat Mar 25 13:56:05 UTC 2023
Author: Vsevolod Stakhov
Date: 2023-03-25 13:49:16 +0000
URL: https://github.com/rspamd/rspamd/commit/266daff34bfd4c7bd6548098ca98ed0e6289488a (HEAD -> master)
[Minor] Neural: Extract lua scripts
---
.luacheckrc | 1 +
lualib/plugins/neural.lua | 109 ++---------------------
lualib/redis_scripts/neural_maybe_invalidate.lua | 25 ++++++
lualib/redis_scripts/neural_maybe_lock.lua | 19 ++++
lualib/redis_scripts/neural_save_unlock.lua | 24 +++++
lualib/redis_scripts/neural_train_size.lua | 20 +++++
6 files changed, 98 insertions(+), 100 deletions(-)
diff --git a/.luacheckrc b/.luacheckrc
index d5a18cc3b..c242bddd8 100644
--- a/.luacheckrc
+++ b/.luacheckrc
@@ -63,6 +63,7 @@ files['/**/lualib/lua_redis.lua'].globals = {
files['/**/lualib/redis_scripts/**'].globals = {
'redis',
'KEYS',
+ 'cjson',
}
files['/**/src/rspamadm/*'].globals = {
diff --git a/lualib/plugins/neural.lua b/lualib/plugins/neural.lua
index 3400f8d27..05dace489 100644
--- a/lualib/plugins/neural.lua
+++ b/lualib/plugins/neural.lua
@@ -96,113 +96,22 @@ local module_config = rspamd_config:get_all_opt(N)
settings = lua_util.override_defaults(settings, module_config)
local redis_params = lua_redis.parse_redis_server('neural')
--- Lua script that checks if we can store a new training vector
--- Uses the following keys:
--- key1 - ann key
--- returns nspam,nham (or nil if locked)
-local redis_lua_script_vectors_len = [[
- local prefix = KEYS[1]
- local locked = redis.call('HGET', prefix, 'lock')
- if locked then
- local host = redis.call('HGET', prefix, 'hostname') or 'unknown'
- return string.format('%s:%s', host, locked)
- end
- local nspam = 0
- local nham = 0
-
- local ret = redis.call('SCARD', prefix .. '_spam_set')
- if ret then nspam = tonumber(ret) end
- ret = redis.call('SCARD', prefix .. '_ham_set')
- if ret then nham = tonumber(ret) end
-
- return {nspam,nham}
-]]
-
--- Lua script to invalidate ANNs by rank
--- Uses the following keys
--- key1 - prefix for keys
--- key2 - number of elements to leave
-local redis_lua_script_maybe_invalidate = [[
- local card = redis.call('ZCARD', KEYS[1])
- local lim = tonumber(KEYS[2])
- if card > lim then
- local to_delete = redis.call('ZRANGE', KEYS[1], 0, card - lim - 1)
- if to_delete then
- for _,k in ipairs(to_delete) do
- local tb = cjson.decode(k)
- if type(tb) == 'table' and type(tb.redis_key) == 'string' then
- redis.call('DEL', tb.redis_key)
- -- Also train vectors
- redis.call('DEL', tb.redis_key .. '_spam_set')
- redis.call('DEL', tb.redis_key .. '_ham_set')
- end
- end
- end
- redis.call('ZREMRANGEBYRANK', KEYS[1], 0, card - lim - 1)
- return to_delete
- else
- return {}
- end
-]]
-
--- Lua script to invalidate ANN from redis
--- Uses the following keys
--- key1 - prefix for keys
--- key2 - current time
--- key3 - key expire
--- key4 - hostname
-local redis_lua_script_maybe_lock = [[
- local locked = redis.call('HGET', KEYS[1], 'lock')
- local now = tonumber(KEYS[2])
- if locked then
- locked = tonumber(locked)
- local expire = tonumber(KEYS[3])
- if now > locked and (now - locked) < expire then
- return {tostring(locked), redis.call('HGET', KEYS[1], 'hostname') or 'unknown'}
- end
- end
- redis.call('HSET', KEYS[1], 'lock', tostring(now))
- redis.call('HSET', KEYS[1], 'hostname', KEYS[4])
- return 1
-]]
-
--- Lua script to save and unlock ANN in redis
--- Uses the following keys
--- key1 - prefix for ANN
--- key2 - prefix for profile
--- key3 - compressed ANN
--- key4 - profile as JSON
--- key5 - expire in seconds
--- key6 - current time
--- key7 - old key
--- key8 - ROC Thresholds
--- key9 - optional PCA
-local redis_lua_script_save_unlock = [[
- local now = tonumber(KEYS[6])
- redis.call('ZADD', KEYS[2], now, KEYS[4])
- redis.call('HSET', KEYS[1], 'ann', KEYS[3])
- redis.call('DEL', KEYS[1] .. '_spam_set')
- redis.call('DEL', KEYS[1] .. '_ham_set')
- redis.call('HDEL', KEYS[1], 'lock')
- redis.call('HDEL', KEYS[7], 'lock')
- redis.call('EXPIRE', KEYS[1], tonumber(KEYS[5]))
- redis.call('HSET', KEYS[1], 'roc_thresholds', KEYS[8])
- if KEYS[9] then
- redis.call('HSET', KEYS[1], 'pca', KEYS[9])
- end
- return 1
-]]
+
+local redis_lua_script_vectors_len = "neural_train_size.lua"
+local redis_lua_script_maybe_invalidate = "neural_maybe_invalidate.lua"
+local redis_lua_script_maybe_lock = "neural_maybe_lock.lua"
+local redis_lua_script_save_unlock = "neural_save_unlock.lua"
local redis_script_id = {}
local function load_scripts()
- redis_script_id.vectors_len = lua_redis.add_redis_script(redis_lua_script_vectors_len,
+ redis_script_id.vectors_len = lua_redis.load_redis_script_from_file(redis_lua_script_vectors_len,
redis_params)
- redis_script_id.maybe_invalidate = lua_redis.add_redis_script(redis_lua_script_maybe_invalidate,
+ redis_script_id.maybe_invalidate = lua_redis.load_redis_script_from_file(redis_lua_script_maybe_invalidate,
redis_params)
- redis_script_id.maybe_lock = lua_redis.add_redis_script(redis_lua_script_maybe_lock,
+ redis_script_id.maybe_lock = lua_redis.load_redis_script_from_file(redis_lua_script_maybe_lock,
redis_params)
- redis_script_id.save_unlock = lua_redis.add_redis_script(redis_lua_script_save_unlock,
+ redis_script_id.save_unlock = lua_redis.load_redis_script_from_file(redis_lua_script_save_unlock,
redis_params)
end
diff --git a/lualib/redis_scripts/neural_maybe_invalidate.lua b/lualib/redis_scripts/neural_maybe_invalidate.lua
new file mode 100644
index 000000000..c54871717
--- /dev/null
+++ b/lualib/redis_scripts/neural_maybe_invalidate.lua
@@ -0,0 +1,25 @@
+-- Lua script to invalidate ANNs by rank
+-- Uses the following keys
+-- key1 - prefix for keys
+-- key2 - number of elements to leave
+
+local card = redis.call('ZCARD', KEYS[1])
+local lim = tonumber(KEYS[2])
+if card > lim then
+ local to_delete = redis.call('ZRANGE', KEYS[1], 0, card - lim - 1)
+ if to_delete then
+ for _,k in ipairs(to_delete) do
+ local tb = cjson.decode(k)
+ if type(tb) == 'table' and type(tb.redis_key) == 'string' then
+ redis.call('DEL', tb.redis_key)
+ -- Also train vectors
+ redis.call('DEL', tb.redis_key .. '_spam_set')
+ redis.call('DEL', tb.redis_key .. '_ham_set')
+ end
+ end
+ end
+ redis.call('ZREMRANGEBYRANK', KEYS[1], 0, card - lim - 1)
+ return to_delete
+else
+ return {}
+end
\ No newline at end of file
diff --git a/lualib/redis_scripts/neural_maybe_lock.lua b/lualib/redis_scripts/neural_maybe_lock.lua
new file mode 100644
index 000000000..7b5c6a60f
--- /dev/null
+++ b/lualib/redis_scripts/neural_maybe_lock.lua
@@ -0,0 +1,19 @@
+-- Lua script lock ANN for learning
+-- Uses the following keys
+-- key1 - prefix for keys
+-- key2 - current time
+-- key3 - key expire
+-- key4 - hostname
+
+local locked = redis.call('HGET', KEYS[1], 'lock')
+local now = tonumber(KEYS[2])
+if locked then
+ locked = tonumber(locked)
+ local expire = tonumber(KEYS[3])
+ if now > locked and (now - locked) < expire then
+ return {tostring(locked), redis.call('HGET', KEYS[1], 'hostname') or 'unknown'}
+ end
+end
+redis.call('HSET', KEYS[1], 'lock', tostring(now))
+redis.call('HSET', KEYS[1], 'hostname', KEYS[4])
+return 1
\ No newline at end of file
diff --git a/lualib/redis_scripts/neural_save_unlock.lua b/lualib/redis_scripts/neural_save_unlock.lua
new file mode 100644
index 000000000..5af1ddcde
--- /dev/null
+++ b/lualib/redis_scripts/neural_save_unlock.lua
@@ -0,0 +1,24 @@
+-- Lua script to save and unlock ANN in redis
+-- Uses the following keys
+-- key1 - prefix for ANN
+-- key2 - prefix for profile
+-- key3 - compressed ANN
+-- key4 - profile as JSON
+-- key5 - expire in seconds
+-- key6 - current time
+-- key7 - old key
+-- key8 - ROC Thresholds
+-- key9 - optional PCA
+local now = tonumber(KEYS[6])
+redis.call('ZADD', KEYS[2], now, KEYS[4])
+redis.call('HSET', KEYS[1], 'ann', KEYS[3])
+redis.call('DEL', KEYS[1] .. '_spam_set')
+redis.call('DEL', KEYS[1] .. '_ham_set')
+redis.call('HDEL', KEYS[1], 'lock')
+redis.call('HDEL', KEYS[7], 'lock')
+redis.call('EXPIRE', KEYS[1], tonumber(KEYS[5]))
+redis.call('HSET', KEYS[1], 'roc_thresholds', KEYS[8])
+if KEYS[9] then
+ redis.call('HSET', KEYS[1], 'pca', KEYS[9])
+end
+return 1
\ No newline at end of file
diff --git a/lualib/redis_scripts/neural_train_size.lua b/lualib/redis_scripts/neural_train_size.lua
new file mode 100644
index 000000000..5a00ae3fc
--- /dev/null
+++ b/lualib/redis_scripts/neural_train_size.lua
@@ -0,0 +1,20 @@
+-- Lua script that checks if we can store a new training vector
+-- Uses the following keys:
+-- key1 - ann key
+-- returns nspam,nham (or nil if locked)
+
+local prefix = KEYS[1]
+local locked = redis.call('HGET', prefix, 'lock')
+if locked then
+ local host = redis.call('HGET', prefix, 'hostname') or 'unknown'
+ return string.format('%s:%s', host, locked)
+end
+local nspam = 0
+local nham = 0
+
+local ret = redis.call('SCARD', prefix .. '_spam_set')
+if ret then nspam = tonumber(ret) end
+ret = redis.call('SCARD', prefix .. '_ham_set')
+if ret then nham = tonumber(ret) end
+
+return {nspam,nham}
\ No newline at end of file
More information about the Commits
mailing list