From 6813fd9ef79dd9a922ddc08d88173a2372406f8a Mon Sep 17 00:00:00 2001
From: Cloud Wu <cloudwu@gmail.com>
Date: Sun, 13 Jul 2014 11:55:25 +0800
Subject: [PATCH] bugfix: skynet.exit redirect error, and add datacenter.wait

---
 lualib/datacenter.lua   |  4 ++
 lualib/skynet.lua       |  2 +-
 service/datacenterd.lua | 81 ++++++++++++++++++++++++++++++++++++++---
 test/testdatacenter.lua | 23 ++++++++++++
 test/testdeadcall.lua   |  4 +-
 5 files changed, 105 insertions(+), 9 deletions(-)
 create mode 100644 test/testdatacenter.lua

diff --git a/lualib/datacenter.lua b/lualib/datacenter.lua
index 347c140e..ebf68ba1 100644
--- a/lualib/datacenter.lua
+++ b/lualib/datacenter.lua
@@ -10,5 +10,9 @@ function datacenter.set(...)
 	return skynet.call("DATACENTER", "lua", "UPDATE", ...)
 end
 
+function datacenter.wait(...)
+	return skynet.call("DATACENTER", "lua", "WAIT", ...)
+end
+
 return datacenter
 
diff --git a/lualib/skynet.lua b/lualib/skynet.lua
index d9c93616..6ccdaea2 100644
--- a/lualib/skynet.lua
+++ b/lualib/skynet.lua
@@ -257,7 +257,7 @@ function skynet.exit()
 		local address = session_coroutine_address[co]
 		local self = skynet.self()
 		if session~=0 and address then
-			skynet.redirect(self, address, "error", session, "")
+			skynet.redirect(address, self, "error", session, "")
 		end
 	end
 	c.command("EXIT")
diff --git a/service/datacenterd.lua b/service/datacenterd.lua
index ad431ae2..1a0fee65 100644
--- a/service/datacenterd.lua
+++ b/service/datacenterd.lua
@@ -2,6 +2,8 @@ local skynet = require "skynet"
 
 local command = {}
 local database = {}
+local wait_queue = {}
+local mode = {}
 
 local function query(db, key, ...)
 	if key == nil then
@@ -22,7 +24,7 @@ local function update(db, key, value, ...)
 	if select("#",...) == 0 then
 		local ret = db[key]
 		db[key] = value
-		return ret
+		return ret, value
 	else
 		if db[key] == nil then
 			db[key] = {}
@@ -31,13 +33,82 @@ local function update(db, key, value, ...)
 	end
 end
 
+local function wakeup(db, key1, key2, value, ...)
+	if key1 == nil then
+		return
+	end
+	local q = db[key1]
+	if q == nil then
+		return
+	end
+	if q[mode] == "queue" then
+		db[key1] = nil
+		if value then
+			-- throw error because can't wake up a branch
+			for _,v in ipairs(q) do
+				local session = v[1]
+				local source = v[2]
+				skynet.redirect(source, 0, "error", session, "")
+			end
+		else
+			return q
+		end
+	else
+		-- it's branch
+		return wakeup(q , key2, value, ...)
+	end
+end
+
 function command.UPDATE(...)
-	return update(database, ...)
+	local ret, value = update(database, ...)
+	if ret or value == nil then
+		return ret
+	end
+	local q = wakeup(wait_queue, ...)
+	if q then
+		for _, v in ipairs(q) do
+			local session = v[1]
+			local source = v[2]
+			skynet.redirect(source, 0, "response", session, skynet.pack(value))
+		end
+	end
+end
+
+local function waitfor(session, source, db, key1, key2, ...)
+	if key2 == nil then
+		-- push queue
+		local q = db[key1]
+		if q == nil then
+			q = { [mode] = "queue" }
+			db[key1] = q
+		else
+			assert(q[mode] == "queue")
+		end
+		table.insert(q, { session, source })
+	else
+		local q = db[key1]
+		if q == nil then
+			q = { [mode] = "branch" }
+			db[key1] = q
+		else
+			assert(q[mode] == "branch")
+		end
+		return waitfor(session, source, q, key2, ...)
+	end
 end
 
 skynet.start(function()
-	skynet.dispatch("lua", function (_, source, cmd, ...)
-		local f = assert(command[cmd])
-		skynet.ret(skynet.pack(f(...)))
+	skynet.dispatch("lua", function (session, source, cmd, ...)
+		if cmd == "WAIT" then
+			local ret = command.QUERY(...)
+			if ret then
+				skynet.ret(skynet.pack(ret))
+			else
+				waitfor(session, source, wait_queue, ...)
+			end
+		else
+			local f = assert(command[cmd])
+			skynet.ret(skynet.pack(f(...)))
+		end
 	end)
 end)
diff --git a/test/testdatacenter.lua b/test/testdatacenter.lua
new file mode 100644
index 00000000..194d2219
--- /dev/null
+++ b/test/testdatacenter.lua
@@ -0,0 +1,23 @@
+local skynet = require "skynet"
+local datacenter = require "datacenter"
+
+local function f1()
+	print("====1==== wait hello")
+	print("\t1>",datacenter.wait ("hello"))
+	print("====1==== wait key.foobar")
+	print("\t1>", pcall(datacenter.wait,"key"))	-- will failed, because "key" is a branch
+	print("\t1>",datacenter.wait ("key", "foobar"))
+end
+
+local function f2()
+	skynet.sleep(10)
+	print("====2==== set key.foobar")
+	datacenter.set("key", "foobar", "bingo")
+end
+
+skynet.start(function()
+	datacenter.set("hello", "world")
+	print(datacenter.get "hello")
+	skynet.fork(f1)
+	skynet.fork(f2)
+end)
diff --git a/test/testdeadcall.lua b/test/testdeadcall.lua
index c5d67b8c..9017d380 100644
--- a/test/testdeadcall.lua
+++ b/test/testdeadcall.lua
@@ -17,9 +17,7 @@ else
 		local test = skynet.newservice(SERVICE_NAME, "test")	-- launch self in test mode
 
 		print(pcall(function()
-			skynet.send(test,"lua", "hello world")
-			skynet.send(test,"lua", "never get there")
-			skynet.call(test,"lua", "fake call")
+			skynet.call(test,"lua", "dead call")
 		end))
 
 		skynet.exit()