From 3ab49352f5a82911812f998bb73e0bd3e9a2b8b2 Mon Sep 17 00:00:00 2001 From: bwang Date: Mon, 14 Sep 2026 20:02:26 +0800 Subject: [PATCH 1/3] fix(keys): compact the key index so dead slots do not accumulate KeyIndex never reuses slots: a key that expires or is removed leaves its slot number behind and key_count only grows. Every full sync walks 0..key_count synchronously, and full syncs happen whenever a worker starts and in every worker after any delete_count bump, which includes re-adding an expired key. With label churn and an exptime the cost grows without bound. remove_expired_keys() now compacts the index once key_count reaches 10000 slots and at least twice the live keys. Under a shared-dict lock it copies the live keys, with their remaining ttl, into a new generation of slots and switches the generation. It then reconciles with what other workers did to the old generation before they switched: renewed expiries are carried over, removed keys are dropped and slots written past the copied range are registered. Finally the old slots are deleted. Workers switch on their next sync(); add() re-registers a key whose slot landed in the old generation after the switch. Generation 0 keeps the existing key names. Compaction yields on a sleep timer, because a posted ngx.sleep(0) event does not keep an exiting worker alive. It is not started by the premature timer run or an exiting worker, is abandoned before the switch when the worker starts exiting, and completes without yielding after the switch. --- prometheus_keys.lua | 311 ++++++++++++++++++++++++++++++++++++++++--- prometheus_test.lua | 318 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 609 insertions(+), 20 deletions(-) diff --git a/prometheus_keys.lua b/prometheus_keys.lua index faeaa66..62ddcee 100644 --- a/prometheus_keys.lua +++ b/prometheus_keys.lua @@ -14,32 +14,83 @@ KeyIndex.__index = KeyIndex -- and the index converges even when far more slots need repairing. local MAX_KEY_COUNT_REPAIRS = 1000 +-- 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 +116,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,6 +211,7 @@ 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 while true do local N = self:sync() if self.index[key] ~= nil then @@ -194,19 +259,27 @@ 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) elseif err ~= "exists" then return "Unexpected error adding a key: " .. err end @@ -223,7 +296,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 +319,9 @@ 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 end end end @@ -256,6 +331,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 +350,199 @@ 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 = {}, {}, {}, {} + 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] = key, M, i + 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 + + -- Until they switched, other workers renewed and removed keys in the old + -- generation: carry over the latest expiry, and drop copies of keys that + -- their old slot no longer holds. + local dropped = false + for j, i in pairs(origin) do + local key = keys[j] + local old_slot, new_slot = from_prefix .. i, to_prefix .. j + if self.dict:get(old_slot) == key then + local ttl = self.dict:ttl(old_slot) + if ttl and ttl > 0 and not self.dict:expire(new_slot, ttl) then + self.dict:add(new_slot, key, ttl) + end + else + if self.dict:get(new_slot) == key 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..fbaa096 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,318 @@ 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 + +-- 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) From 5c017f9ff78d2df66543cd20ebcec06064215d14 Mon Sep 17 00:00:00 2001 From: bwang Date: Tue, 15 Sep 2026 10:21:14 +0800 Subject: [PATCH 2/3] fix(keys): keep slots renewed after a compaction switch - Reconcile the copy by comparing each slot's ttl with the one read while copying - Never shorten a new-generation slot, and keep it when only it was renewed - Drop keys removed from the new generation instead of re-adding them - Add regression tests for renewals and removals after the switch --- prometheus_keys.lua | 33 +++++++++++++------- prometheus_test.lua | 76 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 12 deletions(-) diff --git a/prometheus_keys.lua b/prometheus_keys.lua index 62ddcee..0e523a3 100644 --- a/prometheus_keys.lua +++ b/prometheus_keys.lua @@ -434,7 +434,7 @@ function KeyIndex:compact() local to_prefix, to_count = generation_names(self.prefix, to_gen) self:clear_generation(to_gen) - local keys, index, expire_keys, origin = {}, {}, {}, {} + local keys, index, expire_keys, origin, copied_ttl = {}, {}, {}, {}, {} local M = 0 local function abort(msg) self:delete_slots(to_prefix, 1, M) @@ -456,7 +456,7 @@ function KeyIndex:compact() if not ok or forcible then return abort(err or "copying a slot evicted other entries") end - keys[M], index[key], origin[M] = key, M, i + keys[M], index[key], origin[M], copied_ttl[M] = key, M, i, ttl if exptime then expire_keys[M] = true end @@ -485,20 +485,29 @@ function KeyIndex:compact() self:use_generation(to_gen) self.keys, self.index, self.expire_keys, self.last = keys, index, expire_keys, M - -- Until they switched, other workers renewed and removed keys in the old - -- generation: carry over the latest expiry, and drop copies of keys that - -- their old slot no longer holds. + -- 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 = keys[j] + local key, copied = keys[j], copied_ttl[j] local old_slot, new_slot = from_prefix .. i, to_prefix .. j - if self.dict:get(old_slot) == key then - local ttl = self.dict:ttl(old_slot) - if ttl and ttl > 0 and not self.dict:expire(new_slot, ttl) then - self.dict:add(new_slot, key, ttl) + 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 - else - if self.dict:get(new_slot) == key then + 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 diff --git a/prometheus_test.lua b/prometheus_test.lua index fbaa096..8526b06 100644 --- a/prometheus_test.lua +++ b/prometheus_test.lua @@ -1216,6 +1216,82 @@ function TestKeyIndex:testCompactCarriesRenewalDuringCopy() 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 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() From f8f40efce861d3a4436acfb99d26d3584a72c683 Mon Sep 17 00:00:00 2001 From: bwang Date: Tue, 15 Sep 2026 14:35:15 +0800 Subject: [PATCH 3/3] fix(keys): renew keys in the current generation after a switch - Re-check the generation after renewing an existing slot and renew the key again in the current generation when a compaction switched it meanwhile - Treat only a higher generation as a switch, so an evicted gen key does not make add() loop - Bound the generation switches add() follows per key --- prometheus_keys.lua | 22 +++++++++++++- prometheus_test.lua | 74 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/prometheus_keys.lua b/prometheus_keys.lua index 0e523a3..6ae3315 100644 --- a/prometheus_keys.lua +++ b/prometheus_keys.lua @@ -14,6 +14,10 @@ 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 @@ -212,7 +216,12 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime) 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 @@ -247,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 @@ -264,7 +282,7 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime) if ok then local _, _, forcible2 = self.dict:incr(self.key_count, 1, 0) slot_forcible = slot_forcible or forcible or forcible2 - if (self.dict:get(self.gen_key) or 0) == self.gen then + 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 @@ -280,6 +298,7 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime) -- 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 @@ -322,6 +341,7 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime) -- 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 diff --git a/prometheus_test.lua b/prometheus_test.lua index 8526b06..a3211ac 100644 --- a/prometheus_test.lua +++ b/prometheus_test.lua @@ -1274,6 +1274,80 @@ function TestKeyIndex:testCompactKeepsRenewalAfterOldSlotExpired() 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()