<commit_before>local pl_tablex = require "pl.tablex"
local singletons = require "kong.singletons"
local utils = require "kong.tools.utils"

-- due to startup/require order, cannot use the ones from 'singletons' here
local dns_client = require "resty.dns.client"

local table_concat = table.concat
local crc32 = ngx.crc32_short
local toip = dns_client.toip
local log = ngx.log
local sleep = ngx.sleep
local min = math.min
local max = math.max

local CRIT  = ngx.CRIT
local ERR   = ngx.ERR
local WARN  = ngx.WARN
local DEBUG = ngx.DEBUG
local EMPTY_T = pl_tablex.readonly {}

-- for unit-testing purposes only
local _load_upstreams_dict_into_memory
local _load_upstream_into_memory
local _load_targets_into_memory


--==============================================================================
-- Ring-balancer based resolution
--==============================================================================


-- table holding our balancer objects, indexed by upstream id
local balancers = {}


-- objects whose lifetimes are bound to that of a balancer
local healthcheckers = {}
local healthchecker_callbacks = {}
local target_histories = {}


-- health check API callbacks to be called on healthcheck events
local healthcheck_subscribers = {}


-- Caching logic
--
-- We retain 3 entities in singletons.cache:
--
-- 1) `"balancer:upstreams"` - a list of upstreams
--    to be invalidated on any upstream change
-- 2) `"balancer:upstreams:" .. id` - individual upstreams
--    to be invalidated on individual basis
-- 3) `"balancer:targets:" .. id`
--    target history for an upstream, invalidated:
--    a) along with the upstream it belongs to
--    b) upon any target change for the upstream (can only add entries)
--
-- Distinction between 1 and 2 makes it possible to invalidate individual
-- upstreams, instead of all at once forcing to rebuild all balancers


local function set_balancer(upstream_id, balancer)
  local prev = balancers[upstream_id]
  if prev then
    healthcheckers[prev] = nil
    healthchecker_callbacks[prev] = nil
    target_histories[prev] = nil
  end
  balancers[upstream_id] = balancer
end


local function stop_healthchecker(balancer)
  local healthchecker = healthcheckers[balancer]
  if healthchecker then
    local ok, err = healthchecker:clear()
    if not ok then
      log(ERR, "[healthchecks] error clearing healthcheck data: ", err)
    end
    healthchecker:stop()
    local hc_callback = healthchecker_callbacks[balancer]
    singletons.worker_events.unregister(hc_callback, healthchecker.EVENT_SOURCE)
  end
  healthcheckers[balancer] = nil
end


local get_upstream_by_id
do
  ------------------------------------------------------------------------------
  -- Loads a single upstream entity.
  -- @param upstream_id string
  -- @return the upstream table, or nil+error
  local function load_upstream_into_memory(upstream_id)
    log(DEBUG, "fetching upstream: ", tostring(upstream_id))

    local upstream, err = singletons.db.upstreams:select({id = upstream_id})
    if not upstream then
      return nil, err
    end

    return upstream
  end
  _load_upstream_into_memory = load_upstream_into_memory

  get_upstream_by_id = function(upstream_id)
    local upstream_cache_key = "balancer:upstreams:" .. upstream_id
    return singletons.cache:get(upstream_cache_key, nil,
                                load_upstream_into_memory, upstream_id)
  end
end


local fetch_target_history
do
  ------------------------------------------------------------------------------
  -- Loads the target history from the DB.
  -- @param upstream_id Upstream uuid for which to load the target history
  -- @return The target history array, with target entity tables.
  local function load_targets_into_memory(upstream_id)
    log(DEBUG, "fetching targets for upstream: ", tostring(upstream_id))

    local target_history, err, err_t =
      singletons.db.targets:select_by_upstream_raw({ id = upstream_id })

    if not target_history then
      return nil, err, err_t
    end

    -- perform some raw data updates
    for _, target in ipairs(target_history) do
      -- split `target` field into `name` and `port`
      local port
      target.name, port = string.match(target.target, "^(.-):(%d+)$")
      target.port = tonumber(port)
    end

    return target_history
  end
  _load_targets_into_memory = load_targets_into_memory


  ------------------------------------------------------------------------------
  -- Fetch target history, from cache or the DB.
  -- @param upstream The upstream entity object
  -- @return The target history array, with target entity tables.
  fetch_target_history = function(upstream)
    local targets_cache_key = "balancer:targets:" .. upstream.id
    return singletons.cache:get(targets_cache_key, nil,
                                load_targets_into_memory, upstream.id)
  end
end


--------------------------------------------------------------------------------
-- Applies the history of lb transactions from index `start` forward.
-- @param rb ring balancer object
-- @param history list of targets/transactions to be applied
-- @param start the index where to start in the `history` parameter
-- @return true
local function apply_history(rb, history, start)

  for i = start, #history do
    local target = history[i]

    if target.weight > 0 then
      assert(rb:addHost(target.name, target.port, target.weight))
    else
      assert(rb:removeHost(target.name, target.port))
    end

    target_histories[rb][i] = {
      name = target.name,
      port = target.port,
      weight = target.weight,
      order = target.order,
    }
  end

  return true
end


local function populate_healthchecker(hc, balancer)
  for weight, addr, host in balancer:addressIter() do
    if weight > 0 then
      local ipaddr = addr.ip
      local port = addr.port
      local hostname = host.hostname
      local ok, err = hc:add_target(ipaddr, port, hostname)
      if ok then
        -- Get existing health status which may have been initialized
        -- with data from another worker, and apply to the new balancer.
        local tgt_status = hc:get_target_status(ipaddr, port, hostname)
        balancer:setAddressStatus(tgt_status, ipaddr, port, hostname)

      else
        log(ERR, "[healthchecks] failed adding target: ", err)
      end<commit_msg>fix(balancer) don't set health if target is not created yet<commit_after>local pl_tablex = require "pl.tablex"
local singletons = require "kong.singletons"
local utils = require "kong.tools.utils"

-- due to startup/require order, cannot use the ones from 'singletons' here
local dns_client = require "resty.dns.client"

local table_concat = table.concat
local crc32 = ngx.crc32_short
local toip = dns_client.toip
local log = ngx.log
local sleep = ngx.sleep
local min = math.min
local max = math.max

local CRIT  = ngx.CRIT
local ERR   = ngx.ERR
local WARN  = ngx.WARN
local DEBUG = ngx.DEBUG
local EMPTY_T = pl_tablex.readonly {}

-- for unit-testing purposes only
local _load_upstreams_dict_into_memory
local _load_upstream_into_memory
local _load_targets_into_memory


--==============================================================================
-- Ring-balancer based resolution
--==============================================================================


-- table holding our balancer objects, indexed by upstream id
local balancers = {}


-- objects whose lifetimes are bound to that of a balancer
local healthcheckers = {}
local healthchecker_callbacks = {}
local target_histories = {}


-- health check API callbacks to be called on healthcheck events
local healthcheck_subscribers = {}


-- Caching logic
--
-- We retain 3 entities in singletons.cache:
--
-- 1) `"balancer:upstreams"` - a list of upstreams
--    to be invalidated on any upstream change
-- 2) `"balancer:upstreams:" .. id` - individual upstreams
--    to be invalidated on individual basis
-- 3) `"balancer:targets:" .. id`
--    target history for an upstream, invalidated:
--    a) along with the upstream it belongs to
--    b) upon any target change for the upstream (can only add entries)
--
-- Distinction between 1 and 2 makes it possible to invalidate individual
-- upstreams, instead of all at once forcing to rebuild all balancers


local function set_balancer(upstream_id, balancer)
  local prev = balancers[upstream_id]
  if prev then
    healthcheckers[prev] = nil
    healthchecker_callbacks[prev] = nil
    target_histories[prev] = nil
  end
  balancers[upstream_id] = balancer
end


local function stop_healthchecker(balancer)
  local healthchecker = healthcheckers[balancer]
  if healthchecker then
    local ok, err = healthchecker:clear()
    if not ok then
      log(ERR, "[healthchecks] error clearing healthcheck data: ", err)
    end
    healthchecker:stop()
    local hc_callback = healthchecker_callbacks[balancer]
    singletons.worker_events.unregister(hc_callback, healthchecker.EVENT_SOURCE)
  end
  healthcheckers[balancer] = nil
end


local get_upstream_by_id
do
  ------------------------------------------------------------------------------
  -- Loads a single upstream entity.
  -- @param upstream_id string
  -- @return the upstream table, or nil+error
  local function load_upstream_into_memory(upstream_id)
    log(DEBUG, "fetching upstream: ", tostring(upstream_id))

    local upstream, err = singletons.db.upstreams:select({id = upstream_id})
    if not upstream then
      return nil, err
    end

    return upstream
  end
  _load_upstream_into_memory = load_upstream_into_memory

  get_upstream_by_id = function(upstream_id)
    local upstream_cache_key = "balancer:upstreams:" .. upstream_id
    return singletons.cache:get(upstream_cache_key, nil,
                                load_upstream_into_memory, upstream_id)
  end
end


local fetch_target_history
do
  ------------------------------------------------------------------------------
  -- Loads the target history from the DB.
  -- @param upstream_id Upstream uuid for which to load the target history
  -- @return The target history array, with target entity tables.
  local function load_targets_into_memory(upstream_id)
    log(DEBUG, "fetching targets for upstream: ", tostring(upstream_id))

    local target_history, err, err_t =
      singletons.db.targets:select_by_upstream_raw({ id = upstream_id })

    if not target_history then
      return nil, err, err_t
    end

    -- perform some raw data updates
    for _, target in ipairs(target_history) do
      -- split `target` field into `name` and `port`
      local port
      target.name, port = string.match(target.target, "^(.-):(%d+)$")
      target.port = tonumber(port)
    end

    return target_history
  end
  _load_targets_into_memory = load_targets_into_memory


  ------------------------------------------------------------------------------
  -- Fetch target history, from cache or the DB.
  -- @param upstream The upstream entity object
  -- @return The target history array, with target entity tables.
  fetch_target_history = function(upstream)
    local targets_cache_key = "balancer:targets:" .. upstream.id
    return singletons.cache:get(targets_cache_key, nil,
                                load_targets_into_memory, upstream.id)
  end
end


--------------------------------------------------------------------------------
-- Applies the history of lb transactions from index `start` forward.
-- @param rb ring balancer object
-- @param history list of targets/transactions to be applied
-- @param start the index where to start in the `history` parameter
-- @return true
local function apply_history(rb, history, start)

  for i = start, #history do
    local target = history[i]

    if target.weight > 0 then
      assert(rb:addHost(target.name, target.port, target.weight))
    else
      assert(rb:removeHost(target.name, target.port))
    end

    target_histories[rb][i] = {
      name = target.name,
      port = target.port,
      weight = target.weight,
      order = target.order,
    }
  end

  return true
end


local function populate_healthchecker(hc, balancer)
  for weight, addr, host in balancer:addressIter() do
    if weight > 0 then
      local ipaddr = addr.ip
      local port = addr.port
      local hostname = host.hostname
      local ok, err = hc:add_target(ipaddr, port, hostname)
      if ok then
        -- Get existing health status which may have been initialized
        -- with data from another worker, and apply to the new balancer.
        local tgt_status = hc:get_target_status(ipaddr, port, hostname)
        if tgt_status ~= nil then
          balancer:setAddressStatus(tgt_status, ipaddr, port, hostname)
        end

      else
        log(ERR, "[healthchecks] failed adding target: ", err)
      end
{'fraction_non_alphanumeric': 0.148831005805743, 'fraction_numerical': 0.0018829436686019143, 'mean_word_length': 2.827927927927928, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4978897', 'n_tokens_mistral': 3721, 'n_tokens_neox': 3538, 'n_words': 1161}