mysql-proxy 读写分离脚本一 KDr2

--[[
--
-- author : KDr2 <kdr2@163.com>
-- version 0.01
-- SYNOPSIS:
---  1.维护了一个连接池
---  2.读写分离，简单的将select开头的语句放到slave上执行
---  3.事务支持，所有事务放到master上执行，事务中不更改连接
---  4.简单日志
--
--]]

--- config vars
local min_idle_connections = 1
local max_idle_connections = 2
local log_level=1
--local encoding="utf8"
--- end of config

-- 事务标识，在事务内不归还连接
local transaction_flags={}
setmetatable(transaction_flags,{__index=function() return 0 end})

-- log system
log={
   level={debug=1,info=2,warn=3,error=4},
   funcs={"debug","info","warn","error"},
}
function log.log(level,m)
   if level >= log_level then
      local msg="[" .. os.date("%Y-%m-%d %X") .."] ".. log.funcs[level] .. ": " .. tostring(m)
      print(msg) -- TODO  write msg into a log file.
   end
end
for i,v in ipairs(log.funcs) do
   log[v]=function(m) log.log(log.level[v],m) end
end

-- connect to server
function connect_server()
   log.info(" starting connect_server ... ")
   local least_idle_conns_ndx = 0
   local least_idle_conns = 0
  
   for i = 1, #proxy.backends do
      local s = proxy.backends[i]
      local pool = s.pool
      local cur_idle = pool.users[""].cur_idle_connections

      log.debug("[".. s.address .."].connected_clients = " .. s.connected_clients)
      log.debug("[".. s.address .."].idling_connections = " .. cur_idle)
      log.debug("[".. s.address .."].type = " .. s.type)
      log.debug("[".. s.address .."].state = " .. s.state)

      if s.state ~= proxy.BACKEND_STATE_DOWN then
         -- try to connect to each backend once at least
         if cur_idle == 0 then
            proxy.connection.backend_ndx = i
            log.info("server [".. proxy.backends[i].address .."] open new connection")
            return
         end
         -- try to open at least min_idle_connections
         if least_idle_conns_ndx == 0 or
            ( cur_idle < min_idle_connections and
              cur_idle < least_idle_conns ) then
            least_idle_conns_ndx = i
            least_idle_conns = cur_idle
         end
      end
   end

   if least_idle_conns_ndx > 0 then
      proxy.connection.backend_ndx = least_idle_conns_ndx
   end
  
   if proxy.connection.backend_ndx > 0 then
      local s = proxy.backends[proxy.connection.backend_ndx]
      local pool = s.pool
      local cur_idle = pool.users[""].cur_idle_connections

      if cur_idle >= min_idle_connections then
         -- we have 4 idling connections in the pool, that's good enough
         log.debug("using pooled connection from: " .. proxy.connection.backend_ndx)
         return proxy.PROXY_IGNORE_RESULT
      end
   end
   -- open a new connection
   log.info("opening new connection on: " .. proxy.backends[proxy.connection.backend_ndx].address)
end

---

-- auth.packet is the packet
function read_auth_result( auth )
   if auth.packet:byte() == proxy.MYSQLD_PACKET_OK then
      -- 连接正常
      proxy.connection.backend_ndx = 0
   elseif auth.packet:byte() == proxy.MYSQLD_PACKET_EOF then
      -- we received either a
      -- * MYSQLD_PACKET_ERR and the auth failed or
      -- * MYSQLD_PACKET_EOF which means a OLD PASSWORD (4.0) was sent
      log.error("(read_auth_result) ... not ok yet");
   elseif auth.packet:byte() == proxy.MYSQLD_PACKET_ERR then
      log.error("auth failed!")
   end
end

---
-- read/write splitting
function read_query( packet )
   log.debug("[read_query]")
   log.debug("authed backend = " .. proxy.connection.backend_ndx)
   log.debug("used db = " .. proxy.connection.client.default_db)

   if packet:byte() == proxy.COM_QUIT then
      proxy.response = {
         type = proxy.MYSQLD_PACKET_OK,
      }
      return proxy.PROXY_SEND_RESULT
   end

   if proxy.connection.backend_ndx == 0 then
      local is_read=(string.upper(packet:sub(2))):match("^SELECT")
      local target_type=proxy.BACKEND_TYPE_RW
      if is_read then target_type=proxy.BACKEND_TYPE_RO end
      for i = 1, #proxy.backends do
         local s = proxy.backends[i]
         local pool = s.pool
         local cur_idle = pool.users[proxy.connection.client.username].cur_idle_connections
        
         if cur_idle > 0 and
            s.state ~= proxy.BACKEND_STATE_DOWN and
            s.type == target_type then
            proxy.connection.backend_ndx = i
            break
         end
      end
   end
   -- sync the client-side default_db with the server-side default_db
   if proxy.connection.server and proxy.connection.client.default_db ~= proxy.connection.server.default_db then
      local server_db=proxy.connection.server.default_db
      local client_db=proxy.connection.client.default_db
      local default_db= (#client_db > 0) and client_db or server_db
      if #default_db > 0 then
         proxy.queries:append(2, string.char(proxy.COM_INIT_DB) .. default_db)
         --proxy.queries:append(2, string.char(proxy.COM_QUERY) .. "set names '" .. encoding .."'")
         log.info("change database to " .. default_db);
      end
   end
   if proxy.connection.backend_ndx > 0 then
      log.debug("Query[" .. packet:sub(2) .. "] Target is [" .. proxy.backends[proxy.connection.backend_ndx].address .."]")
   end
   proxy.queries:append(1, packet)
   return proxy.PROXY_SEND_QUERY
end

---
-- as long as we are in a transaction keep the connection
-- otherwise release it so another client can use it
function read_query_result( inj )
   local res      = assert(inj.resultset)
   local flags    = res.flags

   if inj.id ~= 1 then
      -- ignore the result of the USE <default_db>
      return proxy.PROXY_IGNORE_RESULT
   end
   is_in_transaction = flags.in_trans

   if flags.in_trans then
      transaction_flags[proxy.connection.server.thread_id] = transaction_flags[proxy.connection.server.thread_id] + 1
   elseif inj.query:sub(2):lower():match("^%s*commit%s*$") or inj.query:sub(2):lower():match("^%s*rollback%s*$") then
      transaction_flags[proxy.connection.server.thread_id] = transaction_flags[proxy.connection.server.thread_id] - 1
      if transaction_flags[proxy.connection.server.thread_id] < 0 then transaction_flags[proxy.connection.server.thread_id] = 0 end
   end
  
   log.debug("transaction res : " .. tostring(transaction_flags[proxy.connection.server.thread_id]));
   if transaction_flags[proxy.connection.server.thread_id]==0 or transaction_flags[proxy.connection.server.thread_id] == nil then
      -- isnot in a transaction, need to release the backend
      proxy.connection.backend_ndx = 0
   end
end

---
-- close the connections if we have enough connections in the pool
--
-- @return nil - close connection
-- IGNORE_RESULT - store connection in the pool
function disconnect_client()
   log.debug("[disconnect_client]")
   if proxy.connection.backend_ndx == 0 then
      for i = 1, #proxy.backends do
         local s = proxy.backends[i]
         local pool = s.pool
         local cur_idle = pool.users[proxy.connection.client.username].cur_idle_connections
        
         if s.state ~= proxy.BACKEND_STATE_DOWN and
            cur_idle > max_idle_connections then
            -- try to disconnect a backend
            proxy.connection.backend_ndx = i
            log.info("[".. proxy.backends[i].address .."] closing connection, idling: " .. cur_idle)
            return
         end
      end
      return proxy.PROXY_IGNORE_RESULT
   end
end