bugfix: mongo driver runCommand , and fix socket.close

lualib-src/lua-bson.c

@@ -457,6 +457,24 @@ pack_dict(lua_State *L, struct bson *b, bool isarray) {
 	write_length(b, b->size - length, length);
 }
 
+static void
+pack_sorted_dict(lua_State *L, struct bson *b, int n) {
+	int length = reserve_length(b);
+	int i;
+	for (i=0;i<n;i+=2) {
+		size_t sz;
+		const char * key = lua_tolstring(L, i+1, &sz);
+		if (key == NULL) {
+			luaL_error(L, "Argument %d need a string", i+1);
+		}
+		lua_pushvalue(L, i+2);
+		append_one(b, L, key, sz);
+		lua_pop(L,1);
+	}
+	write_byte(b,0);
+	write_length(b, b->size - length, length);
+}
+ 
 static int
 ltostring(lua_State *L) {
 	size_t sz = lua_rawlen(L, 1);
@@ -804,16 +822,8 @@ ldecode(lua_State *L) {
 	return 1;
 }
 
-static int
-lencode(lua_State *L) {
-	struct bson b;
-	bson_create(&b);
-	lua_settop(L,1);
-	luaL_checktype(L, 1, LUA_TTABLE);
-	pack_dict(L, &b, false);
-	void * ud = lua_newuserdata(L, b.size);
-	memcpy(ud, b.ptr, b.size);
-	bson_destroy(&b);
+static void
+bson_meta(lua_State *L) {
 	if (luaL_newmetatable(L, "bson")) {
 		luaL_Reg l[] = {
 			{ "decode", ldecode },
@@ -830,6 +840,36 @@ lencode(lua_State *L) {
 		lua_setfield(L, -2, "__newindex");
 	}
 	lua_setmetatable(L, -2);
+}
+
+static int
+lencode(lua_State *L) {
+	struct bson b;
+	bson_create(&b);
+	lua_settop(L,1);
+	luaL_checktype(L, 1, LUA_TTABLE);
+	pack_dict(L, &b, false);
+	void * ud = lua_newuserdata(L, b.size);
+	memcpy(ud, b.ptr, b.size);
+	bson_destroy(&b);
+	bson_meta(L);
+	return 1;
+}
+
+static int
+lencode_sorted(lua_State *L) {
+	struct bson b;
+	bson_create(&b);
+	int n = lua_gettop(L);
+	if (n%2 != 0) {
+		return luaL_error(L, "Invalid sorted dict");
+	}
+	pack_sorted_dict(L, &b, n);
+	lua_settop(L,1);
+	void * ud = lua_newuserdata(L, b.size);
+	memcpy(ud, b.ptr, b.size);
+	bson_destroy(&b);
+	bson_meta(L);
 	return 1;
 }
 
@@ -1121,6 +1161,7 @@ luaopen_bson(lua_State *L) {
 	}
 	luaL_Reg l[] = {
 		{ "encode", lencode },
+		{ "encode_sorted", lencode_sorted },
 		{ "date", ldate },
 		{ "timestamp", ltimestamp  },
 		{ "regex", lregex },

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