diff --git a/lualib-src/lua-bson.c b/lualib-src/lua-bson.c index bc3b06e4..26226746 100644 --- a/lualib-src/lua-bson.c +++ b/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;isize - 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 }, diff --git a/lualib/mongo.lua b/lualib/mongo.lua index e3e03660..0dd36096 100644 --- a/lualib/mongo.lua +++ b/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) diff --git a/lualib/socket.lua b/lualib/socket.lua index 0461d513..4eaa24d8 100644 --- a/lualib/socket.lua +++ b/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