peilaus alkaen
https://github.com/rspamd/rspamd.git
synced 2024-07-29 20:17:47 +02:00
[Feature] Add framework to manage Redis scripts
This commit is contained in:
vanhempi
9d99bef124
commit
d7f6e6a516
@ -606,4 +606,161 @@ end
|
||||
exports.rspamd_redis_make_request_taskless = redis_make_request_taskless
|
||||
exports.redis_make_request_taskless = redis_make_request_taskless
|
||||
|
||||
local redis_scripts = {
|
||||
}
|
||||
|
||||
local function load_redis_script(script, cfg, ev_base, _)
|
||||
local function merge_tables(t1, t2)
|
||||
for k,v in pairs(t2) do t1[k] = v end
|
||||
end
|
||||
|
||||
local function set_loaded()
|
||||
if script.sha then
|
||||
script.loaded = true
|
||||
end
|
||||
|
||||
local wait_table = {}
|
||||
for _,s in ipairs(script.waitq) do
|
||||
table.insert(wait_table, s)
|
||||
end
|
||||
|
||||
script.waitq = {}
|
||||
|
||||
for _,s in ipairs(wait_table) do
|
||||
s(script)
|
||||
end
|
||||
end
|
||||
local servers = {}
|
||||
|
||||
if script.redis_params.read_servers then
|
||||
merge_tables(servers, script.redis_params.read_servers:all_upstreams())
|
||||
end
|
||||
if script.redis_params.write_servers then
|
||||
merge_tables(servers, script.redis_params.write_servers:all_upstreams())
|
||||
end
|
||||
|
||||
-- Call load script on each server, set loaded flag
|
||||
script.in_flight = #servers
|
||||
for _,s in ipairs(servers) do
|
||||
local function script_cb(err, data)
|
||||
if err then
|
||||
s:fail()
|
||||
else
|
||||
s:ok()
|
||||
script.sha = data -- We assume that sha is the same on all servers
|
||||
end
|
||||
script.in_flight = script.in_flight - 1
|
||||
|
||||
if script.in_flight == 0 then
|
||||
set_loaded(script)
|
||||
end
|
||||
end
|
||||
|
||||
local rspamd_redis = require "rspamd_redis"
|
||||
|
||||
local options = {
|
||||
ev_base = ev_base,
|
||||
config = cfg,
|
||||
callback = script_cb,
|
||||
host = s:get_addr(),
|
||||
timeout = script.redis_params['timeout'],
|
||||
cmd = 'SCRIPT',
|
||||
args = {'LOAD', script.script}
|
||||
}
|
||||
|
||||
if script.redis_params['password'] then
|
||||
options['password'] = script.redis_params['password']
|
||||
end
|
||||
|
||||
if script.redis_params['db'] then
|
||||
options['dbname'] = script.redis_params['db']
|
||||
end
|
||||
|
||||
local ret = rspamd_redis.make_request(options)
|
||||
if not ret then
|
||||
logger.errx('cannot execute redis request to load script')
|
||||
script.in_flight = script.in_flight - 1
|
||||
end
|
||||
end
|
||||
|
||||
if script.in_flight == 0 then
|
||||
set_loaded(script)
|
||||
end
|
||||
end
|
||||
|
||||
local function add_redis_script(script, redis_params)
|
||||
local new_script = {
|
||||
loaded = false,
|
||||
redis_params = redis_params,
|
||||
script = script,
|
||||
waitq = {}, -- callbacks pending for script being loaded
|
||||
id = #redis_scripts + 1
|
||||
}
|
||||
|
||||
-- Register on load function
|
||||
rspamd_config:add_on_load(function(cfg, ev_base, worker)
|
||||
load_redis_script(new_script, cfg, ev_base, worker)
|
||||
end)
|
||||
|
||||
table.insert(redis_scripts, new_script)
|
||||
|
||||
return #redis_scripts
|
||||
end
|
||||
exports.add_redis_script = add_redis_script
|
||||
|
||||
local function exec_redis_script(id, params, callback, args)
|
||||
if not redis_scripts[id] then
|
||||
return false
|
||||
end
|
||||
|
||||
local script = redis_scripts[id]
|
||||
|
||||
local function do_call()
|
||||
local function redis_cb(err, data)
|
||||
if not err then
|
||||
callback(err, data)
|
||||
elseif err == 'NOSCRIPT' then
|
||||
-- Schedule restart
|
||||
table.insert(script.waitq, do_call)
|
||||
if script.in_flight ~= 0 then
|
||||
-- Reload scripts if this has not been initiated yet
|
||||
if params.task then
|
||||
load_redis_script(script, rspamd_config,
|
||||
params.task:get_ev_base(), nil)
|
||||
else
|
||||
load_redis_script(script, rspamd_config,
|
||||
params.ev_base, nil)
|
||||
end
|
||||
end
|
||||
else
|
||||
callback(err, data)
|
||||
end
|
||||
end
|
||||
|
||||
if params.task then
|
||||
if not rspamd_redis_make_request(params.task, script.redis_params,
|
||||
params.key, params.is_write, redis_cb, 'EVALSHA', args) then
|
||||
callback('Cannot make redis request', nil)
|
||||
end
|
||||
else
|
||||
if not redis_make_request_taskless(params.ev_base, rspamd_config,
|
||||
script.redis_params,
|
||||
params.key, params.is_write, redis_cb, 'EVALSHA', args) then
|
||||
callback('Cannot make redis request', nil)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
if not script.loaded then
|
||||
do_call()
|
||||
else
|
||||
-- Delayed until scripts are loaded
|
||||
table.insert(script.waitq, do_call)
|
||||
end
|
||||
|
||||
return true
|
||||
end
|
||||
|
||||
exports.exec_redis_script = exec_redis_script
|
||||
|
||||
return exports
|
||||
|
Ladataan…
Viittaa uudesa ongelmassa
Block a user