Merge branch 'mongors' into dev

lualib/mongo.lua

@@ -1,96 +1,125 @@
 local bson = require "bson"
 local socket = require "socket"
-local socketchannel = require "socketchannel"
+local socketchannel	= require "socketchannel"
 local skynet = require "skynet"
 local driver = require "mongo.driver"
-local md5 = require "md5"
+local md5 =	require	"md5"
 local rawget = rawget
 local assert = assert
 
-local bson_encode = bson.encode
+local bson_encode =	bson.encode
-local bson_encode_order = bson.encode_order
+local bson_encode_order	= bson.encode_order
-local bson_decode = bson.decode
+local bson_decode =	bson.decode
 local empty_bson = bson_encode {}
 
-local mongo = {}
+local mongo	= {}
 mongo.null = assert(bson.null)
 mongo.maxkey = assert(bson.maxkey)
 mongo.minkey = assert(bson.minkey)
 mongo.type = assert(bson.type)
 
 local mongo_cursor = {}
-local cursor_meta = {
+local cursor_meta =	{
-	__index = mongo_cursor,
+	__index	= mongo_cursor,
 }
 
 local mongo_client = {}
 
-local client_meta = {
+local client_meta =	{
-	__index = function(self, key)
+	__index	= function(self, key)
-		return rawget(mongo_client, key) or self:getDB(key)
+		return rawget(mongo_client,	key) or	self:getDB(key)
 	end,
 	__tostring = function (self)
 		local port_string
 		if self.port then
-			port_string = ":" .. tostring(self.port)
+			port_string	= ":" .. tostring(self.port)
 		else
-			port_string = ""
+			port_string	= ""
 		end
 
-		return "[mongo client : " .. self.host .. port_string .."]"
+		return "[mongo client :	" .. self.host .. port_string .."]"
 	end,
-	-- DO NOT need disconnect, because channel will shutdown during gc
+	-- DO NOT need disconnect, because channel will	shutdown during	gc
 }
 
 local mongo_db = {}
 
-local db_meta = {
+local db_meta =	{
-	__index = function (self, key)
+	__index	= function (self, key)
-		return rawget(mongo_db, key) or self:getCollection(key)
+		return rawget(mongo_db,	key) or	self:getCollection(key)
 	end,
 	__tostring = function (self)
-		return "[mongo db : " .. self.name .. "]"
+		return "[mongo db :	" .. self.name .. "]"
 	end
 }
 
 local mongo_collection = {}
-local collection_meta = {
+local collection_meta =	{
-	__index = function(self, key)
+	__index	= function(self, key)
-		return rawget(mongo_collection, key) or self:getCollection(key)
+		return rawget(mongo_collection,	key) or	self:getCollection(key)
-	end ,
+	end	,
 	__tostring = function (self)
-		return "[mongo collection : " .. self.full_name .. "]"
+		return "[mongo collection :	" .. self.full_name	.. "]"
 	end
 }
 
 local function dispatch_reply(so)
-	local len_reply = so:read(4)
+	local len_reply	= so:read(4)
-	local reply = so:read(driver.length(len_reply))
+	local reply	= so:read(driver.length(len_reply))
-	local result = { result = {} }
+	local result = { result	= {} }
-	local succ, reply_id, document, cursor_id, startfrom = driver.reply(reply, result.result)
+	local succ,	reply_id, document,	cursor_id, startfrom = driver.reply(reply, result.result)
-	result.document = document
+	result.document	= document
 	result.cursor_id = cursor_id
 	result.startfrom = startfrom
-	result.data = reply
+	result.data	= reply
 	return reply_id, succ, result
 end
 
-local function mongo_auth(mongoc)
+local function __parse_addr(addr)
-	local user = rawget(mongoc, "username")
+	local host,	port = string.match(addr, "([^:]+):(.+)")
-	local pass = rawget(mongoc, "password")
+	return host, tonumber(port)
+end
+
+local function mongo_auth(mongoc)
+	local user = rawget(mongoc,	"username")
+	local pass = rawget(mongoc,	"password")
 
-	if user == nil or pass == nil then
-		return
-	end
 	return function()
-		assert(mongoc:auth(user, pass))
+		if user	~= nil and pass	~= nil then
+			assert(mongoc:auth(user, pass))
+		end
+		local rs_data =	mongoc:runCommand("ismaster")
+		if rs_data.ok == 1 then
+			if rs_data.hosts then
+				local backup = {}
+				for	_, v in	ipairs(rs_data.hosts) do
+					local host,	port = __parse_addr(v)
+					table.insert(backup, {host = host, port	= port})
+				end
+				mongoc.__sock:changebackup(backup)
+			end
+			if rs_data.ismaster	then
+				return
+			else
+				local host,	port = __parse_addr(rs_data.primary)
+				mongoc.host	= host
+				mongoc.port	= port
+				mongoc.__sock:changehost(host, port)
+			end
+		end
 	end
 end
 
-function mongo.client( conf )
+function mongo.client( conf	)
-	local obj = {
+	local first	= conf
-		host = conf.host,
+	local backup = nil
-		port = conf.port or 27017,
+	if conf.rs then
+		first =	conf.rs[1]
+		backup = conf.rs
+	end
+	local obj =	{
+		host = first.host,
+		port = first.port or 27017,
 	}
 
 	obj.__id = 0
@@ -99,9 +128,10 @@ function mongo.client( conf )
 		port = obj.port,
 		response = dispatch_reply,
 		auth = mongo_auth(obj),
+		backup = backup,
 	}
 	setmetatable(obj, client_meta)
-	obj.__sock:connect(true)	-- try connect only once
+	obj.__sock:connect(true)	-- try connect only	once
 	return obj
 end
 
@@ -109,32 +139,32 @@ function mongo_client:getDB(dbname)
 	local db = {
 		connection = self,
 		name = dbname,
-		full_name = dbname,
+		full_name =	dbname,
 		database = false,
-		__cmd = dbname .. "." .. "$cmd",
+		__cmd =	dbname .. "." .. "$cmd",
 	}
-	db.database = db
+	db.database	= db
 
-	return setmetatable(db, db_meta)
+	return setmetatable(db,	db_meta)
 end
 
 function mongo_client:disconnect()
 	if self.__sock then
 		local so = self.__sock
-		self.__sock = false
+		self.__sock	= false
 		so:close()
 	end
 end
 
 function mongo_client:genId()
 	local id = self.__id + 1
-	self.__id = id
+	self.__id =	id
 	return id
 end
 
 function mongo_client:runCommand(...)
 	if not self.admin then
-		self.admin = self:getDB "admin"
+		self.admin = self:getDB	"admin"
 	end
 	return self.admin:runCommand(...)
 end
@@ -146,14 +176,14 @@ function mongo_client:auth(user,password)
 		return false
 	end
 
-	local key = md5.sumhexa(string.format("%s%s%s",result.nonce,user,password))
+	local key =	md5.sumhexa(string.format("%s%s%s",result.nonce,user,password))
 	local result= self:runCommand ("authenticate",1,"user",user,"nonce",result.nonce,"key",key)
-	return result.ok == 1
+	return result.ok ==	1
 end
 
 function mongo_client:logout()
 	local result = self:runCommand "logout"
-	return result.ok == 1
+	return result.ok ==	1
 end
 
 function mongo_db:runCommand(cmd,cmd_v,...)
@@ -166,18 +196,18 @@ function mongo_db:runCommand(cmd,cmd_v,...)
 	else
 		bson_cmd = bson_encode_order(cmd,cmd_v,...)
 	end
-	local pack = driver.query(request_id, 0, self.__cmd, 0, 1, bson_cmd)
+	local pack = driver.query(request_id, 0, self.__cmd, 0,	1, bson_cmd)
-	-- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data)
+	-- we must hold	req	(req.data),	because	req.document is	a lightuserdata, it's a	pointer	to the string (req.data)
-	local req = sock:request(pack, request_id)
+	local req =	sock:request(pack, request_id)
-	local doc = req.document
+	local doc =	req.document
 	return bson_decode(doc)
 end
 
 function mongo_db:getCollection(collection)
-	local col = {
+	local col =	{
 		connection = self.connection,
 		name = collection,
-		full_name = self.full_name .. "." .. collection,
+		full_name =	self.full_name .. "." .. collection,
 		database = self.database,
 	}
 	self[collection] = setmetatable(col, collection_meta)
@@ -188,20 +218,20 @@ mongo_collection.getCollection = mongo_db.getCollection
 
 function mongo_collection:insert(doc)
 	if doc._id == nil then
-		doc._id = bson.objectid()
+		doc._id	= bson.objectid()
 	end
 	local sock = self.connection.__sock
 	local pack = driver.insert(0, self.full_name, bson_encode(doc))
-	-- flags support 1: ContinueOnError
+	-- flags support 1:	ContinueOnError
 	sock:request(pack)
 end
 
 function mongo_collection:batch_insert(docs)
-	for i=1,#docs do
+	for	i=1,#docs do
 		if docs[i]._id == nil then
-			docs[i]._id = bson.objectid()
+			docs[i]._id	= bson.objectid()
 		end
-		docs[i] = bson_encode(docs[i])
+		docs[i]	= bson_encode(docs[i])
 	end
 	local sock = self.connection.__sock
 	local pack = driver.insert(0, self.full_name, docs)
@@ -209,7 +239,7 @@ function mongo_collection:batch_insert(docs)
 end
 
 function mongo_collection:update(selector,update,upsert,multi)
-	local flags = (upsert and 1 or 0) + (multi and 2 or 0)
+	local flags	= (upsert and 1	or 0) +	(multi and 2 or	0)
 	local sock = self.connection.__sock
 	local pack = driver.update(self.full_name, flags, bson_encode(selector), bson_encode(update))
 	sock:request(pack)
@@ -225,24 +255,24 @@ function mongo_collection:findOne(query, selector)
 	local conn = self.connection
 	local request_id = conn:genId()
 	local sock = conn.__sock
-	local pack = driver.query(request_id, 0, self.full_name, 0, 1, query and bson_encode(query) or empty_bson, selector and bson_encode(selector))
+	local pack = driver.query(request_id, 0, self.full_name, 0,	1, query and bson_encode(query)	or empty_bson, selector	and	bson_encode(selector))
-	-- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data)
+	-- we must hold	req	(req.data),	because	req.document is	a lightuserdata, it's a	pointer	to the string (req.data)
-	local req = sock:request(pack, request_id)
+	local req =	sock:request(pack, request_id)
-	local doc = req.document
+	local doc =	req.document
 	return bson_decode(doc)
 end
 
 function mongo_collection:find(query, selector)
 	return setmetatable( {
 		__collection = self,
-		__query = query and bson_encode(query) or empty_bson,
+		__query	= query	and	bson_encode(query) or empty_bson,
 		__selector = selector and bson_encode(selector),
-		__ptr = nil,
+		__ptr =	nil,
 		__data = nil,
 		__cursor = nil,
 		__document = {},
-		__flags = 0,
+		__flags	= 0,
-	} , cursor_meta)
+	} ,	cursor_meta)
 end
 
 function mongo_cursor:hasNext()
@@ -255,42 +285,42 @@ function mongo_cursor:hasNext()
 		local sock = conn.__sock
 		local pack
 		if self.__data == nil then
-			pack = driver.query(request_id, self.__flags, self.__collection.full_name,0,0,self.__query,self.__selector)
+			pack = driver.query(request_id,	self.__flags, self.__collection.full_name,0,0,self.__query,self.__selector)
 		else
 			if self.__cursor then
 				pack = driver.more(request_id, self.__collection.full_name,0,self.__cursor)
 			else
 				-- no more
-				self.__document = nil
+				self.__document	= nil
-				self.__data = nil
+				self.__data	= nil
 				return false
 			end
 		end
 
 		local ok, result = pcall(sock.request,sock,pack, request_id)
 
-		local doc = result.document
+		local doc =	result.document
 		local cursor = result.cursor_id
 
 		if ok then
 			if doc then
-				self.__document = result.result
+				self.__document	= result.result
-				self.__data = result.data
+				self.__data	= result.data
 				self.__ptr = 1
-				self.__cursor = cursor
+				self.__cursor =	cursor
 				return true
 			else
-				self.__document = nil
+				self.__document	= nil
-				self.__data = nil
+				self.__data	= nil
-				self.__cursor = nil
+				self.__cursor =	nil
 				return false
 			end
 		else
-			self.__document = nil
+			self.__document	= nil
-			self.__data = nil
+			self.__data	= nil
-			self.__cursor = nil
+			self.__cursor =	nil
 			if doc then
-				local err = bson_decode(doc)
+				local err =	bson_decode(doc)
 				error(err["$err"])
 			else
 				error("Reply from mongod error")
@@ -303,11 +333,11 @@ end
 
 function mongo_cursor:next()
 	if self.__ptr == nil then
-		error "Call hasNext first"
+		error "Call	hasNext	first"
 	end
-	local r = bson_decode(self.__document[self.__ptr])
+	local r	= bson_decode(self.__document[self.__ptr])
-	self.__ptr = self.__ptr + 1
+	self.__ptr = self.__ptr	+ 1
-	if self.__ptr > #self.__document then
+	if self.__ptr >	#self.__document then
 		self.__ptr = nil
 	end
 

lualib/socketchannel.lua

@@ -26,6 +26,7 @@ function socket_channel.channel(desc)
 	local c = {
 		__host = assert(desc.host),
 		__port = assert(desc.port),
+		__backup = desc.backup,
 		__auth = desc.auth,
 		__response = desc.response,	-- It's for session mode
 		__request = {},	-- request seq { response func or session }	-- It's for order mode
@@ -152,34 +153,66 @@ local function dispatch_function(self)
 	end
 end
 
+local function connect_backup(self)
+	if self.__backup then
+		for _, addr in ipairs(self.__backup) do
+			local host, port
+			if type(addr) == "table" then
+				host, port = addr.host, addr.port
+			else
+				host = addr
+				port = self.__port
+			end
+			skynet.error("socket: connect to backup host", host, port)
+			local fd = socket.open(host, port)
+			if fd then
+				self.__host = host
+				self.__port = port
+				return fd
+			end
+		end
+	end
+end
+
 local function connect_once(self)
+	if self.__closed then
+		return false
+	end
 	assert(not self.__sock and not self.__authcoroutine)
 	local fd = socket.open(self.__host, self.__port)
 	if not fd then
-		return false
+		fd = connect_backup(self)
+		if not fd then
+			return false
+		end
 	end
-	self.__authcoroutine = coroutine.running()
+
 	self.__sock = setmetatable( {fd} , channel_socket_meta )
 	skynet.fork(dispatch_function(self), self)
 
 	if self.__auth then
+		self.__authcoroutine = coroutine.running()
 		local ok , message = pcall(self.__auth, self)
 		if not ok then
 			close_channel_socket(self)
 			if message ~= socket_error then
+				self.__authcoroutine = false
 				skynet.error("socket: auth failed", message)
 			end
 		end
 		self.__authcoroutine = false
+		if ok and not self.__sock then
+			-- auth may change host, so connect again
+			return connect_once(self)
+		end
 		return ok
 	end
 
-	self.__authcoroutine = false
 	return true
 end
 
 local function try_connect(self , once)
-	local t = 100
+	local t = 0
 	while not self.__closed do
 		if connect_once(self) then
 			if not once then
@@ -292,6 +325,20 @@ function channel:close()
 	end
 end
 
+function channel:changehost(host, port)
+	self.__host = host
+	if port then
+		self.__port = port
+	end
+	if not self.__closed then
+		close_channel_socket(self)
+	end
+end
+
+function channel:changebackup(backup)
+	self.__backup = backup
+end
+
 channel_meta.__gc = channel.close
 
 local function wrapper_socket_function(f)