introduce socket.channel, and rewrite redis/mongo driver

lualib/socket.lua

@@ -93,11 +93,14 @@ end
 
 -- SKYNET_SOCKET_TYPE_ERROR = 5
 socket_message[5] = function(id)
-	print("error on ", id)
 	local s = socket_pool[id]
 	if s == nil then
+		print("socket: error on unknown", id)
 		return
 	end
+	if s.connected then
+		print("socket: error on", id)
+	end
 	s.connected = false
 	wakeup(s)
 end
@@ -146,6 +149,18 @@ function socket.start(id, func)
 	return connect(id, func)
 end
 
+local function clear_socket(id)
+	local s = socket_pool[id]
+	if s then
+		if s.buffer then
+			driver.clear(s.buffer,buffer_pool)
+		end
+		if s.connected then
+			driver.close(id)
+		end
+	end
+end
+
 function socket.close(id)
 	local s = socket_pool[id]
 	if s == nil then
@@ -163,10 +178,9 @@ function socket.close(id)
 		else
 			suspend(s)
 		end
+		s.connected = false
 	end
-	if s.buffer then
-		driver.clear(s.buffer,buffer_pool)
-	end
+	clear_socket(id)
 	assert(s.lock_set == nil or next(s.lock_set) == nil)
 	socket_pool[id] = nil
 end
@@ -235,6 +249,17 @@ function socket.readline(id, sep)
 	end
 end
 
+function socket.block(id)
+	local s = socket_pool[id]
+	if not s or not s.connected then
+		return false
+	end
+	assert(not s.read_required)
+	s.read_required = 0
+	suspend(s)
+	return s.connected
+end
+
 socket.write = assert(driver.send)
 
 function socket.invalid(id)
@@ -251,10 +276,10 @@ function socket.lock(id)
 		lock_set = {}
 		s.lock = lock_set
 	end
-	local co = coroutine.running()
 	if #lock_set == 0 then
-		lock_set[1] = co
+		lock_set[1] = true
 	else
+		local co = coroutine.running()
 		table.insert(lock_set, co)
 		skynet.wait()
 	end
@@ -263,12 +288,9 @@ end
 function socket.unlock(id)
 	local s = socket_pool[id]
 	assert(s)
-	local lock_set = s.lock
-	assert(lock_set)
-	local co = coroutine.running()
-	assert(lock_set[1] == co)
+	local lock_set = assert(s.lock)
 	table.remove(lock_set,1)
-	co = lock_set[1]
+	local co = lock_set[1]
 	if co then
 		skynet.wakeup(co)
 	end
@@ -284,4 +306,251 @@ function socket.abandon(id)
 	socket_pool[id] = nil
 end
 
+-- channel support auto reconnect , and capture socket error in request/response transaction
+-- { host = "", port = , auth = function(so) , response = function(so) session, data }
+
+local channel = {}
+local channel_socket = {}
+local channel_meta = { __index = channel }
+local channel_socket_meta = {
+	__index = channel_socket,
+	__gc = function(cs)
+		local fd = cs[1]
+		cs[1] = false
+		if fd then
+			clear_socket(fd)
+		end
+	end
+}
+local socket_error = channel_socket	-- alias for error object
+
+function socket.channel(desc)
+	local c = {
+		__host = assert(desc.host),
+		__port = assert(desc.port),
+		__auth = desc.auth,
+		__response = desc.response,
+		__request = {},	-- request seq { response func or session }
+		__thread = {}, -- coroutine seq or session->coroutine map
+		__result = {}, -- response result { coroutine -> result }
+		__result_data = {},
+		__connecting = {},
+		__sock = false,
+		__closed = false,
+	}
+
+	return setmetatable(c, channel_meta)
+end
+
+local function close_channel_socket(self)
+	if self.__sock then
+		local so = self.__sock
+		self.__sock = false
+		socket.close(so[1])
+	end
+end
+
+local function wakeup_all(self, errmsg)
+	for i = 1, #self.__thread do
+		local co = self.__thread[i]
+		self.__thread[i] = nil
+		self.__result[co] = socket_error
+		self.__result_data[co] = errmsg
+		skynet.wakeup(co)
+	end
+end
+
+local function dispatch_response(self)
+	local response = self.__response
+	if response then
+		-- response() return session
+		while self.__sock do
+			local ok , session, result_ok, result_data = pcall(response, self.__sock)
+			if ok then
+				local co = self.__thread[session]
+				self.__thread[session] = nil
+				if co then
+					self.__result[co] = result_ok
+					self.__result_data[co] = result_data
+					skynet.wakeup(co)
+				else
+					print("socket: unknown session :", session)
+				end
+			else
+				close_channel_socket(self)
+				local errormsg
+				if session ~= socket_error then
+					errormsg = session
+				end
+				for k,co in pairs(self.__thread) do
+					-- throw error (errormsg)
+					self.__thread[k] = nil
+					self.__result[co] = socket_error
+					self.__result_data[co] = errormsg
+					skynet.wakeup(co)
+				end
+			end
+		end
+	else
+		-- pop response function from __request
+		while self.__sock do
+			local func = table.remove(self.__request, 1)
+			if func == nil then
+				if not socket.block(self.__sock[1]) then
+					close_channel_socket(self)
+					wakeup_all(self)
+				end
+			else
+				local ok, result_ok, result_data = pcall(func, self.__sock)
+				if ok then
+					local co = table.remove(self.__thread, 1)
+					self.__result[co] = result_ok
+					self.__result_data[co] = result_data
+					skynet.wakeup(co)
+				else
+					close_channel_socket(self)
+					local errmsg
+					if result ~= socket_error then
+						errmsg = result_ok
+					end
+					wakeup_all(self, errmsg)
+				end
+			end
+		end
+	end
+end
+
+local function try_connect(self)
+	assert(not self.__sock)
+	local fd = socket.open(self.__host, self.__port)
+	if not fd then
+		return
+	end
+	self.__sock = setmetatable( {fd} , channel_socket_meta )
+	skynet.fork(dispatch_response, self)
+
+	if self.__auth then
+		local ok , message = pcall(self.__auth, self)
+		if not ok then
+			close_channel_socket(self)
+			if message ~= socket_error then
+				error(message)
+			end
+		end
+	end
+end
+
+function channel:connect()
+	if self.__sock then
+		return true
+	end
+	if self.__closed then
+		self.__closed = false
+	end
+	if #self.__connecting > 0 then
+		-- connecting in other coroutine
+		local co = coroutine.running()
+		table.insert(self.__connecting, co)
+		skynet.wait()
+	else
+		self.__connecting[1] = true
+		try_connect(self)
+		self.__connecting[1] = nil
+		for i=2, #self.__connecting do
+			local co = self.__connecting[i]
+			self.__connecting[i] = nil
+			skynet.wakeup(co)
+		end
+	end
+
+	if self.__sock then
+		return true
+	else
+		return false
+	end
+end
+
+local function reconnect_channel(self)
+	local t = 100
+	while not self.__closed do
+		if self:connect() then
+			return
+		end
+		if t > 1000 then
+			print("socket: try to reconnect", self.__host, self.__port)
+			skynet.sleep(t)
+			t = 0
+		else
+			skynet.sleep(t)
+		end
+		t = t + 100
+	end
+end
+
+function channel:request(request, response)
+	if not self.__sock then
+		assert(not self.__closed)
+		reconnect_channel(self)
+	end
+
+	if not socket.write(self.__sock[1], request) then
+		return self:request(request, response)
+	end
+
+	if response == nil then
+		-- no response
+		return
+	end
+
+	local co = coroutine.running()
+
+	if self.__response then
+		-- response is session
+		self.__thread[response] = co
+	else
+		-- response is a function, push it to __request
+		table.insert(self.__request, response)
+		table.insert(self.__thread, co)
+	end
+	skynet.wait()
+
+	local result = self.__result[co]
+	self.__result[co] = nil
+	local result_data = self.__result_data[co]
+	self.__result_data[co] = nil
+
+	if result == socket_error then
+		if result_data then
+			print("socket: dispatch", request, result_data)
+		end
+		return self:request(request, response)
+	else
+		assert(result, result_data)
+		return result_data
+	end
+end
+
+function channel:close()
+	if not self.__closed then
+		self.__closed = true
+		close_channel_socket(self)
+	end
+end
+
+channel_meta.__gc = channel.close
+
+local function wrapper_socket_function(f)
+	return function(self, ...)
+		local result = f(self[1], ...)
+		if not result then
+			error(socket_error)
+		else
+			return result
+		end
+	end
+end
+
+channel_socket.read = wrapper_socket_function(socket.read)
+channel_socket.readline = wrapper_socket_function(socket.readline)
+
 return socket