EndGame/resty/lrucache.lua

-- Copyright (C) Yichun Zhang (agentzh)


local ffi = require "ffi"
local ffi_new = ffi.new
local ffi_sizeof = ffi.sizeof
local ffi_cast = ffi.cast
local ffi_fill = ffi.fill
local ngx_now = ngx.now
local uintptr_t = ffi.typeof("uintptr_t")
local setmetatable = setmetatable
local tonumber = tonumber
local type = type
local new_tab
do
    local ok
    ok, new_tab = pcall(require, "table.new")
    if not ok then
        new_tab = function(narr, nrec) return {} end
    end
end


if string.find(jit.version, " 2.0", 1, true) then
    ngx.log(ngx.ALERT, "use of lua-resty-lrucache with LuaJIT 2.0 is ",
                       "not recommended; use LuaJIT 2.1+ instead")
end


local ok, tb_clear = pcall(require, "table.clear")
if not ok then
    local pairs = pairs
    tb_clear = function (tab)
        for k, _ in pairs(tab) do
            tab[k] = nil
        end
    end
end


-- queue data types
--
-- this queue is a double-ended queue and the first node
-- is reserved for the queue itself.
-- the implementation is mostly borrowed from nginx's ngx_queue_t data
-- structure.

ffi.cdef[[
    typedef struct lrucache_queue_s  lrucache_queue_t;
    struct lrucache_queue_s {
        double             expire;  /* in seconds */
        lrucache_queue_t  *prev;
        lrucache_queue_t  *next;
        uint32_t           user_flags;
    };
]]

local queue_arr_type = ffi.typeof("lrucache_queue_t[?]")
local queue_type = ffi.typeof("lrucache_queue_t")
local NULL = ffi.null


-- queue utility functions

local function queue_insert_tail(h, x)
    local last = h[0].prev
    x.prev = last
    last.next = x
    x.next = h
    h[0].prev = x
end


local function queue_init(size)
    if not size then
        size = 0
    end
    local q = ffi_new(queue_arr_type, size + 1)
    ffi_fill(q, ffi_sizeof(queue_type, size + 1), 0)

    if size == 0 then
        q[0].prev = q
        q[0].next = q

    else
        local prev = q[0]
        for i = 1, size do
          local e = q + i
          e.user_flags = 0
          prev.next = e
          e.prev = prev
          prev = e
        end

        local last = q[size]
        last.next = q
        q[0].prev = last
    end

    return q
end


local function queue_is_empty(q)
    -- print("q: ", tostring(q), "q.prev: ", tostring(q), ": ", q == q.prev)
    return q == q[0].prev
end


local function queue_remove(x)
    local prev = x.prev
    local next = x.next

    next.prev = prev
    prev.next = next

    -- for debugging purpose only:
    x.prev = NULL
    x.next = NULL
end


local function queue_insert_head(h, x)
    x.next = h[0].next
    x.next.prev = x
    x.prev = h
    h[0].next = x
end


local function queue_last(h)
    return h[0].prev
end


local function queue_head(h)
    return h[0].next
end


-- true module stuffs

local _M = {
    _VERSION = '0.10'
}
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"
    end

    local self = {
        hasht = {},
        free_queue = queue_init(size),
        cache_queue = queue_init(),
        key2node = {},
        node2key = {},
        num_items = 0,
        max_items = size,
    }
    return setmetatable(self, mt)
end


function _M.count(self)
    return self.num_items
end


function _M.capacity(self)
    return self.max_items
end


function _M.get(self, key)
    local hasht = self.hasht
    local val = hasht[key]
    if val == nil then
        return nil
    end

    local node = self.key2node[key]

    -- print(key, ": moving node ", tostring(node), " to cache queue head")
    local cache_queue = self.cache_queue
    queue_remove(node)
    queue_insert_head(cache_queue, node)

    if node.expire >= 0 and node.expire < ngx_now() then
        -- print("expired: ", node.expire, " > ", ngx_now())
        return nil, val, node.user_flags
    end

    return val, nil, node.user_flags
end


function _M.delete(self, key)
    self.hasht[key] = nil

    local key2node = self.key2node
    local node = key2node[key]

    if not node then
        return false
    end

    key2node[key] = nil
    self.node2key[ptr2num(node)] = nil

    queue_remove(node)
    queue_insert_tail(self.free_queue, node)
    self.num_items = self.num_items - 1
    return true
end


function _M.set(self, key, value, ttl, flags)
    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

        else
            -- take a free queue node
            node = queue_head(free_queue)
            -- only add count if we are not evicting
            self.num_items = self.num_items + 1
            -- print(key, ": get a new free node: ", tostring(node))
        end

        node2key[ptr2num(node)] = key
        key2node[key] = node
    end

    queue_remove(node)
    queue_insert_head(self.cache_queue, node)

    if ttl then
        node.expire = ngx_now() + ttl
    else
        node.expire = -1
    end

    if type(flags) == "number" and flags >= 0 then
        node.user_flags = flags

    else
        node.user_flags = 0
    end
end


function _M.get_keys(self, max_count, res)
    if not max_count or max_count == 0 then
        max_count = self.num_items
    end

    if not res then
        res = new_tab(max_count + 1, 0) -- + 1 for trailing hole
    end

    local cache_queue = self.cache_queue
    local node2key = self.node2key

    local i = 0
    local node = queue_head(cache_queue)

    while node ~= cache_queue do
        if i >= max_count then
            break
        end

        i = i + 1
        res[i] = node2key[ptr2num(node)]
        node = node.next
    end

    res[i + 1] = nil

    return res
end


function _M.flush_all(self)
    tb_clear(self.hasht)
    tb_clear(self.node2key)
    tb_clear(self.key2node)

    local cache_queue = self.cache_queue
    local free_queue = self.free_queue

    -- splice the cache_queue into free_queue
    if not queue_is_empty(cache_queue) then
        local free_head = free_queue[0]
        local free_last = free_head.prev

        local cache_head = cache_queue[0]
        local cache_first = cache_head.next
        local cache_last = cache_head.prev

        free_last.next = cache_first
        cache_first.prev = free_last

        cache_last.next = free_head
        free_head.prev = cache_last

        cache_head.next = cache_queue
        cache_head.prev = cache_queue
    end
end


return _M