From c19846c3a5ec2ab509b08f10dc1d8e24616777a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=91=E9=A3=8E?= <cloudwu@gmail.com>
Date: Sat, 15 Sep 2012 16:05:25 +0800
Subject: [PATCH] add trace into simpledb

---
 lualib-src/trace_service.c |  6 +++++-
 lualib/skynet.lua          | 31 ++++++++++++++-----------------
 service/simpledb.lua       |  8 ++++++++
 3 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/lualib-src/trace_service.c b/lualib-src/trace_service.c
index f6040581..be076e02 100644
--- a/lualib-src/trace_service.c
+++ b/lualib-src/trace_service.c
@@ -1,4 +1,5 @@
 #include "trace_service.h"
+
 #include <stdlib.h>
 #include <string.h>
 #include <assert.h>
@@ -82,6 +83,9 @@ trace_new(struct trace_pool *p) {
 void 
 trace_register(struct trace_pool *p, int session) {
 	struct trace_info *t = p->current;
+	if (t == NULL) {
+		return;
+	}
 	int hash = session % HASH_SIZE;
 	assert(t->session == 0);
 	t->session = session;
@@ -97,13 +101,13 @@ trace_yield(struct trace_pool *p) {
 	struct trace_info *t = p->current;
 	if (t == NULL)
 		return NULL;
-	p->current = NULL;
 
 	diff_time(&t->ti,&t->ti_sec,&t->ti_nsec);
 
 	if (t->session == 0) {
 		return t;
 	} else {
+		p->current = NULL;
 		return NULL;
 	}
 }
diff --git a/lualib/skynet.lua b/lualib/skynet.lua
index 13ffea96..348e7b0b 100644
--- a/lualib/skynet.lua
+++ b/lualib/skynet.lua
@@ -24,6 +24,9 @@ local session_coroutine_address = {}
 local wakeup_session = {}
 local sleep_session = {}
 
+local trace_handle
+local trace_func = function() end
+
 -- suspend is function
 local suspend
 
@@ -45,8 +48,10 @@ function suspend(co, result, command, param, size)
 		error(debug.traceback(co,command))
 	end
 	if command == "CALL" then
+		c.trace_register(trace_handle, param)
 		session_id_coroutine[param] = co
 	elseif command == "SLEEP" then
+		c.trace_register(trace_handle, param)
 		session_id_coroutine[param] = co
 		sleep_session[co] = param
 	elseif command == "RETURN" then
@@ -228,6 +233,7 @@ end
 local function dispatch_message(prototype, msg, sz, session, source, ...)
 	-- PTYPE_RESPONSE = 1, read skynet.h
 	if prototype == 1 then
+		c.trace_switch(trace_handle, session)
 		local co = session_id_coroutine[session]
 		if co == "BREAK" then
 			session_id_coroutine[session] = nil
@@ -250,6 +256,11 @@ local function dispatch_message(prototype, msg, sz, session, source, ...)
 			error(string.format("Can't dispatch type %s : ", p.name))
 		end
 	end
+	local info = c.trace_yield(trace_handle)
+	if info then
+		local ti = c.trace_delete(trace_handle, info)
+		trace_func(info, ti)
+	end
 	while true do
 		local key,co = next(fork_queue)
 		if co == nil then
@@ -483,8 +494,6 @@ local function init_template(start)
     init_all()
 end
 
-local trace_handle
-
 function skynet.start(start_func)
 	c.callback(dispatch_message)
 	trace_handle = assert(c.stat "trace")
@@ -505,24 +514,12 @@ function skynet.filter(f ,start_func)
 	end)
 end
 
-function skynet.trace_new()
+function skynet.trace()
 	return c.trace_new(trace_handle)
 end
 
-function skynet.trace_delete(info)
-	return c.trace_delete(info)
-end
-
-function skynet.trace_switch(session)
-	return c.trace_switch(trace_handle, session)
-end
-
-function skynet.trace_yield()
-	return c.trace_yield(trace_handle)
-end
-
-function skynet.trace_register(session)
-	return c.trace_register(trace_handle, session)
+function skynet.trace_callback(func)
+	trace_func = func
 end
 
 return skynet
diff --git a/service/simpledb.lua b/service/simpledb.lua
index ff2f912c..090402a7 100644
--- a/service/simpledb.lua
+++ b/service/simpledb.lua
@@ -13,8 +13,16 @@ function command.SET(key, value)
 	skynet.ret(last)
 end
 
+local trace_cache = {}
+
+skynet.trace_callback(function(handle, ti)
+	print("TRACE",trace_cache[handle],ti)
+	trace_cache[handle] = nil
+end)
+
 skynet.start(function()
 	skynet.dispatch("text", function(session, address, message)
+		trace_cache[skynet.trace()] = message
 		print("simpledb",message, skynet.address(address), session)
 		local cmd, key , value = string.match(message, "(%w+) (%w+) ?(.*)")
 		local f = command[cmd]