Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor(tracing): simplified the dynamic hooks code #12461

Merged
merged 1 commit into from
Mar 5, 2024
Merged

refactor(tracing): simplified the dynamic hooks code #12461

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion kong-3.7.0-0.rockspec
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -610,6 +610,5 @@ build = {
["kong.timing.hooks.socket"] = "kong/timing/hooks/socket.lua",

["kong.dynamic_hook"] = "kong/dynamic_hook/init.lua",
["kong.dynamic_hook.wrap_function_gen"] = "kong/dynamic_hook/wrap_function_gen.lua",
}
}
222 changes: 173 additions & 49 deletions kong/dynamic_hook/init.lua
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,25 @@
local warp_function_gen = require("kong.dynamic_hook.wrap_function_gen")

local ngx = ngx

local _M = {
local ngx = ngx
local type = type
local pcall = pcall
local select = select
local ipairs = ipairs
local assert = assert
local ngx_log = ngx.log
local ngx_WARN = ngx.WARN
local ngx_get_phase = ngx.get_phase


local _M = {
TYPE = {
BEFORE = 1,
AFTER = 2,
AFTER = 2,
BEFORE_MUT = 3,
AFTER_MUT = 4,
AFTER_MUT = 4,
},
}

local pcall = pcall

local non_function_hooks = {
local NON_FUNCTION_HOOKS = {
--[[
[group_name] = {
[hook_name] = <function>,
Expand All @@ -23,39 +29,152 @@ local non_function_hooks = {
--]]
}

local always_enabled_groups = {}

local wrap_functions = {
[0] = warp_function_gen.generate_wrap_function(0),
[1] = warp_function_gen.generate_wrap_function(1),
[2] = warp_function_gen.generate_wrap_function(2),
[3] = warp_function_gen.generate_wrap_function(3),
[4] = warp_function_gen.generate_wrap_function(4),
[5] = warp_function_gen.generate_wrap_function(5),
[6] = warp_function_gen.generate_wrap_function(6),
[7] = warp_function_gen.generate_wrap_function(7),
[8] = warp_function_gen.generate_wrap_function(8),
["varargs"] = warp_function_gen.generate_wrap_function("varargs"),
}

local ALWAYS_ENABLED_GROUPS = {}


local function should_execute_original_func(group_name)
if ALWAYS_ENABLED_GROUPS[group_name] then
return
end

local phase = ngx_get_phase()
if phase == "init" or phase == "init_worker" then
return true
end

local dynamic_hook = ngx.ctx.dynamic_hook
if not dynamic_hook then
return true
end

local enabled_groups = dynamic_hook.enabled_groups
if not enabled_groups[group_name] then
return true
end
end


local function execute_hook_vararg(hook, hook_type, group_name, ...)
if not hook then
return
end
local ok, err = pcall(hook, ...)
if not ok then
ngx_log(ngx_WARN, "failed to run ", hook_type, " hook of ", group_name, ": ", err)
end
end


local function execute_hooks_vararg(hooks, hook_type, group_name, ...)
if not hooks then
return
end
for _, hook in ipairs(hooks) do
execute_hook_vararg(hook, hook_type, group_name, ...)
end
end


local function execute_after_hooks_vararg(handlers, group_name, ...)
execute_hook_vararg(handlers.after_mut, "after_mut", group_name, ...)
execute_hooks_vararg(handlers.afters, "after", group_name, ...)
return ...
end


local function wrap_function_vararg(group_name, original_func, handlers)
return function (...)
if should_execute_original_func(group_name) then
return original_func(...)
end
execute_hooks_vararg(handlers.befores, "before", group_name, ...)
return execute_after_hooks_vararg(handlers, group_name, original_func(...))
end
end


local function execute_hook(hook, hook_type, group_name, a1, a2, a3, a4, a5, a6, a7, a8)
if not hook then
bungle marked this conversation as resolved.
Show resolved Hide resolved
return
end
local ok, err = pcall(hook, a1, a2, a3, a4, a5, a6, a7, a8)
if not ok then
ngx_log(ngx_WARN, "failed to run ", hook_type, " hook of ", group_name, ": ", err)
end
end


local function execute_hooks(hooks, hook_type, group_name, a1, a2, a3, a4, a5, a6, a7, a8)
if not hooks then
return
end
for _, hook in ipairs(hooks) do
execute_hook(hook, hook_type, group_name, a1, a2, a3, a4, a5, a6, a7, a8)
end
end


local function execute_original_func(max_args, original_func, a1, a2, a3, a4, a5, a6, a7, a8)
if max_args == 0 then
return original_func()
elseif max_args == 1 then
return original_func(a1)
elseif max_args == 2 then
return original_func(a1, a2)
elseif max_args == 3 then
return original_func(a1, a2, a3)
elseif max_args == 4 then
return original_func(a1, a2, a3, a4)
elseif max_args == 5 then
return original_func(a1, a2, a3, a4, a5)
elseif max_args == 6 then
return original_func(a1, a2, a3, a4, a5, a6)
elseif max_args == 7 then
return original_func(a1, a2, a3, a4, a5, a6, a7)
else
return original_func(a1, a2, a3, a4, a5, a6, a7, a8)
end
end


local function wrap_function(max_args, group_name, original_func, handlers)
return function(a1, a2, a3, a4, a5, a6, a7, a8)
if should_execute_original_func(group_name) then
a1, a2, a3, a4, a5, a6, a7, a8 = execute_original_func(max_args, original_func, a1, a2, a3, a4, a5, a6, a7, a8)

else
execute_hook(handlers.before_mut, "before_mut", group_name, a1, a2, a3, a4, a5, a6, a7, a8)
execute_hooks(handlers.befores, "before", group_name, a1, a2, a3, a4, a5, a6, a7, a8)
a1, a2, a3, a4, a5, a6, a7, a8 = execute_original_func(max_args, original_func, a1, a2, a3, a4, a5, a6, a7, a8)
execute_hook(handlers.after_mut, "after_mut", group_name, a1, a2, a3, a4, a5, a6, a7, a8)
execute_hooks(handlers.afters, "after", group_name, a1, a2, a3, a4, a5, a6, a7, a8)
end
return a1, a2, a3, a4, a5, a6, a7, a8
end
end


function _M.hook_function(group_name, parent, child_key, max_args, handlers)
assert(type(parent) == "table", "parent must be a table")
assert(type(child_key) == "string", "child_key must be a string")

if type(max_args) == "string" then
assert(max_args == "varargs", "max_args must be a number or \"varargs\"")
local is_varargs = max_args == "varargs"
if is_varargs then
assert(handlers.before_mut == nil, "before_mut is not supported for varargs functions")

else
assert(type(max_args) == "number", "max_args must be a number or \"varargs\"")
assert(max_args >= 0 and max_args <= 8, "max_args must be >= 0")
assert(type(max_args) == "number", 'max_args must be a number or "varargs"')
assert(max_args >= 0 and max_args <= 8, 'max_args must be >= 0 and <= 8, or "varargs"')
end

local old_func = parent[child_key]
assert(type(old_func) == "function", "parent[" .. child_key .. "] must be a function")
local original_func = parent[child_key]
assert(type(original_func) == "function", "parent[" .. child_key .. "] must be a function")

parent[child_key] = wrap_functions[max_args](always_enabled_groups, group_name, old_func, handlers)
if is_varargs then
parent[child_key] = wrap_function_vararg(group_name, original_func, handlers)
else
parent[child_key] = wrap_function(max_args, group_name, original_func, handlers)
end
end


Expand All @@ -64,18 +183,18 @@ function _M.hook(group_name, hook_name, handler)
assert(type(hook_name) == "string", "hook_name must be a string")
assert(type(handler) == "function", "handler must be a function")

local hooks = non_function_hooks[group_name]
local hooks = NON_FUNCTION_HOOKS[group_name]
if not hooks then
hooks = {}
non_function_hooks[group_name] = hooks
NON_FUNCTION_HOOKS[group_name] = hooks
end

hooks[hook_name] = handler
end


function _M.is_group_enabled(group_name)
if always_enabled_groups[group_name] then
if ALWAYS_ENABLED_GROUPS[group_name] then
return true
end

Expand All @@ -93,12 +212,12 @@ function _M.is_group_enabled(group_name)
end


function _M.run_hooks(ctx, group_name, hook_name, ...)
function _M.run_hooks(group_name, hook_name, a1, a2, a3, a4, a5, a6, a7, a8, ...)
if not _M.is_group_enabled(group_name) then
return
end

local hooks = non_function_hooks[group_name]
local hooks = NON_FUNCTION_HOOKS[group_name]
if not hooks then
return
end
Expand All @@ -108,30 +227,35 @@ function _M.run_hooks(ctx, group_name, hook_name, ...)
return
end

local ok, err = pcall(handler, ...)
local argc = select("#", ...)
local ok, err
if argc == 0 then
ok, err = pcall(handler, a1, a2, a3, a4, a5, a6, a7, a8)
else
ok, err = pcall(handler, a1, a2, a3, a4, a5, a6, a7, a8, ...)
end
if not ok then
ngx.log(ngx.WARN,
string.format("failed to run dynamic hook %s.%s: %s",
group_name, hook_name, err))
ngx_log(ngx_WARN, "failed to run dynamic hook ", group_name, ".", hook_name, ": ", err)
end
end


function _M.enable_on_this_request(group_name)
local info = ngx.ctx.dynamic_hook
if not info then
info = {
enabled_groups = {},
function _M.enable_on_this_request(group_name, ngx_ctx)
ngx_ctx = ngx_ctx or ngx.ctx
if ngx_ctx.dynamic_hook then
ngx_ctx.dynamic_hook.enabled_groups[group_name] = true
else
ngx_ctx.dynamic_hook = {
enabled_groups = {
[group_name] = true
},
}
ngx.ctx.dynamic_hook = info
end

info.enabled_groups[group_name] = true
end


function _M.always_enable(group_name)
always_enabled_groups[group_name] = true
ALWAYS_ENABLED_GROUPS[group_name] = true
end


Expand Down
Loading
Loading