service/redis_pool.lua

redis_pool.lua

local skynet = require "skynet"
local redis = require "redis"
local parse_redis_url = require "redis_url".parse

local sort = table.sort
local unpack = table.unpack

local conf
local pool = {}
local named = {}
local pointer = 1

local CALL = {}
local CAST = {}

local function sortbycounter(a, b)
  return a.__pool_counter < b.__pool_counter
end

function CALL.open(newconf)
  conf = newconf
  if conf then
    conf.pool_size = conf.pool_size or 5
  else
    conf = parse_redis_url(skynet.getenv("redis_url"))
    conf.pool_size = tonumber(skynet.getenv("redis_pool_size")) or 5
  end

  pool = {
    n = conf.pool_size
  }
  for i = 1, pool.n do
    local db = redis.connect(conf)
    db.__pool_handle = i
    db.__pool_counter = 0
    pool[i] = db
  end
end

function CAST.close()
  skynet.exit()
end

-- Checkout a connection, returns connection handle.
function CALL.checkout()
  -- find the least used one in next 3 connections
  local selected
  if pool.n <= 3 then
    selected = pool
  else
    selected = { nil, nil, nil }
    for i = 1, 3 do
      if pointer > pool.n then
        pointer = 1
      end

      selected[i] = pool[pointer]
      pointer = pointer + 1
    end
  end

  sort(selected, sortbycounter)
  local db = selected[1]
  pointer = db.__pool_handle + 1
  db.__pool_counter = db.__pool_counter + 1
  return db.__pool_handle
end

function CALL.named(name)
  assert(named[name] == nil, "Duplicate redis connection name")
  named[name] = redis.connect(conf)
  return name
end

function CAST.checkin(handle)
  if type(handle) == "number" then
    local db = assert(pool[handle], "Unknown redis connection")
    db.__pool_counter = db.__pool_counter - 1
  else
    local db = assert(named[handle], "Unknown redis connection")
    db:disconnect()
    named[handle] = nil
  end
end

skynet.start(function()
  skynet.dispatch("lua", function(session, address, cmd, ...) -- luacheck: no unused args
    local f
    if session ~= 0 then
      local db = pool[cmd] or named[cmd]
      if db then
        local args = {...}
        f = db[args[1]]
        args[1] = db
        return skynet.ret(skynet.pack(f(unpack(args))))
      else
        f = assert(CALL[cmd], "unknown redis_pool call command " .. cmd)
      end
    else
      f = assert(CAST[cmd], "unknown redis_pool cast command " .. cmd)
      return f(...)
    end

    skynet.ret(skynet.pack(f(...)))
  end)
end)