-- HAProxy LUA extensions
-- REF: https://www.haproxy.com/blog/5-ways-to-extend-haproxy-with-lua/
-- MAN: haproxy-lua

-- Fetches

-- Suffix the given key to the transaction source (IP address)
--
-- Parameters:
--  - key: suffix key variable name (e.g. req.website)
local function thr_src_add(txn, key)
  local key = txn:get_var(key)
  return txn.f:src()..'/'..key
end
core.register_fetches('thr-src-add', thr_src_add);

-- Return the throttling key corresponding to the given SLO
--
-- Parameters:
--   - slo:     SLO (numeric ID) variable name (e.g. req.slo)
--   - website: website (name/ID) variable name (e.g. req.website)
--   - owner:   owner (name/ID) variable name (e.g. req.owner)
local function thr_key(txn, slo, website, owner)
  local slo_id = txn:get_var(slo)
  local slo_type = slo_id // 1000
  if slo_type == 1 then
    -- source
    return txn.f:src()..'//'
  elseif slo_type == 2 then
    -- website
    return '//'..txn:get_var(website)
  elseif slo_type == 3 then
    -- source+website
    return txn.f:src()..'//'..txn:get_var(website)
  elseif slo_type == 4 then
    -- owner
    return '/'..txn:get_var(owner)..'/'
  elseif slo_type == 5 then
    -- source+owner
    return txn.f:src()..'/'..txn:get_var(owner)..'/'
  else
    -- global
    return '//'
  end
end
core.register_fetches('thr-key', thr_key);

-- Return the (re-)estimated quantity of (milli)clients (sharing the same
--   throttling bucket/key)
--
-- Parameters:
--   - slo:  target SLO rate variable name (e.g. txn.slo_rate)
--   - rate: actual incoming requests rate variable name (e.g. txn.thr_rate)
--   - N:    current (milli)clients quantity variable name (e.g. txn.thr_clients)
local function thr_clients(txn, slo, rate, N)
  local slo = txn:get_var(slo)
  local rate = txn:get_var(rate)
  local N = txn:get_var(N)
  -- the very first fetch from GPT0 will return 0
  if N < 1000 then N = 1000 end
  -- re-estimate the quantity of clients
  N = N*(1+1000*(rate-slo)/(N*slo))//1
  if N < 1000 then
    -- less than 1000 (milli)client - 1 client - makes no sense
    N = 1000
  elseif N > 1000000000 then
    -- more than 1000000000 (milli)client - 1'000'000 clients - neither
    N = 1000000000
  end
  return N
end
core.register_fetches('thr-clients', thr_clients);


-- Actions

-- Delay (sleep) by time targeting the SLO rate, taking into account the
--   actual throttling rate
--
-- Parameters:
--  - slo:  target SLO rate variable name (e.g. txn.slo_rate)
--  - rate: actual incoming requests rate variable name (e.g. txn.thr_rate)
--  - N:    (milli)clients quantity variable name (e.g. txn.thr_clients)
local function thr_delay(txn, slo, rate, N)
  local slo = txn:get_var(slo)
  local rate = txn:get_var(rate)
  local N = txn:get_var(N)
  core.msleep(math.min(5000, N*(1.0+0.5*(math.random()-0.5))*math.min(1.0,rate/slo)/slo)//1)
end
core.register_action('thr-delay', {'http-req', 'http-res'}, thr_delay, 3);