fix(keys): prevent infinite loop in KeyIndex:add() when key_count falls behind occupied slots

prometheus_keys.lua

@@ -7,6 +7,13 @@
 local KeyIndex = {}
 KeyIndex.__index = KeyIndex
 
+-- Upper bound on how far a single add() call may advance key_count while it
+-- repairs a counter that has fallen behind occupied slots (see the comment in
+-- add()). It bounds the worst-case latency of one call; the advanced counter
+-- is shared through the dict, so later calls resume where this one stopped
+-- and the index converges even when far more slots need repairing.
+local MAX_KEY_COUNT_REPAIRS = 1000
+
 
 -- check and remove expired keys
 local function remove_expired_keys(_, self)
@@ -123,6 +130,9 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime)
   end
 
   for _, key in pairs(keys) do
+    local retried = false
+    local repairs = 0
+    local repair_forcible = false
     while true do
       local N = self:sync()
       if self.index[key] ~= nil then
@@ -158,6 +168,14 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime)
           end
         end
         if not expired then
+          if repair_forcible then
+            -- the key was adopted from an occupied slot after repair
+            -- increments that forcibly displaced other entries; report the
+            -- eviction just like the new-slot success path does.
+            return (err_msg_lru_eviction .. "; key index: adopted key after " ..
+              "key_count repair: idx=" .. self.key_prefix .. self.index[key] ..
+              ", key=" .. key)
+          end
           break
         end
       end
@@ -170,14 +188,51 @@ function KeyIndex:add(key_or_keys, err_msg_lru_eviction, exptime)
         if exptime and exptime > 0 then
           self.expire_keys[N] = true
         end
-        if forcible or forcible2 then
+        if forcible or forcible2 or repair_forcible then
           return (err_msg_lru_eviction .. "; key index: add key: idx=" ..
                   self.key_prefix .. N .. ", key=" .. key)
         end
         break
       elseif err ~= "exists" then
         return "Unexpected error adding a key: " .. err
       end
+
+      -- "exists": slot N is already occupied although key_count reported N-1.
+      -- Once per key this can be a benign race with another worker that has
+      -- created slot N but not incremented key_count yet, so retry and let
+      -- sync() pick the new slot up. If it repeats, key_count has fallen
+      -- behind the occupied slots: it is an ordinary shared-dict node, so on
+      -- a full dict it can be LRU-evicted (it is only refreshed when new keys
+      -- are registered, so it goes cold under steady traffic) and incr() then
+      -- re-creates it at 1, far below the surviving slots. Retrying the same
+      -- slot forever would spin the worker at 100% CPU with the shared-dict
+      -- 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
+        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
+          -- the add() error path above, instead of burning the repair budget
+          -- on retries that cannot succeed.
+          return "Unexpected error advancing key_count: " .. incr_err
+        end
+        if forcible3 then
+          -- re-creating an evicted key_count displaced another entry; surface
+          -- it through the LRU-eviction warning on the success path, like
+          -- forcible/forcible2.
+          repair_forcible = true
+        end
+        -- The cap counts attempts, not successes: it exists to guarantee the
+        -- loop terminates.
+        repairs = repairs + 1
+        if repairs >= MAX_KEY_COUNT_REPAIRS then
+          return (err_msg_lru_eviction .. "; key index: key_count fell " ..
+            "behind occupied slots; advanced it by " .. repairs ..
+            " without finding a free slot, dropping key: " .. key)
+        end
+      end
+      retried = true
     end
   end
 end

prometheus_test.lua

@@ -24,12 +24,17 @@ function SimpleDict:add(k, v, exptime)
   if k == "willnotfitk" or v == "willnotfitv" then
     forcible = true
   end
+  self:get(k)  -- prunes the key if it has expired, like ngx.shared.DICT does
+  if self.dict and self.dict[k] then
+    return false, "exists", false  -- match ngx.shared.DICT:add on present keys
+  end
   self:set(k, v, exptime)
   return true, nil, forcible  -- ok, err, forcible
 end
 function SimpleDict:incr(k, v, init)
   local forcible = false
-  if k == "willnotfitk" or v == "willnotfitv" then
+  if k == "willnotfitk" or v == "willnotfitv" or
+     (self.forcible_keys and self.forcible_keys[k]) then
     forcible = true
   end
   if not self.dict[k] then self.dict[k] = {value = init} end
@@ -877,6 +882,96 @@ function TestKeyIndex:testListDeduplicatesStaleSlot()
   luaunit.assertEquals(keys[1], "dup")
 end
 
+-- Regression test for apache/apisix#12275 (workers pinned at 100% CPU).
+-- key_count is an ordinary shared-dict node: it is only refreshed when a new
+-- key is registered, so under steady traffic it goes cold and, once the dict
+-- is full, gets LRU-evicted while higher-numbered slots survive. sync() then
+-- reads it as 0 and add() keeps calling dict:add() on the same occupied slot:
+-- "exists" -> retry -> "exists", forever, without yielding. Before the fix
+-- this call never returned and the worker spun at 100% CPU; the fix advances
+-- key_count past occupied slots so add() terminates and repairs the counter.
+function TestKeyIndex:testAddAfterKeyCountEvicted()
+  for i = 1, 50 do
+    self.key_index:add("key" .. i, "eviction_err")
+  end
+  luaunit.assertEquals(self.dict:get("_prefix_key_count"), 50)
+
+  -- the shared dict force-evicts its LRU-tail node, which is key_count
+  self.dict:delete("_prefix_key_count")
+
+  local err = self.key_index:add("newkey", "eviction_err")
+  luaunit.assertEquals(err, nil)
+  luaunit.assertEquals(self.dict:get("_prefix_key_51"), "newkey")
+  luaunit.assertEquals(self.dict:get("_prefix_key_count"), 51)
+  luaunit.assertEquals(self.key_index.index["newkey"], 51)
+end
+
+-- Same eviction, seen from a freshly started worker (e.g. respawned after a
+-- crash, or after nginx reload) whose local index is empty: it must adopt the
+-- surviving slots while repairing the counter instead of spinning.
+function TestKeyIndex:testAddAfterKeyCountEvictedFreshWorker()
+  for i = 1, 50 do
+    self.key_index:add("key" .. i, "eviction_err")
+  end
+  self.dict:delete("_prefix_key_count")
+
+  local worker2 = require('prometheus_keys').new(self.dict, "_prefix_", 1)
+  local err = worker2:add("newkey", "eviction_err")
+  luaunit.assertEquals(err, nil)
+  luaunit.assertEquals(self.dict:get("_prefix_key_51"), "newkey")
+  luaunit.assertEquals(self.dict:get("_prefix_key_count"), 51)
+  -- surviving slots were adopted into the local index along the way
+  luaunit.assertEquals(worker2.index["key7"], 7)
+  luaunit.assertEquals(worker2.index["newkey"], 51)
+end
+
+-- One add() call bounds its repair work (MAX_KEY_COUNT_REPAIRS = 1000) and
+-- degrades by dropping the sample, like the existing dict-full path. The
+-- advanced counter is shared, so a later call resumes and converges.
+function TestKeyIndex:testKeyCountRepairIsBounded()
+  local keys = {}
+  for i = 1, 1500 do
+    keys[i] = "key" .. i
+  end
+  self.key_index:add(keys, "eviction_err")
+  self.dict:delete("_prefix_key_count")
+
+  local err = self.key_index:add("newkey", "eviction_err")
+  luaunit.assertStrContains(err, "key_count fell behind occupied slots")
+  -- bounded progress was made and persisted for the next call
+  luaunit.assertEquals(self.dict:get("_prefix_key_count"), 1000)
+
+  err = self.key_index:add("newkey", "eviction_err")
+  luaunit.assertEquals(err, nil)
+  luaunit.assertEquals(self.dict:get("_prefix_key_1501"), "newkey")
+  luaunit.assertEquals(self.dict:get("_prefix_key_count"), 1501)
+end
+
+-- If the repair increments forcibly displaced other entries and the requested
+-- key then turns out to already exist (it gets adopted from an occupied slot
+-- during the repair crawl), the eviction must still be reported through the
+-- LRU-eviction error, exactly like the new-slot success path reports
+-- forcible/forcible2.
+function TestKeyIndex:testKeyCountRepairReportsForcibleEviction()
+  for i = 1, 50 do
+    self.key_index:add("key" .. i, "eviction_err")
+  end
+  self.dict:delete("_prefix_key_count")
+  -- re-creating/incrementing key_count now displaces other entries
+  self.dict.forcible_keys = {["_prefix_key_count"] = true}
+
+  -- a worker with an empty local index registers a key that already exists in
+  -- the dict at a slot ahead of the (rewound) counter: the repair crawl adopts
+  -- it and exits through the existing-key break.
+  local worker2 = require('prometheus_keys').new(self.dict, "_prefix_", 1)
+  local err = worker2:add("key30", "eviction_err")
+  luaunit.assertStrContains(err,
+    "eviction_err; key index: adopted key after key_count repair: " ..
+    "idx=_prefix_key_30, key=key30")
+  luaunit.assertEquals(worker2.index["key30"], 30)
+  luaunit.assertEquals(self.dict:get("_prefix_key_count"), 30)
+end
+
 function TestKeyIndex:testSync()
   self.key_index:sync()
   luaunit.assertEquals(ngx.logs, nil)