feat(ai-proxy-multi): add support for healthcheck (apache#12509)

Author: Revolyssup

apisix/healthcheck_manager.lua

@@ -26,6 +26,8 @@ local healthcheck
 local events = require("apisix.events")
 local tab_clone = core.table.clone
 local timer_every = ngx.timer.every
+local ngx_re      = require('ngx.re')
+local jp = require("jsonpath")
 local string_sub     = string.sub
 
 local _M = {}
@@ -58,8 +60,17 @@ local function remove_etcd_prefix(key)
     return string_sub(key, #prefix + 1)
 end
 
+local function parse_path(resource_full_path)
+    local resource_path_parts = ngx_re.split(resource_full_path, "#")
+    local resource_path = resource_path_parts[1] or resource_full_path
+    local resource_sub_path = resource_path_parts[2] or ""
+    return resource_path, resource_sub_path
+end
 
 local function fetch_latest_conf(resource_path)
+    -- if resource path contains json path, extract out the prefix
+    -- for eg: extracts /routes/1 from /routes/1#plugins.abc
+    resource_path = parse_path(resource_path)
     local resource_type, id
     -- Handle both formats:
     -- 1. /<etcd-prefix>/<resource_type>/<id>
@@ -206,6 +217,15 @@ function _M.upstream_version(index, nodes_ver)
 end
 
 
+local function get_plugin_name(path)
+    -- Extract JSON path (after '#') or use full path
+    local json_path = path:match("#(.+)$") or path
+    -- Match plugin name in the JSON path segment
+    return json_path:match("^plugins%['([^']+)'%]")
+        or json_path:match('^plugins%["([^"]+)"%]')
+        or json_path:match("^plugins%.([^%.]+)")
+end
+
 local function timer_create_checker()
     if core.table.nkeys(waiting_pool) == 0 then
         return
@@ -224,7 +244,21 @@ local function timer_create_checker()
             if not res_conf then
                 goto continue
             end
-            local upstream = res_conf.value.upstream or res_conf.value
+            local upstream
+            local plugin_name = get_plugin_name(resource_path)
+            if plugin_name and plugin_name ~= "" then
+                local _, sub_path = parse_path(resource_path)
+                local json_path = "$." .. sub_path
+                --- the users of the API pass the jsonpath(in resourcepath) to
+                --- upstream_constructor_config which is passed to the
+                --- callback construct_upstream to create an upstream dynamically
+                local upstream_constructor_config = jp.value(res_conf.value, json_path)
+                local plugin = require("apisix.plugins." .. plugin_name)
+                upstream = plugin.construct_upstream(upstream_constructor_config)
+                upstream.resource_key = resource_path
+            else
+                upstream = res_conf.value.upstream or res_conf.value
+            end
             local new_version = _M.upstream_version(res_conf.modifiedIndex, upstream._nodes_ver)
             core.log.info("checking waiting pool for resource: ", resource_path,
                     " current version: ", new_version, " requested version: ", resource_ver)

apisix/http/service.lua

@@ -17,6 +17,7 @@
 local core   = require("apisix.core")
 local apisix_upstream = require("apisix.upstream")
 local plugin_checker = require("apisix.plugin").plugin_checker
+local plugin = require("apisix.plugin")
 local services
 local error = error
 
@@ -46,6 +47,9 @@ local function filter(service)
         return
     end
 
+
+    plugin.set_plugins_meta_parent(service.value.plugins, service)
+
     apisix_upstream.filter_upstream(service.value.upstream, service)
 
     core.log.info("filter service: ", core.json.delay_encode(service, true))

apisix/plugin.lua

@@ -34,6 +34,8 @@ local type          = type
 local local_plugins = core.table.new(32, 0)
 local tostring      = tostring
 local error         = error
+local getmetatable  = getmetatable
+local setmetatable  = setmetatable
 -- make linter happy to avoid error: getting the Lua global "load"
 -- luacheck: globals load, ignore lua_load
 local lua_load          = load
@@ -1234,6 +1236,30 @@ function _M.run_plugin(phase, plugins, api_ctx)
     return api_ctx, plugin_run
 end
 
+function _M.set_plugins_meta_parent(plugins, parent)
+    if not plugins then
+        return
+    end
+    for _, plugin_conf in pairs(plugins) do
+        if not plugin_conf._meta then
+            plugin_conf._meta = {}
+        end
+        if not plugin_conf._meta.parent then
+            local parent_info = {
+                resource_key = parent.key,
+                resource_version = tostring(parent.modifiedIndex)
+            }
+            local mt_table = getmetatable(plugin_conf._meta)
+            if mt_table then
+                mt_table.parent = parent_info
+            else
+                plugin_conf._meta = setmetatable(plugin_conf._meta,
+                                                    { __index = {parent = parent_info} })
+            end
+        end
+    end
+end
+
 
 function _M.run_global_rules(api_ctx, global_rules, phase_name)
     if global_rules and #global_rules > 0 then

apisix/plugins/ai-proxy-multi.lua

@@ -19,6 +19,10 @@ local core = require("apisix.core")
 local schema = require("apisix.plugins.ai-proxy.schema")
 local base   = require("apisix.plugins.ai-proxy.base")
 local plugin = require("apisix.plugin")
+local ipmatcher  = require("resty.ipmatcher")
+local healthcheck_manager = require("apisix.healthcheck_manager")
+local tonumber = tonumber
+local pairs = pairs
 
 local require = require
 local pcall = pcall
@@ -118,17 +122,100 @@ local function transform_instances(new_instances, instance)
     new_instances[instance.priority][instance.name] = instance.weight
 end
 
+local function parse_domain_for_node(node)
+    local host = node.domain or node.host
+    if not ipmatcher.parse_ipv4(host)
+       and not ipmatcher.parse_ipv6(host)
+    then
+        node.domain = host
 
-local function create_server_picker(conf, ups_tab)
+        local ip, err = core.resolver.parse_domain(host)
+        if ip then
+            node.host = ip
+        end
+
+        if err then
+            core.log.error("dns resolver domain: ", host, " error: ", err)
+        end
+    end
+end
+
+
+local function resolve_endpoint(instance_conf)
+    local endpoint = core.table.try_read_attr(instance_conf, "override", "endpoint")
+    local scheme, host, port, _ = endpoint:match("^(https?)://([^:/]+):?(%d*)(/?.*)$")
+    if port == "" then
+        port = (scheme == "https") and "443" or "80"
+    end
+    local node = {
+        host = host,
+        port = tonumber(port),
+        scheme = scheme,
+    }
+    parse_domain_for_node(node)
+    return node
+end
+
+
+local function get_checkers_status_ver(checkers)
+    local status_ver_total = 0
+    for _, checker in pairs(checkers) do
+        status_ver_total = status_ver_total + checker.status_ver
+    end
+    return status_ver_total
+end
+
+
+
+local function fetch_health_instances(conf, checkers)
+    local instances = conf.instances
+    local new_instances = core.table.new(0, #instances)
+    if not checkers then
+        for _, ins in ipairs(conf.instances) do
+            transform_instances(new_instances, ins)
+        end
+        return new_instances
+    end
+
+    for _, ins in ipairs(instances) do
+        local checker = checkers[ins.name]
+        if checker then
+            local host = ins.checks and ins.checks.active and ins.checks.active.host
+            local port = ins.checks and ins.checks.active and ins.checks.active.port
+
+            local node = resolve_endpoint(ins)
+            local ok, err = checker:get_target_status(node.host, port or node.port, host)
+            if ok then
+                transform_instances(new_instances, ins)
+            elseif err then
+                core.log.error("failed to get health check target status, addr: ",
+                    node.host, ":", port or node.port, ", host: ", host, ", err: ", err)
+            end
+        else
+            transform_instances(new_instances, ins)
+        end
+    end
+
+    if core.table.nkeys(new_instances) == 0 then
+        core.log.warn("all upstream nodes is unhealthy, use default")
+        for _, ins in ipairs(instances) do
+            transform_instances(new_instances, ins)
+        end
+    end
+
+    return new_instances
+end
+
+
+local function create_server_picker(conf, ups_tab, checkers)
     local picker = pickers[conf.balancer.algorithm] -- nil check
     if not picker then
         pickers[conf.balancer.algorithm] = require("apisix.balancer." .. conf.balancer.algorithm)
         picker = pickers[conf.balancer.algorithm]
     end
-    local new_instances = {}
-    for _, ins in ipairs(conf.instances) do
-        transform_instances(new_instances, ins)
-    end
+
+    local new_instances = fetch_health_instances(conf, checkers)
+    core.log.info("fetch health instances: ", core.json.delay_encode(new_instances))
 
     if #new_instances._priority_index > 1 then
         core.log.info("new instances: ", core.json.delay_encode(new_instances))
@@ -149,11 +236,57 @@ local function get_instance_conf(instances, name)
 end
 
 
+function _M.construct_upstream(instance)
+    local upstream = {}
+    local node = resolve_endpoint(instance)
+    if not node then
+        return nil, "failed to resolve endpoint for instance: " .. instance.name
+    end
+
+    if not node.host or not node.port then
+        return nil, "invalid upstream node: " .. core.json.encode(node)
+    end
+
+    parse_domain_for_node(node)
+
+    local node = {
+        host = node.host,
+        port = node.port,
+        scheme = node.scheme,
+        weight = instance.weight or 1,
+        priority = instance.priority or 0,
+        name = instance.name,
+    }
+    upstream.nodes = {node}
+    upstream.checks = instance.checks
+    return upstream
+end
+
+
 local function pick_target(ctx, conf, ups_tab)
+    local checkers
+    for i, instance in ipairs(conf.instances) do
+        if instance.checks then
+            -- json path is 0 indexed so we need to decrement i
+            local resource_path = conf._meta.parent.resource_key ..
+                                  "#plugins['ai-proxy-multi'].instances[" .. i-1 .. "]"
+            local resource_version = conf._meta.parent.resource_version
+            local checker = healthcheck_manager.fetch_checker(resource_path, resource_version)
+            checkers = checkers or {}
+            checkers[instance.name] = checker
+        end
+    end
+
+    local version = plugin.conf_version(conf)
+    if checkers then
+        local status_ver = get_checkers_status_ver(checkers)
+        version = version .. "#" .. status_ver
+    end
+
     local server_picker = ctx.server_picker
     if not server_picker then
-        server_picker = lrucache_server_picker(ctx.matched_route.key, plugin.conf_version(conf),
-                                               create_server_picker, conf, ups_tab)
+        server_picker = lrucache_server_picker(ctx.matched_route.key, version,
+                                               create_server_picker, conf, ups_tab, checkers)
     end
     if not server_picker then
         return nil, nil, "failed to fetch server picker"

apisix/plugins/ai-proxy/schema.lua

@@ -14,6 +14,8 @@
 -- See the License for the specific language governing permissions and
 -- limitations under the License.
 --
+local schema_def = require("apisix.schema_def")
+
 local _M = {}
 
 local auth_item_schema = {
@@ -88,6 +90,13 @@ local ai_instance_schema = {
                     },
                 },
             },
+            checks = {
+                type = "object",
+                properties = {
+                    active = schema_def.health_checker_active,
+                },
+                required = {"active"}
+            }
         },
         required = {"name", "provider", "auth", "weight"}
     },

apisix/router.lua

@@ -18,6 +18,7 @@ local require = require
 local http_route = require("apisix.http.route")
 local apisix_upstream = require("apisix.upstream")
 local core    = require("apisix.core")
+local set_plugins_meta_parent = require("apisix.plugin").set_plugins_meta_parent
 local str_lower = string.lower
 local ipairs  = ipairs
 
@@ -33,6 +34,8 @@ local function filter(route)
         return
     end
 
+    set_plugins_meta_parent(route.value.plugins, route)
+
     if route.value.host then
         route.value.host = str_lower(route.value.host)
     elseif route.value.hosts then