diff --git a/prometheus_keys.lua b/prometheus_keys.lua index faeaa66..6ae3315 100644 --- a/prometheus_keys.lua +++ b/prometheus_keys.lua @@ -14,32 +14,87 @@ KeyIndex.__index = KeyIndex -- and the index converges even when far more slots need repairing. local MAX_KEY_COUNT_REPAIRS = 1000 +-- Upper bound on how many generation switches a single add() call follows +-- for one key before it gives up on that key. +local MAX_GENERATION_RETRIES = 3 + +-- Slots are never reused: every key that expires or is removed leaves a dead +-- slot behind, and every full sync (a worker starting, or any delete) walks +-- all of them. Once key_count reaches COMPACT_MIN_SLOTS and is at least +-- COMPACT_RATIO times the live keys, remove_expired_keys() copies the live +-- keys into a fresh generation of slots and drops the old one, so key_count +-- stays proportional to the live keys instead of to the keys ever created. +local COMPACT_MIN_SLOTS = 10000 +local COMPACT_RATIO = 2 +-- Slots handled between yields while compacting. +local COMPACT_BATCH = 1000 +-- The compaction lock and the old generation's key_count expire after this +-- many seconds; the lock is refreshed after every batch. +local COMPACT_LOCK_TTL = 60 +-- An add() in another worker writes its slot before it increments key_count, +-- so up to one slot per concurrent writer may sit past key_count. Scans that +-- must not miss such slots look this far past key_count. +local COMPACT_TAIL_SLACK = 64 +local COMPACT_ERR_PREFIX = "key index compaction: " + -- check and remove expired keys -local function remove_expired_keys(_, self) +local function remove_expired_keys(premature, self) + if premature then + return + end self:remove_expired_keys() end +-- Generation 0 keeps the historical key names, so workers running a version +-- without compaction still share the same index while nginx reloads. +local function generation_names(prefix, gen) + if gen == 0 then + return prefix .. "key_", prefix .. "key_count" + end + return prefix .. gen .. "_key_", prefix .. gen .. "_key_count" +end + +-- ngx.sleep(0) is a posted event, which does not keep an exiting worker alive: +-- the worker would exit with the compaction suspended and its lock held. A +-- sleep timer does, and an exiting worker finishes without yielding. +local function yield() + if ngx.get_phase() == "timer" and not ngx.worker.exiting() then + ngx.sleep(0.001) + end +end + function KeyIndex.new(shared_dict, prefix, remove_expired_keys_interval) local self = setmetatable({}, KeyIndex) self.dict = shared_dict - self.key_prefix = prefix .. "key_" + self.prefix = prefix + self.gen_key = prefix .. "gen" + self.lock_key = prefix .. "compact_lock" self.delete_count = prefix .. "delete_count" - self.key_count = prefix .. "key_count" - self.last = 0 self.deleted = 0 self.not_expired_index = 1 - self.keys = {} - self.index = {} - self.expire_keys = {} + self.compact_min_slots = COMPACT_MIN_SLOTS + self:use_generation(0) ngx.timer.every(remove_expired_keys_interval or 600, remove_expired_keys, self) return self end +-- Points this worker at generation `gen` with an empty local view; the next +-- sync() loads it. +function KeyIndex:use_generation(gen) + self.gen = gen + self.key_prefix, self.key_count = generation_names(self.prefix, gen) + self.last = 0 + self.keys = {} + self.index = {} + self.expire_keys = {} +end + -- check and remove expired keys function KeyIndex:remove_expired_keys() + self:sync() for i, _ in pairs(self.expire_keys) do -- Read i-th key. If it is nil or ttl is < 0, it means it was expired local ttl, err = self.dict:ttl(self.key_prefix .. i) @@ -65,10 +120,23 @@ function KeyIndex:remove_expired_keys() -- entries are indistinguishable from absent ones through every dict API, -- reclaiming them here cannot change any observable behaviour. self.dict:flush_expired() + + self:compact() end -- Loads new keys that might have been added by other workers since last sync. function KeyIndex:sync() + local gen = self.dict:get(self.gen_key) or 0 + if gen > self.gen then + -- Another worker compacted the index; the slots tracked locally are + -- being deleted. + self:use_generation(gen) + elseif gen < self.gen then + -- The generation only ever grows, so a lower value means its node was + -- evicted from a full dict. + self.dict:set(self.gen_key, self.gen) + end + local delete_count = self.dict:get(self.delete_count) or 0 local N = self.dict:get(self.key_count) or 0 if self.deleted ~= delete_count then @@ -147,7 +215,13 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime) local retried = false local repairs = 0 local repair_forcible = false + local slot_forcible = false + local switches = 0 while true do + if switches > MAX_GENERATION_RETRIES then + return ("key index: generation switched " .. switches .. + " times while adding key, dropping key: " .. key) + end local N = self:sync() if self.index[key] ~= nil then -- key already exists, if has exptime, set expire @@ -182,6 +256,15 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime) end end if not expired then + if exptime and (self.dict:get(self.gen_key) or 0) > self.gen then + -- A compaction switched generations after this worker synced and + -- may have reconciled the copy before this renewal of the old + -- slot, so renew the key again in the current generation. A lower + -- generation only means its node was evicted (see sync()). + switches = switches + 1 + retried = false + goto continue + end if repair_forcible then -- the key was adopted from an occupied slot after repair -- increments that forcibly displaced other entries; report the @@ -194,19 +277,28 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime) end end N = N+1 - local ok, err, forcible = self.dict:add(self.key_prefix .. N, key, exptime) + local slot = self.key_prefix .. N + local ok, err, forcible = self.dict:add(slot, key, exptime) if ok then local _, _, forcible2 = self.dict:incr(self.key_count, 1, 0) - self.keys[N] = key - self.index[key] = N - if exptime and exptime > 0 then - self.expire_keys[N] = true - end - if forcible or forcible2 or repair_forcible then - return (err_msg_lru_eviction .. "; key index: add key: idx=" .. - self.key_prefix .. N .. ", key=" .. key) + slot_forcible = slot_forcible or forcible or forcible2 + if (self.dict:get(self.gen_key) or 0) <= self.gen then + self.keys[N] = key + self.index[key] = N + if exptime and exptime > 0 then + self.expire_keys[N] = true + end + if slot_forcible or repair_forcible then + return (err_msg_lru_eviction .. "; key index: add key: idx=" .. + slot .. ", key=" .. key) + end + break end - break + -- A compaction switched generations after this worker synced, so the + -- slot may land in the generation being dropped after its keys were + -- copied. Delete it and register the key in the current generation. + self.dict:delete(slot) + switches = switches + 1 elseif err ~= "exists" then return "Unexpected error adding a key: " .. err end @@ -223,7 +315,7 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime) -- lock held hot (apache/apisix#12275). Advance the counter past the -- occupied slot instead: the next sync() adopts that slot's occupant -- and progress resumes. - if retried then + if not ok and retried then local _, incr_err, forcible3 = self.dict:incr(self.key_count, 1, 0) if incr_err then -- hard failure (e.g. "no memory"): give up immediately, mirroring @@ -246,7 +338,10 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime) " without finding a free slot, dropping key: " .. key) end end - retried = true + -- after a generation switch the slots start over, so an earlier + -- "exists" in the old generation says nothing about the new one + retried = not ok + ::continue:: end end end @@ -256,6 +351,7 @@ end -- Args: -- key: String value of the key, must exists in this index. function KeyIndex:remove(key, err_msg_lru_eviction) + self:sync() local i = self.index[key] if i then self.index[key] = nil @@ -274,4 +370,208 @@ function KeyIndex:remove(key, err_msg_lru_eviction) end end -return KeyIndex \ No newline at end of file +-- Yields, then extends the compaction lock. Returns why the compaction has to +-- stop, or nil to go on. +function KeyIndex:refresh_compact_lock(token) + yield() + if ngx.worker.exiting() then + return "worker is exiting" + end + if self.dict:get(self.lock_key) ~= token then + return "lost lock" + end + self.dict:expire(self.lock_key, COMPACT_LOCK_TTL) +end + +function KeyIndex:release_compact_lock(token) + if self.dict:get(self.lock_key) == token then + self.dict:delete(self.lock_key) + end +end + +-- Deletes slots first..last of the generation named by `slot_prefix`, +-- yielding between batches. +function KeyIndex:delete_slots(slot_prefix, first, last) + for i = first, last do + self.dict:delete(slot_prefix .. i) + if i % COMPACT_BATCH == 0 then + yield() + end + end +end + +-- Deletes the leftovers of a compaction into generation `gen` that stopped +-- before switching to it; those slots are written densely from 1. +function KeyIndex:clear_generation(gen) + local slot_prefix, count_key = generation_names(self.prefix, gen) + local misses, i = 0, 1 + while misses < COMPACT_TAIL_SLACK do + if self.dict:get(slot_prefix .. i) == nil then + misses = misses + 1 + else + misses = 0 + self.dict:delete(slot_prefix .. i) + end + i = i + 1 + end + self.dict:delete(count_key) +end + +-- Copies the live keys into generation self.gen + 1 and switches every worker +-- to it once most slots of the current generation are dead. Runs in at most +-- one worker at a time. +-- +-- The copy is written before the generation switch, so other workers never +-- list a partial index. Changes other workers make to the old generation +-- while it is copied are applied afterwards: keys removed in the meantime are +-- dropped from the copy, and slots written past the copied range are added. +-- add() re-registers keys whose slot it wrote into the old generation after +-- the switch. +function KeyIndex:compact() + local N = self:sync() + if N < self.compact_min_slots or ngx.worker.exiting() then + return + end + local live = 0 + for _ in pairs(self.index) do + live = live + 1 + end + if N < live * COMPACT_RATIO then + return + end + + local token = ngx.worker.pid() .. ":" .. N + local locked, lock_err = self.dict:add(self.lock_key, token, COMPACT_LOCK_TTL) + if not locked then + if lock_err ~= "exists" then + ngx.log(ngx.ERR, COMPACT_ERR_PREFIX, "failed to take lock: ", lock_err) + end + return + end + + local from_prefix, from_count = self.key_prefix, self.key_count + local to_gen = self.gen + 1 + local to_prefix, to_count = generation_names(self.prefix, to_gen) + self:clear_generation(to_gen) + + local keys, index, expire_keys, origin, copied_ttl = {}, {}, {}, {}, {} + local M = 0 + local function abort(msg) + self:delete_slots(to_prefix, 1, M) + self.dict:delete(to_count) + self:release_compact_lock(token) + if not ngx.worker.exiting() then + ngx.log(ngx.ERR, COMPACT_ERR_PREFIX, msg) + end + end + + for i = 1, N do + local slot = from_prefix .. i + local key = self.dict:get(slot) + local ttl = key and self.dict:ttl(slot) + if ttl and not index[key] then + local exptime = ttl > 0 and ttl or nil + M = M + 1 + local ok, err, forcible = self.dict:set(to_prefix .. M, key, exptime) + if not ok or forcible then + return abort(err or "copying a slot evicted other entries") + end + keys[M], index[key], origin[M], copied_ttl[M] = key, M, i, ttl + if exptime then + expire_keys[M] = true + end + end + if i % COMPACT_BATCH == 0 then + local stop = self:refresh_compact_lock(token) + if stop then + return abort(stop .. " while copying slots") + end + end + end + + local ok, err, forcible = self.dict:set(to_count, M) + if not ok or forcible then + return abort(err or "writing key_count evicted other entries") + end + local stop = self:refresh_compact_lock(token) + if stop then + return abort(stop .. " before switching generation") + end + ok, err = self.dict:set(self.gen_key, to_gen) + if not ok then + return abort(err) + end + + self:use_generation(to_gen) + self.keys, self.index, self.expire_keys, self.last = keys, index, expire_keys, M + + -- Workers renew and remove keys in the old generation until they switch, + -- and in the new one afterwards. A ttl above the one read while copying + -- (0 for permanent keys) means the slot was renewed: keep the key while + -- either slot holds it unrenewed or renewed, never shorten the new slot, + -- and drop keys whose unrenewed slot is gone from either generation. + local dropped = false + for j, i in pairs(origin) do + local key, copied = keys[j], copied_ttl[j] + local old_slot, new_slot = from_prefix .. i, to_prefix .. j + local old_ttl = self.dict:get(old_slot) == key and self.dict:ttl(old_slot) + local new_ttl = self.dict:get(new_slot) == key and self.dict:ttl(new_slot) + local old_renewed = old_ttl and old_ttl > copied + local keep + if new_ttl then + keep = old_ttl or new_ttl > copied + if old_renewed and new_ttl > 0 and new_ttl < old_ttl then + self.dict:expire(new_slot, old_ttl) + end + elseif old_renewed then + keep = self.dict:add(new_slot, key, old_ttl) + end + if not keep then + if new_ttl then + self.dict:delete(new_slot) + end + if self.index[key] == j then + self.index[key] = nil + end + self.keys[j] = nil + self.expire_keys[j] = nil + dropped = true + end + if j % COMPACT_BATCH == 0 then + yield() + end + end + if dropped then + self.deleted = self.deleted + 1 + self.dict:incr(self.delete_count, 1, 0) + end + + local last = N + repeat + local first = last + 1 + last = (self.dict:get(from_count) or 0) + COMPACT_TAIL_SLACK + for i = first, last do + local slot = from_prefix .. i + local key = self.dict:get(slot) + local ttl = key and self.dict:ttl(slot) + if ttl then + local add_err = self:add(key, COMPACT_ERR_PREFIX .. "evicted entries", + ttl > 0 and ttl or nil) + if add_err then + ngx.log(ngx.ERR, COMPACT_ERR_PREFIX, add_err) + end + end + if i % COMPACT_BATCH == 0 then + yield() + end + end + until (self.dict:get(from_count) or 0) + COMPACT_TAIL_SLACK <= last + + self:delete_slots(from_prefix, 1, last) + -- Not deleted: a late add() in another worker would re-create it without + -- an expiry. + self.dict:expire(from_count, COMPACT_LOCK_TTL) + self:release_compact_lock(token) +end + +return KeyIndex diff --git a/prometheus_test.lua b/prometheus_test.lua index b485b0a..a3211ac 100644 --- a/prometheus_test.lua +++ b/prometheus_test.lua @@ -123,6 +123,12 @@ Nginx.worker = {} function Nginx.worker.id() return 'testworker' end +function Nginx.worker.pid() + return 1 +end +function Nginx.worker.exiting() + return false +end function Nginx.sleep() end Nginx.timer = {} function Nginx.timer.every(_, _, _) end @@ -1043,6 +1049,468 @@ function TestKeyIndex:testKeyCountRepairReportsForcibleEviction() luaunit.assertEquals(self.dict:get("_prefix_key_count"), 30) end +local function count_nodes(dict, prefix) + local n = 0 + for k in pairs(dict.dict) do + if k:find(prefix, 1, true) == 1 then n = n + 1 end + end + return n +end + +-- Registers one permanent key and `churn` keys expiring after 1 second, and +-- lowers the compaction threshold below that slot count. +local function churn_slots(key_index, churn) + key_index.compact_min_slots = 10 + key_index:add("permanent", "eviction_err") + for i = 1, churn do + key_index:add("churn" .. i, "eviction_err", 1) + end +end + +-- Slots are never reused, so without compaction every full sync walks every +-- slot ever allocated. Once most slots are dead the live keys move to a new +-- generation and the old one is deleted. +function TestKeyIndex:testCompactDropsDeadSlots() + churn_slots(self.key_index, 30) + self.key_index:add("long", "eviction_err", 100) + luaunit.assertEquals(self.dict:get("_prefix_key_count"), 32) + sleep(2) + + self.key_index:remove_expired_keys() + + luaunit.assertEquals(ngx.logs, nil) + luaunit.assertEquals(self.dict:get("_prefix_gen"), 1) + luaunit.assertEquals(self.dict:get("_prefix_1_key_count"), 2) + luaunit.assertEquals(self.dict:get("_prefix_1_key_1"), "permanent") + luaunit.assertEquals(self.dict:get("_prefix_1_key_2"), "long") + -- expiry carries over: none for permanent keys, the remaining ttl otherwise + luaunit.assertEquals(self.dict:ttl("_prefix_1_key_1"), 0) + luaunit.assertTrue(self.dict:ttl("_prefix_1_key_2") > 90) + -- of generation 0 only its key_count is left, and it expires + luaunit.assertEquals(count_nodes(self.dict, "_prefix_key_"), 1) + luaunit.assertTrue(self.dict:ttl("_prefix_key_count") > 0) + luaunit.assertNil(self.dict:get("_prefix_compact_lock")) + + local keys = self.key_index:list() + table.sort(keys) + luaunit.assertEquals(keys, {"long", "permanent"}) + luaunit.assertNil(self.key_index:add("new", "eviction_err")) + luaunit.assertEquals(self.dict:get("_prefix_1_key_3"), "new") +end + +function TestKeyIndex:testCompactOnlyWhenMostSlotsAreDead() + for i = 1, 20 do + self.key_index:add("key" .. i, "eviction_err") + end + -- below the default slot threshold + self.key_index:remove_expired_keys() + luaunit.assertNil(self.dict:get("_prefix_gen")) + + -- above the threshold, but all slots are live + self.key_index.compact_min_slots = 10 + self.key_index:remove_expired_keys() + luaunit.assertNil(self.dict:get("_prefix_gen")) + luaunit.assertEquals(self.dict:get("_prefix_key_count"), 20) +end + +function TestKeyIndex:testCompactSkipsWhileAnotherWorkerHoldsLock() + churn_slots(self.key_index, 30) + sleep(2) + self.dict:set("_prefix_compact_lock", "other") + + self.key_index:remove_expired_keys() + + luaunit.assertNil(self.dict:get("_prefix_gen")) + luaunit.assertEquals(self.dict:get("_prefix_key_1"), "permanent") + luaunit.assertEquals(self.dict:get("_prefix_compact_lock"), "other") +end + +function TestKeyIndex:testOtherWorkersFollowCompaction() + churn_slots(self.key_index, 30) + self.key_index:add("gone", "eviction_err") + local worker2 = require('prometheus_keys').new(self.dict, "_prefix_", 1) + luaunit.assertEquals(#worker2:list(), 32) + sleep(2) + + self.key_index:remove_expired_keys() + + local keys = worker2:list() + table.sort(keys) + luaunit.assertEquals(keys, {"gone", "permanent"}) + luaunit.assertEquals(worker2.gen, 1) + + -- remove() on a worker that has not synced since the switch must remove the + -- key from the new generation, not from the deleted slots it still tracks + local worker3 = require('prometheus_keys').new(self.dict, "_prefix_", 1) + worker3.gen, worker3.last = 0, 32 + worker3.keys, worker3.index = {[32] = "gone"}, {gone = 32} + luaunit.assertNil(worker3:remove("gone")) + luaunit.assertEquals(self.key_index:list(), {"permanent"}) + + -- a worker started after the switch loads only the new generation + local worker4 = require('prometheus_keys').new(self.dict, "_prefix_", 1) + luaunit.assertEquals(worker4:list(), {"permanent"}) +end + +-- Another worker writes its slot before it increments key_count; the copy +-- stops at key_count, so such a slot must still reach the new generation. +function TestKeyIndex:testCompactKeepsSlotWrittenPastKeyCount() + churn_slots(self.key_index, 30) + sleep(2) + self.dict:add("_prefix_key_32", "late") + + self.key_index:remove_expired_keys() + + local keys = self.key_index:list() + table.sort(keys) + luaunit.assertEquals(keys, {"late", "permanent"}) + luaunit.assertNil(self.dict:get("_prefix_key_32")) +end + +-- A key removed from the old generation after it was copied must not come +-- back through the copy. +function TestKeyIndex:testCompactDropsKeyRemovedDuringCopy() + churn_slots(self.key_index, 30) + self.key_index:add("doomed", "eviction_err") + local worker2 = require('prometheus_keys').new(self.dict, "_prefix_", 1) + worker2:sync() + sleep(2) + + -- worker2 removes the key once the copy is written, before the switch + local set = SimpleDict.set + self.dict.set = function(d, k, v, exptime) + if k == "_prefix_1_key_count" then + d.set = nil + worker2:remove("doomed") + end + return set(d, k, v, exptime) + end + self.key_index:remove_expired_keys() + + luaunit.assertEquals(self.dict:get("_prefix_gen"), 1) + luaunit.assertEquals(self.key_index:list(), {"permanent"}) + luaunit.assertEquals(worker2:list(), {"permanent"}) + luaunit.assertNil(self.dict:get("_prefix_1_key_2")) +end + +-- Until the switch, other workers renew keys in the old generation; the copy +-- must end up with the latest expiry, not the one read while copying. +function TestKeyIndex:testCompactCarriesRenewalDuringCopy() + churn_slots(self.key_index, 30) + self.key_index:add("renewed", "eviction_err", 5) + sleep(2) + + local set = SimpleDict.set + self.dict.set = function(d, k, v, exptime) + if k == "_prefix_1_key_count" then + d.set = nil + d:expire("_prefix_key_32", 100) + end + return set(d, k, v, exptime) + end + self.key_index:remove_expired_keys() + + luaunit.assertEquals(self.dict:get("_prefix_gen"), 1) + local slot = "_prefix_1_key_" .. self.key_index.index["renewed"] + luaunit.assertEquals(self.dict:get(slot), "renewed") + luaunit.assertTrue(self.dict:ttl(slot) > 90) +end + +-- Runs `after_switch` right after the compaction switches to generation 1, +-- before it reconciles the copy with the old generation. +local function on_generation_switch(dict, after_switch) + local set = SimpleDict.set + dict.set = function(d, k, v, exptime) + local ok, err, forcible = set(d, k, v, exptime) + if k == "_prefix_gen" then + d.set = nil + after_switch() + end + return ok, err, forcible + end +end + +-- After the switch, other workers renew keys in the new generation; the +-- expiry left in the old one must not shorten them. +function TestKeyIndex:testCompactKeepsRenewalAfterSwitch() + churn_slots(self.key_index, 30) + self.key_index:add("renewed", "eviction_err", 5) + local worker2 = require('prometheus_keys').new(self.dict, "_prefix_", 1) + sleep(2) + + on_generation_switch(self.dict, function() + luaunit.assertNil(worker2:add("renewed", "eviction_err", 100)) + end) + self.key_index:remove_expired_keys() + + local slot = "_prefix_1_key_" .. self.key_index.index["renewed"] + luaunit.assertEquals(self.dict:get(slot), "renewed") + luaunit.assertTrue(self.dict:ttl(slot) > 90) +end + +-- The old slot of a key renewed after the switch may expire before the copy is +-- reconciled; the renewed slot must survive. +function TestKeyIndex:testCompactKeepsRenewalAfterOldSlotExpired() + churn_slots(self.key_index, 30) + self.key_index:add("renewed", "eviction_err", 5) + local worker2 = require('prometheus_keys').new(self.dict, "_prefix_", 1) + sleep(2) + + on_generation_switch(self.dict, function() + luaunit.assertNil(worker2:add("renewed", "eviction_err", 100)) + self.dict:delete("_prefix_key_32") + end) + self.key_index:remove_expired_keys() + + local slot = "_prefix_1_key_" .. self.key_index.index["renewed"] + luaunit.assertEquals(self.dict:get(slot), "renewed") + luaunit.assertTrue(self.dict:ttl(slot) > 90) + local keys = worker2:list() + table.sort(keys) + luaunit.assertEquals(keys, {"permanent", "renewed"}) + local worker3 = require('prometheus_keys').new(self.dict, "_prefix_", 1) + keys = worker3:list() + table.sort(keys) + luaunit.assertEquals(keys, {"permanent", "renewed"}) +end + +-- A worker that synced before the switch may renew the old slot after the +-- copy was reconciled and before the old slots are deleted; the renewal must +-- reach the new slot. +function TestKeyIndex:testAddRenewsNewSlotWhenSwitchFollowsSync() + churn_slots(self.key_index, 30) + self.key_index:add("renewed", "eviction_err", 5) + local worker2 = require('prometheus_keys').new(self.dict, "_prefix_", 1) + sleep(2) + + self.key_index.delete_slots = function() end + self.dict.expire = function(d, k, exptime) + if k == "_prefix_key_32" then + d.expire = nil + self.key_index:remove_expired_keys() + luaunit.assertEquals(d:get("_prefix_gen"), 1) + end + return SimpleDict.expire(d, k, exptime) + end + luaunit.assertNil(worker2:add("renewed", "eviction_err", 100)) + + luaunit.assertEquals(worker2.gen, 1) + local slot = "_prefix_1_key_" .. self.key_index.index["renewed"] + luaunit.assertEquals(self.dict:get(slot), "renewed") + luaunit.assertTrue(self.dict:ttl(slot) > 90) +end + +-- A generation evicted from a full dict that cannot be written back is not a +-- generation switch; add() must not loop on it. +function TestKeyIndex:testAddDoesNotLoopOnEvictedGeneration() + churn_slots(self.key_index, 30) + self.key_index:add("renewed", "eviction_err", 100) + sleep(2) + self.key_index:remove_expired_keys() + luaunit.assertEquals(self.dict:get("_prefix_gen"), 1) + + local calls = 0 + local set = SimpleDict.set + self.dict.set = function(d, k, v, exptime) + if k == "_prefix_gen" then + calls = calls + 1 + luaunit.assertTrue(calls < 10, "add() keeps looping") + return false, "no memory", false + end + return set(d, k, v, exptime) + end + self.dict:delete("_prefix_gen") + + luaunit.assertNil(self.key_index:add("renewed", "eviction_err", 100)) + luaunit.assertNil(self.key_index:add("new", "eviction_err", 100)) + + luaunit.assertEquals(self.key_index.gen, 1) + luaunit.assertEquals(self.dict:get("_prefix_1_key_3"), "new") + local keys = self.key_index:list() + table.sort(keys) + luaunit.assertEquals(keys, {"new", "permanent", "renewed"}) +end + +-- Generations that keep switching under add() must not keep it looping. +function TestKeyIndex:testAddGivesUpAfterRepeatedGenerationSwitches() + self.dict.add = function(d, k, v, exptime) + local ok, err, forcible = SimpleDict.add(d, k, v, exptime) + if ok and k:match("key_%d+$") then + d:incr("_prefix_gen", 1, 0) + end + return ok, err, forcible + end + + local err = self.key_index:add("key", "eviction_err") + + luaunit.assertStrContains(err, "generation switched 4 times") + self.dict.add = nil + luaunit.assertEquals(self.key_index:list(), {}) +end + +-- A key removed from the new generation after the switch must not come back +-- from its old slot. +function TestKeyIndex:testCompactDropsKeyRemovedAfterSwitch() + churn_slots(self.key_index, 30) + self.key_index:add("doomed", "eviction_err", 100) + local worker2 = require('prometheus_keys').new(self.dict, "_prefix_", 1) + sleep(2) + + on_generation_switch(self.dict, function() + luaunit.assertNil(worker2:remove("doomed")) + end) + self.key_index:remove_expired_keys() + + luaunit.assertEquals(self.key_index:list(), {"permanent"}) + luaunit.assertEquals(worker2:list(), {"permanent"}) + luaunit.assertNil(self.dict:get("_prefix_1_key_2")) +end + +-- A slot written into the old generation after the switch would be deleted +-- with it, so add() registers the key again in the current generation. +function TestKeyIndex:testAddReRegistersKeyAfterGenerationSwitch() + local incr = SimpleDict.incr + self.dict.incr = function(d, k, v, init) + if k == "_prefix_key_count" then + d.incr = nil + d:set("_prefix_gen", 1) + end + return incr(d, k, v, init) + end + + luaunit.assertNil(self.key_index:add("key", "eviction_err")) + + luaunit.assertNil(self.dict:get("_prefix_key_1")) + luaunit.assertEquals(self.dict:get("_prefix_1_key_1"), "key") + luaunit.assertEquals(self.key_index.gen, 1) + luaunit.assertEquals(self.key_index:list(), {"key"}) +end + +-- A compaction that stopped before switching leaves slots of the next +-- generation behind; they must not leak into the next attempt. +function TestKeyIndex:testCompactClearsLeftoverGeneration() + churn_slots(self.key_index, 30) + sleep(2) + self.dict:set("_prefix_1_key_1", "stale1") + self.dict:set("_prefix_1_key_2", "stale2") + + self.key_index:remove_expired_keys() + + luaunit.assertEquals(self.dict:get("_prefix_1_key_count"), 1) + luaunit.assertEquals(self.dict:get("_prefix_1_key_1"), "permanent") + luaunit.assertNil(self.dict:get("_prefix_1_key_2")) + luaunit.assertEquals(self.key_index:list(), {"permanent"}) +end + +-- The generation only grows; if its node is evicted from a full dict, a +-- worker that knows the current generation writes it back. +function TestKeyIndex:testSyncRestoresEvictedGeneration() + churn_slots(self.key_index, 30) + sleep(2) + self.key_index:remove_expired_keys() + self.dict:delete("_prefix_gen") + + self.key_index:sync() + + luaunit.assertEquals(self.dict:get("_prefix_gen"), 1) + luaunit.assertEquals(self.key_index:list(), {"permanent"}) +end + +function TestPrometheus:testCompactionKeepsMetricData() + self.p.key_index.compact_min_slots = 10 + local keep = self.p:counter("keep_total", "Kept", {"id"}) + local churn = self.p:counter("churn_total", "Churned", {"id"}, 1) + keep:inc(5, {"a"}) + for i = 1, 30 do + churn:inc(1, {tostring(i)}) + end + self.p._counter:sync() + local expected = {} + for _, line in ipairs(self.p:metric_data()) do + if not line:find("churn_total", 1, true) then + table.insert(expected, line) + end + end + sleep(2) + + self.p.key_index:remove_expired_keys() + + luaunit.assertEquals(self.dict:get("__ngx_prom__gen"), 1) + luaunit.assertEquals(self.p:metric_data(), expected) + keep:inc(1, {"b"}) + self.p._counter:sync() + luaunit.assertEquals(self.dict:get('keep_total{id="b"}'), 1) + luaunit.assertNotNil(self.p.key_index.index['keep_total{id="b"}']) +end + +-- Timers still pending when a worker exits run once with premature = true; +-- that run must not start a compaction. +function TestKeyIndex:testPrematureTimerDoesNotCompact() + local every = ngx.timer.every + local callback + ngx.timer.every = function(_, cb, _) callback = cb end + local key_index = require('prometheus_keys').new(self.dict, "_prefix_", 1) + ngx.timer.every = every + churn_slots(key_index, 30) + sleep(2) + + callback(true, key_index) + luaunit.assertNil(self.dict:get("_prefix_gen")) + + callback(false, key_index) + luaunit.assertEquals(self.dict:get("_prefix_gen"), 1) +end + +function TestKeyIndex:testExitingWorkerDoesNotStartCompaction() + churn_slots(self.key_index, 30) + sleep(2) + + local written = {} + local set = SimpleDict.set + self.dict.set = function(d, k, v, exptime) + written[#written + 1] = k + return set(d, k, v, exptime) + end + local exiting = ngx.worker.exiting + ngx.worker.exiting = function() return true end + local ok, err = pcall(self.key_index.remove_expired_keys, self.key_index) + ngx.worker.exiting = exiting + self.dict.set = nil + luaunit.assertTrue(ok, err) + + -- nothing was copied, not merely rolled back + luaunit.assertEquals(written, {}) + luaunit.assertNil(self.dict:get("_prefix_gen")) + luaunit.assertNil(self.dict:get("_prefix_compact_lock")) +end + +-- A worker told to exit while copying resumes from its sleep timer and gives +-- up before switching generations, releasing the lock and its copy. +function TestKeyIndex:testWorkerExitingDuringCopyAbortsCompaction() + churn_slots(self.key_index, 30) + sleep(2) + + local get_phase, ngx_sleep, exiting = ngx.get_phase, ngx.sleep, ngx.worker.exiting + local exiting_now, sleeps = false, {} + ngx.get_phase = function() return "timer" end + ngx.sleep = function(s) + table.insert(sleeps, s) + exiting_now = true + end + ngx.worker.exiting = function() return exiting_now end + local ok, err = pcall(self.key_index.remove_expired_keys, self.key_index) + ngx.get_phase, ngx.sleep, ngx.worker.exiting = get_phase, ngx_sleep, exiting + luaunit.assertTrue(ok, err) + + luaunit.assertEquals(sleeps, {0.001}) + luaunit.assertNil(self.dict:get("_prefix_gen")) + luaunit.assertNil(self.dict:get("_prefix_compact_lock")) + luaunit.assertNil(self.dict:get("_prefix_1_key_1")) + luaunit.assertNil(self.dict:get("_prefix_1_key_count")) + luaunit.assertEquals(self.dict:get("_prefix_key_1"), "permanent") + luaunit.assertEquals(self.key_index:list(), {"permanent"}) +end + function TestKeyIndex:testSync() self.key_index:sync() luaunit.assertEquals(ngx.logs, nil)