93 lines
2.9 KiB
Lua
93 lines
2.9 KiB
Lua
-- Based on https://github.com/Klaessen/openresty-loadbalancers/blob/main/sticky-balancer.lua
|
|
|
|
local _M = {}
|
|
|
|
local balancer = require "ngx.balancer"
|
|
local cookie_name = "SERVER_ID"
|
|
|
|
-- Define backend server based on their capabilities (IP, port, weight, number of retries before timeout, duration of timeout, current number of fails, fail timestamp)
|
|
local servers = {
|
|
{ "127.0.0.1", 10062, weight = 1, max_fails = 3, fail_timeout = 30, fail_count = 0, last_fail_time = 0 },
|
|
{ "127.0.0.1", 10072, weight = 1, max_fails = 3, fail_timeout = 30, fail_count = 0, last_fail_time = 0 },
|
|
{ "127.0.0.1", 10082, weight = 1, max_fails = 3, fail_timeout = 30, fail_count = 0, last_fail_time = 0 },
|
|
}
|
|
|
|
-- Put the ammount of servers in the dictionary
|
|
local s = ngx.shared.servers
|
|
s:set("invidious-servers", #servers)
|
|
|
|
-- Generate a weighted server list based on weights
|
|
local function generate_weighted_server_list(servers)
|
|
local weighted_servers = {}
|
|
for _, server in ipairs(servers) do
|
|
for i = 1, server.weight do
|
|
table.insert(weighted_servers, server)
|
|
end
|
|
end
|
|
return weighted_servers
|
|
end
|
|
|
|
local weighted_servers = generate_weighted_server_list(servers)
|
|
|
|
-- Hash function to select server
|
|
local function hash(key, num_buckets)
|
|
local hash = ngx.crc32_long(key)
|
|
return (hash % num_buckets) + 1
|
|
end
|
|
|
|
-- Check if a server is available based on max_fails and fail_timeout
|
|
local function is_server_available(server)
|
|
if server.fail_count >= server.max_fails then
|
|
if (ngx.now() - server.last_fail_time) < server.fail_timeout then
|
|
return false
|
|
else
|
|
server.fail_count = 0
|
|
server.last_fail_time = 0
|
|
end
|
|
end
|
|
return true
|
|
end
|
|
|
|
-- Select server based on cookie or assign a new one
|
|
local function select_server()
|
|
local cookie = ngx.var["cookie_" .. cookie_name]
|
|
local server_index
|
|
|
|
math.randomseed(os.time())
|
|
|
|
if cookie then
|
|
server_index = tonumber(cookie)
|
|
ngx.header["X-Server-Id"] = server_index
|
|
else
|
|
-- server_index = hash(ngx.var.remote_addr, #weighted_servers)
|
|
server_index = math.random(#servers)
|
|
ngx.header["Set-Cookie"] = cookie_name .. "=" .. server_index .. "; domain=inv.nadekonw7plitnjuawu6ytjsl7jlglk2t6pyq6eftptmiv3dvqndwvyd.onion; Path=/; HttpOnly; SameSite=Lax"
|
|
end
|
|
|
|
local server = weighted_servers[server_index]
|
|
if is_server_available(server) then
|
|
return server
|
|
else
|
|
-- If the selected server is not available, find another one
|
|
for _, s in ipairs(weighted_servers) do
|
|
if is_server_available(s) then
|
|
return s
|
|
end
|
|
end
|
|
end
|
|
ngx.log(ngx.ERR, "No available servers")
|
|
return nil
|
|
end
|
|
|
|
local server = select_server()
|
|
if not server then
|
|
ngx.exit(502)
|
|
return
|
|
end
|
|
local ok, err = balancer.set_current_peer(server[1], server[2])
|
|
if not ok then
|
|
ngx.log(ngx.ERR, "Failed to set the current peer: ", err)
|
|
server.fail_count = server.fail_count + 1
|
|
server.last_fail_time = ngx.now()
|
|
return ngx.exit(500)
|
|
end
|