--[[ $%BEGINLICENSE%$
 Copyright (C) 2007-2008 MySQL AB, 2008 Sun Microsystems, Inc

 This program is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation; version 2 of the License.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

 $%ENDLICENSE%$ --]]

---
-- slave aware load-balancer
--
-- * take slaves which are too far behind out of the rotation
-- * when we bring back the slave make sure it doesn't get 
--   behind again too early (slow start)
-- * fallback to the master of all slaves are too far behind
-- * check if the backends are really configures as slave
-- * check if the io-thread is running at all
--
-- it is unclear how to 
-- - handle requests sent to a backend with a io-thread off
-- - should we just close the connection to the client
--   and let it retry ?
-- - how to activate the backend again ?
-- - if a slave is lagging, we don't send queries to it
--   how do we figure out it is back again ?

require("proxy.auto-config")
-- require("proxy.strict")

local commands    = require("proxy.commands")

if not proxy.global.config.ro_balance then
	proxy.global.config.ro_balance = { 
		max_seconds_lag = 10,      -- 10 seconds
		max_bytes_lag = 10 * 1024, -- 10k
		check_timeout = 2,
		is_debug = false
	}
end

if not proxy.global.lag then
	proxy.global.lag = { }
end

local config      = proxy.global.config.ro_balance
local backend_lag = proxy.global.lag

---
-- pick a backend server
--
-- we take the lag into account
-- * ignore backends which aren't configured
-- * ignore backends which don't have a io-thread running
-- * ignore backends that are too far behind
function connect_server()
	local fallback_ndx
	local slave_ndx 
	local unknown_slave_ndx 
	local slave_bytes_lag 

	for b_ndx = 1, #proxy.global.backends do
		local backend = proxy.global.backends[b_ndx]

		if backend.state ~= proxy.BACKEND_STATE_DOWN then
			if backend.type == proxy.BACKEND_TYPE_RW then
				-- the fallback
				fallback_ndx = b_ndx
			else
				--- 
				-- connect to the backends first we don't know yet
				if not backend_lag[backend.dst.name] then
					unknown_slave_ndx = b_ndx
				elseif backend_lag[backend.dst.name].state == "running" then
					if not slave_bytes_lag then
						slave_ndx = b_ndx
						slave_bytes_lag = backend_lag[backend.dst.name].slave_bytes_lag
					elseif backend_lag[backend.dst.name].slave_bytes_lag < slave_bytes_lag then
						slave_ndx = b_ndx
						slave_bytes_lag = backend_lag[backend.dst.name].slave_bytes_lag
					end
				end
			end
		end
	end

	proxy.connection.backend_ndx = unknown_slave_ndx or slave_ndx or fallback_ndx

	if config.is_debug then
		print("(connect-server) using backend: " .. proxy.global.backends[proxy.connection.backend_ndx].dst.name)
	end
end

function read_query(packet)
	local cmd      = commands.parse(packet)
	local ret      = proxy.global.config:handle(cmd)
	if ret then return ret end

	--
	-- check the backend periodicly for its lag
	--
	
	-- translate the backend_ndx into its dst.name
	local backend_addr = proxy.global.backends[proxy.connection.backend_ndx].dst.name
	backend_lag[backend_addr] = backend_lag[backend_addr] or {
		state = "unchecked"
	}

	local now = os.time()

	if config.is_debug then
		print("(read_query) we are on backend: ".. backend_addr .. 
			" in state: " .. backend_lag[backend_addr].state)
	end

	if backend_lag[backend_addr].state == "unchecked" or 
	   (backend_lag[backend_addr].state == "running" and 
	    now - backend_lag[backend_addr].check_ts > config.check_timeout) then
		if config.is_debug then
			print("(read-query) unchecked, injecting a SHOW SLAVE STATUS")
		end
		proxy.queries:append(2, string.char(proxy.COM_QUERY) .. "SHOW SLAVE STATUS", { resultset_is_needed=true }  )
		proxy.queries:append(1, packet, { resultset_is_needed=true })

		return proxy.PROXY_SEND_QUERY
	elseif proxy.global.backends[proxy.connection.backend_ndx].type == proxy.BACKEND_TYPE_RW or  -- master
	       backend_lag[backend_addr].state == "running" then -- good slave
		-- pass through
		return
	else
		-- looks like this is a bad backend
		-- let's get the client to connect to another backend
		--
		-- ... by closing the connection
		proxy.response = {
			type = proxy.MYSQLD_PACKET_ERR,
			errmsg = "slave-state is " .. backend_lag[backend_addr].state,
			sqlstate = "08S01"
		}

		proxy.queries:reset()

		-- and now we have to tell the proxy to close the connection
		--
		proxy.connection.connection_close = true

		return proxy.PROXY_SEND_RESULT
	end
end

function read_query_result(inj)
	-- pass through the client query
	if inj.id == 1 then
		return 
	end

	---
	-- parse the SHOW SLAVE STATUS
	--
	-- what can happen ?
	-- * no permissions (ERR)
	-- * if not slave, result is empty
	local res = inj.resultset
	local fields = res.fields
	local show_slave_status = {}

	-- turn the resultset into local hash
	for row in res.rows do
		for field_id, field in pairs(row) do
			show_slave_status[fields[field_id].name] = tostring(field)
			if config.is_debug then
				print(("[%d] '%s' = '%s'"):format(field_id, fields[field_id].name, tostring(field)))
			end
		end
	end

	local backend_addr = proxy.global.backends[proxy.connection.backend_ndx].dst.name
	backend_lag[backend_addr].check_ts = os.time()
	backend_lag[backend_addr].state = nil

	if not show_slave_status["Master_Host"] then
		-- this backend is not a slave
		backend_lag[backend_addr].state = "noslave"
		return proxy.PROXY_IGNORE_RESULT
	end

	if show_slave_status["Master_Log_File"] == show_slave_status["Relay_Master_Log_File"] then
		-- ok, we use the same relay-log for reading and writing
		backend_lag[backend_addr].slave_bytes_lag = tonumber(show_slave_status["Exec_Master_Log_Pos"]) - tonumber(show_slave_status["Read_Master_Log_Pos"])
	else
		backend_lag[backend_addr].slave_bytes_lag = nil
	end

	if show_slave_status["Seconds_Behind_Master"] then
		backend_lag[backend_addr].seconds_lag = tonumber(show_slave_status["Seconds_Behind_Master"])
	else
		backend_lag[backend_addr].seconds_lag = nil
	end

	if show_slave_status["Slave_IO_Running"] == "No" then
		backend_lag[backend_addr].state = "noiothread"
	end

	backend_lag[backend_addr].state = backend_lag[backend_addr].state or "running"

	return proxy.PROXY_IGNORE_RESULT
end