bugfix: mongo driver runCommand , and fix socket.close

lualib/mongo.lua

@@ -6,6 +6,7 @@ local rawget = rawget
 local assert = assert
 
 local bson_encode = bson.encode
+local bson_encode_sorted = bson.encode_sorted
 local bson_decode = bson.decode
 local empty_bson = bson_encode {}
 
@@ -96,11 +97,17 @@ local function reply_queue(obj)
 	while true do
 		local len_reply = socket.read(sock, 4)
 		if not len_reply then
+			if obj.__sock == nil then
+				return
+			end
 			sock = reconnect(obj)
 		else
 			local length = driver.length(len_reply)
 			local reply = socket.read(sock, length)
 			if not reply then
+				if obj.__sock == nil then
+					return
+				end
 				sock = reconnect(obj)
 			else
 				local succ, reply_id, document, cursor_id, startfrom = driver.reply(reply, tmp)
@@ -142,8 +149,9 @@ end
 
 function mongo_client:disconnect()
 	if self.__sock then
-		socket.close(self.__sock)
+		local so = self.__sock
 		self.__sock = nil
+		socket.close(so)
 	end
 end
 
@@ -153,11 +161,11 @@ function mongo_client:genId()
 	return id
 end
 
-function mongo_client:runCommand(cmd)
+function mongo_client:runCommand(...)
 	if not self.admin then
 		self.admin = self:getDB "admin"
 	end
-	return self.admin:runCommand(cmd)
+	return self.admin:runCommand(...)
 end
 
 local function get_reply(conn, request_id, result)
@@ -168,11 +176,17 @@ local function get_reply(conn, request_id, result)
 	return r.data, r.succ, r.document, r.cursor_id, r.startfrom
 end
 
-function mongo_db:runCommand(cmd)
+function mongo_db:runCommand(cmd,cmd_v,...)
 	local conn = self.connection
 	local request_id = conn:genId()
 	local sock = conn.__sock
-	local pack = driver.query(request_id, 0, self.__cmd, 0, 1, bson_encode(cmd))
+	local bson_cmd
+	if not cmd_v then
+		bson_cmd = bson_encode_sorted(cmd,1)
+	else
+		bson_cmd = bson_encode_sorted(cmd,cmd_v,...)
+	end
+	local pack = driver.query(request_id, 0, self.__cmd, 0, 1, bson_cmd)
 	-- todo: check send
 	assert(socket.write(sock, pack), "write fail")
 	local _, succ, doc = get_reply(conn,request_id)

lualib/socket.lua

@@ -153,7 +153,14 @@ function socket.close(id)
 	end
 	if s.connected then
 		driver.close(s.id)
-		suspend(s)
+		if s.co then
+			-- reading this socket on another coroutine
+			assert(not s.closing)
+			s.closing = coroutine.running()
+			skynet.wait()
+		else
+			suspend(s)
+		end
 	end
 	if s.buffer then
 		driver.clear(s.buffer,buffer_pool)
@@ -162,6 +169,13 @@ function socket.close(id)
 	socket_pool[id] = nil
 end
 
+local function close_socket(s)
+	if s.closing then
+		skynet.wakeup(s.closing)
+	end
+	return driver.readall(s.buffer, buffer_pool)
+end
+
 function socket.read(id, sz)
 	local s = socket_pool[id]
 	assert(s)
@@ -170,7 +184,7 @@ function socket.read(id, sz)
 		return ret
 	end
 	if not s.connected then
-		return false, driver.readall(s.buffer, buffer_pool)
+		return false, close_socket(s)
 	end
 
 	assert(not s.read_required)
@@ -180,7 +194,7 @@ function socket.read(id, sz)
 	if ret then
 		return ret
 	else
-		return false, driver.readall(s.buffer, buffer_pool)
+		return false, close_socket(s)
 	end
 end
 
@@ -188,14 +202,14 @@ function socket.readall(id)
 	local s = socket_pool[id]
 	assert(s)
 	if not s.connected then
-		local r = driver.readall(s.buffer, buffer_pool)
+		local r = close_socket(s)
 		return r ~= "" and r
 	end
 	assert(not s.read_required)
 	s.read_required = true
 	suspend(s)
 	assert(s.connected == false)
-	return driver.readall(s.buffer, buffer_pool)
+	return close_socket(s)
 end
 
 function socket.readline(id, sep)
@@ -207,7 +221,7 @@ function socket.readline(id, sep)
 		return ret
 	end
 	if not s.connected then
-		return false, driver.readall(s.buffer, buffer_pool)
+		return false, close_socket(s)
 	end
 	assert(not s.read_required)
 	s.read_required = sep
@@ -215,7 +229,7 @@ function socket.readline(id, sep)
 	if s.connected then
 		return driver.readline(s.buffer, buffer_pool, sep)
 	else
-		return false, driver.readall(s.buffer, buffer_pool)
+		return false, close_socket(s)
 	end
 end