Add mainthread api

Author: cloudwu

lualib/service.lua

@@ -743,6 +743,22 @@ end
 
 ltask.running = coroutine_running
 
+do
+	local enter = assert(ltask.mainthread_enter)
+	local leave = assert(ltask.mainthread_leave)
+
+	function ltask.mainthread_run(f, ...)
+		enter()
+		continue_session()
+		local ok, err = pcall(f, ...)
+		leave()
+		continue_session()
+		if not ok then
+			error(err)
+		end
+	end
+end
+
 local yieldable_require; do
 	local require = _G.require
 	local loaded = package.loaded

src/ltask.c

@@ -44,6 +44,43 @@ LUAMOD_API int luaopen_ltask_root(lua_State *L);
 
 #endif
 
+#define MAINTHREAD_STATUS_NONE 0
+#define MAINTHREAD_STATUS_READY 1
+#define MAINTHREAD_STATUS_YIELD 2
+#define MAINTHREAD_STATUS_FINISH 3
+#define MAINTHREAD_STATUS_INVALID 4
+
+struct mainthread_session {
+	service_id srv;
+	struct cond ev;
+	int status;
+};
+
+static void
+mainthread_init(struct mainthread_session *mt) {
+	mt->srv.id = 0;
+	mt->status = MAINTHREAD_STATUS_NONE;
+	cond_create(&mt->ev);
+}
+
+static void
+mainthread_deinit(struct mainthread_session *mt) {
+	assert(mt->status == MAINTHREAD_STATUS_NONE || mt->status == MAINTHREAD_STATUS_YIELD);
+	mt->status = MAINTHREAD_STATUS_INVALID;
+	cond_release(&mt->ev);
+}
+
+static inline int
+mainthread_current(struct mainthread_session *mt, service_id id) {
+	return mt->srv.id == id.id && mt->status == MAINTHREAD_STATUS_READY;
+}
+
+static inline void
+mainthread_trigger(struct mainthread_session *mt) {
+	cond_trigger_begin(&mt->ev);
+	cond_trigger_end(&mt->ev, 1);
+}
+
 struct ltask {
 	const struct ltask_config *config;
 	struct worker_thread *workers;
@@ -58,6 +95,7 @@ struct ltask {
 	struct logqueue *lqueue;
 	struct queue *external_message;
 	struct message *external_last_message;
+	struct mainthread_session mt;
 	atomic_int schedule_owner;
 	atomic_int active_worker;
 	atomic_int thread_count;
@@ -170,7 +208,7 @@ collect_done_job(struct ltask *task, service_id done_job[]) {
 	for (i=0;i<worker_n;i++) {
 		struct worker_thread * w = &task->workers[i];
 		service_id job = worker_done_job(w);
-		if (job.id) {
+		if (job.id && !mainthread_current(&task->mt, job)) {
 			debug_printf(task->logger, "Service %x is done", job.id);
 			done_job[done_job_n++] = job;
 		}
@@ -654,6 +692,8 @@ thread_worker(void *ud) {
 					} else {
 						service_send_signal(P, id);
 					}
+				} else if (mainthread_current(&w->task->mt, id)) {
+					mainthread_trigger(&w->task->mt);
 				} else {
 					service_status_set(P, id, SERVICE_STATUS_DONE);
 				}
@@ -777,6 +817,8 @@ ltask_init(lua_State *L) {
 	}
 
 	task->blocked_service = 0;
+	
+	mainthread_init(&task->mt);
 
 	return 1;
 }
@@ -901,6 +943,7 @@ ltask_wait(lua_State *L) {
 	}
 	message_delete(ctx->task->external_last_message);
 	queue_delete(ctx->task->external_message);
+	mainthread_deinit(&ctx->task->mt);
 	return 0;
 }
 
@@ -1096,6 +1139,53 @@ ltask_log_sender(lua_State *L) {
 	return 2;
 }
 
+// run service in mainthread
+
+static int
+lmainthread_wait(lua_State *L) {
+	struct ltask *task = (struct ltask *)get_ptr(L, "LTASK_GLOBAL");
+	struct mainthread_session * mt = &task->mt;
+	if (mt->status != MAINTHREAD_STATUS_NONE && mt->status != MAINTHREAD_STATUS_READY)
+		return luaL_error(L, "mainthread in use");
+	struct service_pool * P = task->services;
+	int finish = 0;
+	do {
+		cond_wait_begin(&mt->ev);
+		cond_wait(&mt->ev);
+		cond_wait_end(&mt->ev);
+		if (mt->status != MAINTHREAD_STATUS_READY)
+			luaL_error(L, "mainthread not ready");
+		
+		service_id id = mt->srv;
+		
+		debug_printf(task->logger, "service %x run in mainthread", id.id);
+		if (service_resume(P, id)) {
+			// dead
+			debug_printf(task->logger, "service %x is dead in mainthread", id.id);
+			service_status_set(P, id, SERVICE_STATUS_DEAD);
+		} else {
+			service_status_set(P, id, SERVICE_STATUS_SCHEDULE);
+			switch (mt->status) {
+			case MAINTHREAD_STATUS_YIELD:
+				debug_printf(task->logger, "service %x yield from mainthread", id.id);
+				break;
+			case MAINTHREAD_STATUS_FINISH:
+				debug_printf(task->logger, "service %x finish in mainthread", id.id);
+				finish = 1;
+				break;
+			default:
+				debug_printf(task->logger, "service %x is dead in mainthread", id.id);
+				service_status_set(P, id, SERVICE_STATUS_DEAD);
+				schedule_back(task, id);
+				return luaL_error(L, "Can't yield in mainthread (status = %d), kill service (%d)", mt->status, id.id);
+			}
+			schedule_back(task, id);
+		}
+		mt->status = MAINTHREAD_STATUS_NONE;
+	} while (!finish);
+	return 0;
+}
+
 LUAMOD_API int
 luaopen_ltask_bootstrap(lua_State *L) {
 	static atomic_int init = 0;
@@ -1120,6 +1210,7 @@ luaopen_ltask_bootstrap(lua_State *L) {
 		{ "unpack_remove", luaseri_unpack_remove },
 		{ "external_sender", ltask_external_sender },
 		{ "log_sender", ltask_log_sender },
+		{ "mainthread_wait", lmainthread_wait },
 		{ NULL, NULL },
 	};
 
@@ -1563,6 +1654,30 @@ ltask_debuglog(lua_State *L) {
 
 #endif
 
+static int
+mainthread_change_status(lua_State *L, int status) {
+	const struct service_ud *S = getS(L);
+	struct ltask *task = S->task;
+	task->mt.srv = S->id;
+	task->mt.status = status;
+	return 0;
+}
+
+static int
+lmainthread_enter(lua_State *L) {
+	return mainthread_change_status(L, MAINTHREAD_STATUS_READY);
+}
+
+static int
+lmainthread_leave(lua_State *L) {
+	return mainthread_change_status(L, MAINTHREAD_STATUS_FINISH);
+}
+
+static int
+lmainthread_yield(lua_State *L) {
+	return mainthread_change_status(L, MAINTHREAD_STATUS_YIELD);
+}
+
 LUAMOD_API int
 luaopen_ltask(lua_State *L) {
 	luaL_checkversion(L);
@@ -1600,6 +1715,9 @@ luaopen_ltask(lua_State *L) {
 		{ "debuglog", ltask_debuglog },
 		{ "eventinit", ltask_eventinit },
 		{ "eventreset", ltask_eventreset },
+		{ "mainthread_enter", lmainthread_enter },
+		{ "mainthread_leave", lmainthread_leave },
+		{ "mainthread_yield", lmainthread_yield },
 		{ NULL, NULL },
 	};