try_register

lib/resty/checkups/api.lua

@@ -170,6 +170,16 @@ function _M.get_ups_timeout(skey)
 end
 
 
+function _M.get_ups(skey)
+    if not skey then
+        return
+    end
+
+    local ups = base.upstream.checkups[skey]
+    return ups
+end
+
+
 function _M.create_checker()
     local phase = get_phase()
     if phase ~= "init_worker" then
@@ -233,14 +243,25 @@ local function gen_upstream(skey, upstream)
             return nil, "cluster invalid"
         end
     else
-        -- only servers
         local dyupstream, err = dyconfig.do_get_upstream(skey)
         if err then
             return nil, err
         end
 
         dyupstream = dyupstream or {}
-        dyupstream.cluster = upstream
+        if upstream.servers then
+            -- store config
+            for k, v in pairs(upstream) do
+                if k ~= "servers" then
+                    dyupstream[k] = v
+                else
+                    dyupstream.cluster = { { servers = v } }
+                end
+            end
+        else
+            -- only cluster
+            dyupstream.cluster = upstream
+        end
         ups = dyupstream
     end
 
@@ -322,4 +343,14 @@ function _M.delete_upstream(skey)
 end
 
 
+function _M.try_register(name, module)
+    return try.register(name, module)
+end
+
+
+function _M.try_unregister(name)
+    return try.unregister(name)
+end
+
+
 return _M

lib/resty/checkups/consistent_hash.lua

@@ -4,6 +4,8 @@ local floor      = math.floor
 local str_byte   = string.byte
 local tab_sort   = table.sort
 local tab_insert = table.insert
+local ipairs     = ipairs
+local type       = type
 
 local _M = { _VERSION = "0.11" }
 
@@ -22,7 +24,7 @@ local function hash_string(str)
 end
 
 
-local function init_consistent_hash_state(servers)
+local function init_state(servers)
     local weight_sum = 0
     for _, srv in ipairs(servers) do
         weight_sum = weight_sum + (srv.weight or 1)
@@ -62,37 +64,52 @@ local function binary_search(circle, key)
 end
 
 
-function _M.next_consistent_hash_server(servers, peer_cb, hash_key)
-    local is_tab = require "resty.checkups.base".is_tab
-    servers.chash = is_tab(servers.chash) and servers.chash
-                    or init_consistent_hash_state(servers)
+local function next_server(servers, peer_cb, opts)
+    servers.chash = type(servers.chash) == "table" and servers.chash
+                    or init_state(servers)
 
     local chash = servers.chash
     if chash.members == 1 then
         if peer_cb(1, servers[1]) then
-            return servers[1]
+            return 1, servers[1]
         end
 
-        return nil, "consistent hash: no servers available"
+        return nil, nil, nil, "consistent hash: no servers available"
     end
 
     local circle = chash.circle
-    local st = binary_search(circle, hash_string(hash_key))
+    local st = binary_search(circle, hash_string(opts.hash_key))
     local size = #circle
     local ed = st + size - 1
     for i = st, ed do  -- TODO: algorithm O(n)
         local idx = circle[(i - 1) % size + 1][2]
         if peer_cb(idx, servers[idx]) then
-            return servers[idx]
+            return idx, servers[idx]
         end
     end
 
-    return nil, "consistent hash: no servers available"
+    return nil, nil, nil, "consistent hash: no servers available"
 end
 
 
-function _M.free_consitent_hash_server(srv, failed)
-    return
+local function gen_opts(ups, opts)
+    local key
+    local mode = ups.mode
+    if mode == "hash" then
+        key = opts.hash_key or ngx.var.uri
+    elseif mode == "url_hash" then
+        key = ngx.var.uri
+    elseif mode == "ip_hash" then
+        key = ngx.var.remote_addr
+    elseif mode == "header_hash" then
+        key = ngx.var.http_x_hash_key or ngx.var.uri
+    end
+    return { hash_key=key }
+end
+
+function _M.ipairsrvs(servers, peer_cb, ups, opts)
+    local mopts = gen_opts(ups, opts)
+    return function() return next_server(servers, peer_cb, mopts) end
 end
 
 

lib/resty/checkups/round_robin.lua

@@ -13,19 +13,20 @@ return:
     - (table) server
     - (string) error
 --]]
-function _M.next_round_robin_server(servers, peer_cb)
+local function next_server(servers, peer_cb)
     local srvs_cnt = #servers
 
     if srvs_cnt == 1 then
         if peer_cb(1, servers[1]) then
-            return servers[1], nil
+            return 1, servers[1], nil
         end
 
-        return nil, "round robin: no servers available"
+        return nil, nil, nil, "round robin: no servers available"
     end
 
     -- select round robin server
     local best
+    local best_idx
     local max_weight
     local weight_sum = 0
     for idx = 1, srvs_cnt do
@@ -46,22 +47,22 @@ function _M.next_round_robin_server(servers, peer_cb)
             if not max_weight or srv.current_weight > max_weight then
                 max_weight = srv.current_weight
                 best = srv
+                best_idx = idx
             end
         end
     end
 
     if not best then
-        return nil, "round robin: no servers available"
+        return nil, nil, nil, "round robin: no servers available"
     end
 
     best.current_weight = best.current_weight - weight_sum
 
-    return best, nil
+    return best_idx, best
 end
 
 
-
-function _M.free_round_robin_server(srv, failed)
+function _M.free_server(srv, failed)
     if not failed then
         return
     end
@@ -70,4 +71,9 @@ function _M.free_round_robin_server(srv, failed)
 end
 
 
+function _M.ipairsrvs(servers, peer_cb, ups, opts)
+    return function() return next_server(servers, peer_cb) end
+end
+
+
 return _M