profile support skynet coroutine

lualib-src/lua-profile.c

@@ -1,3 +1,4 @@
+#include <stdio.h>
 #include <lua.h>
 #include <lauxlib.h>
 
@@ -11,6 +12,8 @@
 #define NANOSEC 1000000000
 #define MICROSEC 1000000
 
+// #define DEBUG_LOG
+
 static double
 get_time() {
 #if  !defined(__APPLE__)
@@ -47,7 +50,11 @@ diff_time(double start) {
 
 static int
 lstart(lua_State *L) {
-	lua_pushthread(L);
+	if (lua_type(L,1) == LUA_TTHREAD) {
+		lua_settop(L,1);
+	} else {
+		lua_pushthread(L);
+	}
 	lua_rawget(L, lua_upvalueindex(2));
 	if (!lua_isnil(L, -1)) {
 		return luaL_error(L, "Thread %p start profile more than once", lua_topointer(L, 1));
@@ -57,7 +64,11 @@ lstart(lua_State *L) {
 	lua_rawset(L, lua_upvalueindex(2));
 
 	lua_pushthread(L);
-	lua_pushnumber(L, get_time());
+	double ti = get_time();
+#ifdef DEBUG_LOG
+	fprintf(stderr, "PROFILE [%p] start\n", L);
+#endif
+	lua_pushnumber(L, ti);
 	lua_rawset(L, lua_upvalueindex(1));
 
 	return 0;
@@ -65,9 +76,15 @@ lstart(lua_State *L) {
 
 static int
 lstop(lua_State *L) {
-	lua_pushthread(L);
+	if (lua_type(L,1) == LUA_TTHREAD) {
+		lua_settop(L,1);
+	} else {
+		lua_pushthread(L);
+	}
 	lua_rawget(L, lua_upvalueindex(1));
-	luaL_checktype(L, -1, LUA_TNUMBER);
+	if (lua_type(L, -1) != LUA_TNUMBER) {
+		return luaL_error(L, "Call profile.start() before profile.stop()");
+	} 
 	double ti = diff_time(lua_tonumber(L, -1));
 	lua_pushthread(L);
 	lua_rawget(L, lua_upvalueindex(2));
@@ -81,14 +98,20 @@ lstop(lua_State *L) {
 	lua_pushnil(L);
 	lua_rawset(L, lua_upvalueindex(2));
 
-	lua_pushnumber(L, ti + total_time);
+	total_time += ti;
+	lua_pushnumber(L, total_time);
+#ifdef DEBUG_LOG
+	fprintf(stderr, "PROFILE [%p] stop (%lf / %lf)\n", L, ti, total_time);
+#endif
 
 	return 1;
 }
 
 static int
-lresume(lua_State *L) {
-	lua_pushvalue(L,1);
+timing_resume(lua_State *L) {
+#ifdef DEBUG_LOG
+	lua_State *from = lua_tothread(L, -1);
+#endif
 	lua_rawget(L, lua_upvalueindex(2));
 	if (lua_isnil(L, -1)) {		// check total time
 		lua_pop(L,1);
@@ -96,6 +119,9 @@ lresume(lua_State *L) {
 		lua_pop(L,1);
 		lua_pushvalue(L,1);
 		double ti = get_time();
+#ifdef DEBUG_LOG
+		fprintf(stderr, "PROFILE [%p] resume\n", from);
+#endif
 		lua_pushnumber(L, ti);
 		lua_rawset(L, lua_upvalueindex(1));	// set start time
 	}
@@ -106,8 +132,25 @@ lresume(lua_State *L) {
 }
 
 static int
-lyield(lua_State *L) {
-	lua_pushthread(L);
+lresume(lua_State *L) {
+	lua_pushvalue(L,1);
+	
+	return timing_resume(L);
+}
+
+static int
+lresume_co(lua_State *L) {
+	luaL_checktype(L, 2, LUA_TTHREAD);
+	lua_rotate(L, 2, -1);
+
+	return timing_resume(L);
+}
+
+static int
+timing_yield(lua_State *L) {
+#ifdef DEBUG_LOG
+	lua_State *from = lua_tothread(L, -1);
+#endif
 	lua_rawget(L, lua_upvalueindex(2));	// check total time
 	if (lua_isnil(L, -1)) {
 		lua_pop(L,1);
@@ -120,7 +163,11 @@ lyield(lua_State *L) {
 		double starttime = lua_tonumber(L, -1);
 		lua_pop(L,1);
 
-		ti += diff_time(starttime);
+		double diff = diff_time(starttime);
+		ti += diff;
+#ifdef DEBUG_LOG
+		fprintf(stderr, "PROFILE [%p] yield (%lf/%lf)\n", from, diff, ti);
+#endif
 
 		lua_pushthread(L);
 		lua_pushnumber(L, ti);
@@ -132,6 +179,21 @@ lyield(lua_State *L) {
 	return co_yield(L);
 }
 
+static int
+lyield(lua_State *L) {
+	lua_pushthread(L);
+
+	return timing_yield(L);
+}
+
+static int
+lyield_co(lua_State *L) {
+	luaL_checktype(L, 1, LUA_TTHREAD);
+	lua_rotate(L, 1, -1);
+	
+	return timing_yield(L);
+}
+
 int
 luaopen_profile(lua_State *L) {
 	luaL_checkversion(L);
@@ -140,6 +202,8 @@ luaopen_profile(lua_State *L) {
 		{ "stop", lstop },
 		{ "resume", lresume },
 		{ "yield", lyield },
+		{ "resume_co", lresume_co },
+		{ "yield_co", lyield_co },
 		{ NULL, NULL },
 	};
 	luaL_newlibtable(L,l);
@@ -154,7 +218,7 @@ luaopen_profile(lua_State *L) {
 	lua_setmetatable(L, -3); 
 	lua_setmetatable(L, -3);
 
-	lua_pushnil(L);
+	lua_pushnil(L);	// cfunction (coroutine.resume or coroutine.yield)
 	luaL_setfuncs(L,l,3);
 
 	int libtable = lua_gettop(L);
@@ -166,20 +230,33 @@ luaopen_profile(lua_State *L) {
 	if (co_resume == NULL)
 		return luaL_error(L, "Can't get coroutine.resume");
 	lua_pop(L,1);
+
 	lua_getfield(L, libtable, "resume");
 	lua_pushcfunction(L, co_resume);
 	lua_setupvalue(L, -2, 3);
 	lua_pop(L,1);
 
+	lua_getfield(L, libtable, "resume_co");
+	lua_pushcfunction(L, co_resume);
+	lua_setupvalue(L, -2, 3);
+	lua_pop(L,1);
+
 	lua_getfield(L, -1, "yield");
 
 	lua_CFunction co_yield = lua_tocfunction(L, -1);
 	if (co_yield == NULL)
 		return luaL_error(L, "Can't get coroutine.yield");
 	lua_pop(L,1);
+
 	lua_getfield(L, libtable, "yield");
 	lua_pushcfunction(L, co_yield);
 	lua_setupvalue(L, -2, 3);
+	lua_pop(L,1);
+
+	lua_getfield(L, libtable, "yield_co");
+	lua_pushcfunction(L, co_yield);
+	lua_setupvalue(L, -2, 3);
+	lua_pop(L,1);
 
 	lua_settop(L, libtable);
 

lualib/skynet/coroutine.lua

@@ -5,6 +5,7 @@ local coroutine = coroutine
 local coroutine_resume = coroutine.resume
 local coroutine_yield = coroutine.yield
 local coroutine_status = coroutine.status
+local coroutine_running = coroutine.running
 
 local select = select
 local skynetco = {}
@@ -23,30 +24,25 @@ function skynetco.create(f)
 	return co
 end
 
-function skynetco.isskynetcoroutine(co)
-	co = co or coroutine.running()
-	return skynet_coroutines[co] ~= nil
-end
-
 do -- begin skynetco.resume
 
 	local profile = require "profile"
-	-- skynet use profile.resume/yield instead of coroutine.resume/yield
-	-- read skynet.lua for detail
-	local skynet_resume = profile.resume
-	local skynet_yield = profile.yield
+	-- skynet use profile.resume_co/yield_co instead of coroutine.resume/yield
+
+	local skynet_resume = profile.resume_co
+	local skynet_yield = profile.yield_co
 
 	local function unlock(co, ...)
 		skynet_coroutines[co] = true
 		return ...
 	end
 
-	local function skynet_yielding(co, ...)
+	local function skynet_yielding(co, from, ...)
 		skynet_coroutines[co] = false
-		return unlock(co, skynet_resume(co, skynet_yield(...)))
+		return unlock(co, skynet_resume(co, from, skynet_yield(from, ...)))
 	end
 
-	local function resume(co, ok, ...)
+	local function resume(co, from, ok, ...)
 		if not ok then
 			return ok, ...
 		elseif coroutine_status(co) == "dead" then
@@ -57,10 +53,13 @@ do -- begin skynetco.resume
 			return true, select(2, ...)
 		else
 			-- blocked in skynet framework, so raise the yielding message
-			return resume(co, skynet_yielding(co, ...))
+			return resume(co, from, skynet_yielding(co, from, ...))
 		end
 	end
 
+	-- record the root of coroutine caller (It should be a skynet thread)
+	local coroutine_caller = setmetatable({} , { __mode = "kv" })
+
 function skynetco.resume(co, ...)
 	local co_status = skynet_coroutines[co]
 	if not co_status then
@@ -75,7 +74,19 @@ function skynetco.resume(co, ...)
 			return false, "cannot resume none skynet coroutine"
 		end
 	end
-	return resume(co, coroutine_resume(co, ...))
+	local from = coroutine_running()
+	local caller = coroutine_caller[from] or from
+	coroutine_caller[co] = caller
+	return resume(co, caller, coroutine_resume(co, ...))
+end
+
+function skynetco.thread(co)
+	co = co or coroutine_running()
+	if skynet_coroutines[co] ~= nil then
+		return coroutine_caller[co] , false
+	else
+		return co, true
+	end
 end
 
 end -- end of skynetco.resume

test/testcoroutine.lua

@@ -1,6 +1,7 @@
 local skynet = require "skynet"
 -- You should use skynet.coroutine instead of origin coroutine in skynet
 local coroutine = require "skynet.coroutine"
+local profile = require "profile"
 
 local function status(co)
 	repeat
@@ -18,21 +19,35 @@ end
 
 local function test(n)
 	local co = coroutine.running()
-	print ("begin", coroutine.isskynetcoroutine(co))
+	print ("begin", co, coroutine.thread(co))	-- false
 	skynet.fork(status, co)
 	for i=1,n do
 		skynet.sleep(100)
 		coroutine.yield(i)
 	end
-	print "end"
+	print ("end", co)
 end
 
-skynet.start(function()
-	print("Is the main thead a skynet coroutine ?", coroutine.isskynetcoroutine())	-- always false
-	print(coroutine.resume(coroutine.running()))	-- always return false
+local function main()
 	local f = coroutine.wrap(test)
+	coroutine.yield "begin"
 	for i=1,3 do
 		local n = f(5)
 		print("main thread",n)
 	end
+	coroutine.yield "end"
+	print("main thread time:", profile.stop(coroutine.thread()))
+end
+
+skynet.start(function()
+	print("Main thead :", coroutine.thread())	-- true
+	print(coroutine.resume(coroutine.running()))	-- always return false
+
+	profile.start()
+
+	local f = coroutine.wrap(main)
+	print("main step", f())
+	print("main step", f())
+	print("main step", f())
+--	print("main thread time:", profile.stop())
 end)