add increment method (incr)

openresty · Nov 25, 2015 · 968e280 · 968e280
1 parent b1beb00
commit 968e280
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 31 deletions.
diff --git a/lib/resty/lrucache.lua b/lib/resty/lrucache.lua
@@ -111,6 +111,43 @@ local function queue_head(h)
 end
 
 
+local function ptr2num(ptr)
+    return tonumber(ffi_cast(uintptr_t, ptr))
+end
+
+local function handle_lru(self, key)
+    local key2node = self.key2node
+    local node = key2node[key]
+    if not node then
+        local free_queue = self.free_queue
+        local node2key = self.node2key
+
+        if queue_is_empty(free_queue) then
+            -- evict the least recently used key
+            -- assert(not queue_is_empty(self.cache_queue))
+            node = queue_last(self.cache_queue)
+
+            local oldkey = node2key[ptr2num(node)]
+            -- print(key, ": evicting oldkey: ", oldkey, ", oldnode: ",
+            --         tostring(node))
+            if oldkey then
+                self.hasht[oldkey] = nil
+                key2node[oldkey] = nil
+            end
+
+        else
+            -- take a free queue node
+            node = queue_head(free_queue)
+            -- print(key, ": get a new free node: ", tostring(node))
+        end
+
+        node2key[ptr2num(node)] = key
+        key2node[key] = node
+    end
+
+    return node
+end
+
 -- true module stuffs
 
 local _M = {
@@ -119,11 +156,6 @@ local _M = {
 local mt = { __index = _M }
 
 
-local function ptr2num(ptr)
-    return tonumber(ffi_cast(uintptr_t, ptr))
-end
-
-
 function _M.new(size)
     if size < 1 then
         return nil, "size too small"
@@ -181,40 +213,31 @@ function _M.delete(self, key)
     return true
 end
 
-
 function _M.set(self, key, value, ttl)
     local hasht = self.hasht
     hasht[key] = value
 
-    local key2node = self.key2node
-    local node = key2node[key]
-    if not node then
-        local free_queue = self.free_queue
-        local node2key = self.node2key
-
-        if queue_is_empty(free_queue) then
-            -- evict the least recently used key
-            -- assert(not queue_is_empty(self.cache_queue))
-            node = queue_last(self.cache_queue)
-
-            local oldkey = node2key[ptr2num(node)]
-            -- print(key, ": evicting oldkey: ", oldkey, ", oldnode: ",
-            --         tostring(node))
-            if oldkey then
-                hasht[oldkey] = nil
-                key2node[oldkey] = nil
-            end
+    local node = handle_lru(self,key)
+    queue_remove(node)
+    queue_insert_head(self.cache_queue, node)
 
-        else
-            -- take a free queue node
-            node = queue_head(free_queue)
-            -- print(key, ": get a new free node: ", tostring(node))
-        end
+    if ttl then
+        node.expire = ngx_now() + ttl
+    else
+        node.expire = -1
+    end
+end
 
-        node2key[ptr2num(node)] = key
-        key2node[key] = node
+function _M.incr(self, key, ttl, by)
+    local hasht = self.hasht
+    local value = hasht[key]
+    if value == nil then
+        return value
     end
+    value = value + by
+    hasht[key] = value
 
+    local node = handle_lru(self,key)
     queue_remove(node)
     queue_insert_head(self.cache_queue, node)
 

diff --git a/lib/resty/lrucache/pureffi.lua b/lib/resty/lrucache/pureffi.lua
@@ -530,5 +530,33 @@ function _M.set(self, key, value, ttl)
     end
 end
 
+function _M.incr(self, key, ttl, by)
+    by = by or 1
+
+    if type(key) ~= "string" then
+        key = tostring(key)
+    end
+
+    local node_id = find_key(self, key)
+    local node
+    if not node_id then
+        return nil
+    end
+
+    node = self.node_v + node_id
+    local value = self.val_v[node_id] + by
+    self.val_v[node_id] = value
+
+    queue_remove(node)
+    queue_insert_head(self.cache_queue, node)
+
+    if ttl then
+        node.expire = ngx_now() + ttl
+    else
+        node.expire = -1
+    end
+
+    return value
+end
 
 return _M