From 1d0d2678ed8d6a9c4f60261b03085995bb20cbce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=91=E9=A3=8E?= <cloudwu@gmail.com>
Date: Wed, 31 Jul 2013 16:41:46 +0800
Subject: [PATCH] mongo driver support request_id

---
 lualib-src/lua-mongo.c | 19 ++++++++++++
 lualib/mongo.lua       | 66 ++++++++++++++++++++++++++++--------------
 lualib/skynet.lua      | 20 +++++++++----
 3 files changed, 77 insertions(+), 28 deletions(-)

diff --git a/lualib-src/lua-mongo.c b/lualib-src/lua-mongo.c
index 39e8d1a5..a0e8a0c0 100644
--- a/lualib-src/lua-mongo.c
+++ b/lualib-src/lua-mongo.c
@@ -521,6 +521,24 @@ reply_length(lua_State *L) {
 	return 1;
 }
 
+static int
+copy_result(lua_State *L) {
+	if (lua_type(L,2) == LUA_TNIL)
+		return 0;
+	int n1 = lua_rawlen(L,1);
+	int n2 = lua_rawlen(L,2);
+	int i;
+	for (i=0;i<n1;i++) {
+		lua_rawgeti(L,1,i+1);
+		lua_rawseti(L,2,i+1);
+	}
+	for (i=n1;i<n2;i++) {
+		lua_pushnil(L);
+		lua_rawseti(L,2,i+1);
+	}
+	return 0;
+}
+
 int
 luaopen_mongo_driver(lua_State *L) {
 	luaL_checkversion(L);
@@ -533,6 +551,7 @@ luaopen_mongo_driver(lua_State *L) {
 		{ "update", op_update },
 		{ "insert", op_insert },
 		{ "length", reply_length },
+		{ "copy_result", copy_result },
 		{ NULL, NULL },
 	};
 
diff --git a/lualib/mongo.lua b/lualib/mongo.lua
index 9060f3b6..01a4fbac 100644
--- a/lualib/mongo.lua
+++ b/lualib/mongo.lua
@@ -1,5 +1,6 @@
 local bson = require "bson"
 local socket = require "socket"
+local skynet = require "skynet"
 local driver = require "mongo.driver"
 local rawget = rawget
 local assert = assert
@@ -61,11 +62,38 @@ local collection_meta = {
 	end
 }
 
+local function reply_queue(obj)
+	local sock = obj.__sock
+	local tmp = {}
+	local request_set = obj.__request
+	while true do
+		-- todo: reconnect
+		local length = driver.length(socket.read(sock, 4))
+		local reply = socket.read(sock, length)
+		if reply == nil then
+			--todo : reconnect
+			break
+		end
+		local succ, reply_id, document, cursor_id, startfrom = driver.reply(reply, tmp)
+		local result = assert(request_set[reply_id])
+		driver.copy_result(tmp, result.result)
+		result.succ = succ
+		result.document = document
+		result.cursor_id = cursor_id
+		result.startfrom = startfrom
+		result.data = reply
+		skynet.wakeup(result.co)
+	end
+end
+
 function mongo.client( obj )
 	obj.port = obj.port or 27017
 	obj.__id = 0
 	obj.__sock = assert(socket.open(obj.host, obj.port),"Connect failed")
-	return setmetatable(obj, client_meta)
+	obj.__request = {}
+	setmetatable(obj, client_meta)
+	skynet.fork(reply_queue, obj)
+	return obj
 end
 
 function mongo_client:getDB(dbname)
@@ -101,24 +129,23 @@ function mongo_client:runCommand(cmd)
 	return self.admin:runCommand(cmd)
 end
 
-local function get_reply(sock, result)
-	local length = driver.length(socket.read(sock, 4))
-	local reply = socket.read(sock, length)
-	return reply, driver.reply(reply, result)
+local function get_reply(conn, request_id, result)
+	local r = { result = result , co = coroutine.running() }
+	conn.__request[request_id] = r
+	skynet.wait()
+	conn.__request[request_id] = nil
+	return r.data, r.succ, r.document, r.cursor_id, r.startfrom
 end
 
 function mongo_db:runCommand(cmd)
-	local request_id = self.connection:genId()
-	local sock = self.connection.__sock
-	socket.lock(sock)
+	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))
 	-- todo: check send
 	socket.write(sock, pack)
 
-	local _, succ, reply_id, doc = get_reply(sock)
-	socket.unlock(sock)
-	assert(request_id == reply_id, "Reply from mongod error")
-	-- todo: check succ
+	local _, succ, doc = get_reply(conn,request_id)
 	return bson_decode(doc)
 end
 
@@ -175,17 +202,15 @@ function mongo_collection:delete(selector, single)
 end
 
 function mongo_collection:findOne(query, selector)
-	local request_id = self.connection:genId()
-	local sock = self.connection.__sock
-	socket.lock(sock)
+	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))
 
 	-- todo: check send
 	socket.write(sock, pack)
 
-	local _, succ, reply_id, doc = get_reply(sock)
-	socket.unlock(sock)
-	assert(request_id == reply_id, "Reply from mongod error")
+	local _, succ, doc = get_reply(conn, request_id)
 	-- todo: check succ
 	return bson_decode(doc)
 end
@@ -225,13 +250,10 @@ function mongo_cursor:hasNext()
 			end
 		end
 
-		socket.lock(sock)
 		--todo: check send
 		socket.write(sock, pack)
 
-		local data, succ, reply_id, doc, cursor = get_reply(sock, self.__document)
-		socket.unlock(sock)
-		assert(request_id == reply_id, "Reply from mongod error")
+		local data, succ, doc, cursor = get_reply(conn, request_id, self.__document)
 		if succ then
 			if doc then
 				self.__data = data
diff --git a/lualib/skynet.lua b/lualib/skynet.lua
index 108fd743..db060b36 100644
--- a/lualib/skynet.lua
+++ b/lualib/skynet.lua
@@ -122,7 +122,7 @@ function skynet.sleep(ti)
 	local session = c.command("TIMEOUT",tostring(ti))
 	assert(session)
 	session = tonumber(session)
-	local ret = coroutine.yield("SLEEP", session)
+	local ret = coroutine_yield("SLEEP", session)
 	sleep_session[coroutine.running()] = nil
 	if ret == true then
 		c.trace_switch(trace_handle, session)
@@ -133,10 +133,18 @@ end
 function skynet.yield()
 	local session = c.command("TIMEOUT","0")
 	assert(session)
-	coroutine.yield("SLEEP", tonumber(session))
+	coroutine_yield("SLEEP", tonumber(session))
 	sleep_session[coroutine.running()] = nil
 end
 
+function skynet.wait()
+	local session = c.genid()
+	coroutine_yield("SLEEP", session)
+	local co = coroutine.running()
+	sleep_session[co] = nil
+	session_id_coroutine[session] = nil
+end
+
 function skynet.register(name)
 	c.command("REG", name)
 end
@@ -224,25 +232,25 @@ skynet.tostring = assert(c.tostring)
 function skynet.call(addr, typename, ...)
 	local p = proto[typename]
 	local session = c.send(addr, p.id , nil , p.pack(...))
-	return p.unpack(coroutine.yield("CALL", session))
+	return p.unpack(coroutine_yield("CALL", session))
 end
 
 function skynet.blockcall(addr, typename , ...)
 	local p = proto[typename]
 	c.command("LOCK")
 	local session = c.send(addr, p.id , nil , p.pack(...))
-	return p.unpack(coroutine.yield("CALL", session))
+	return p.unpack(coroutine_yield("CALL", session))
 end
 
 function skynet.rawcall(addr, typename, msg, sz)
 	local p = proto[typename]
 	local session = c.send(addr, p.id , nil , msg, sz)
-	return coroutine.yield("CALL", session)
+	return coroutine_yield("CALL", session)
 end
 
 function skynet.ret(msg, sz)
 	msg = msg or ""
-	coroutine.yield("RETURN", msg, sz)
+	coroutine_yield("RETURN", msg, sz)
 end
 
 function skynet.wakeup(co)