triton: rewrited context sleeping implementation

Instead of entering working thread into sleep triton saves machine context and stack on sleep and restores context/stack on wakeup.
This saves costly new thread allocation.

3 files changed, 119 insertions, 70 deletions

diff --git a/accel-pppd/triton/triton.c b/accel-pppd/triton/triton.c
index 7a494a5..ce054ce 100644
--- a/accel-pppd/triton/triton.c
+++ b/accel-pppd/triton/triton.c
@@ -5,6 +5,9 @@
 #include <string.h>
 #include <unistd.h>
 #include <limits.h>
+#include <assert.h>
+#include <ucontext.h>
+#include <setjmp.h>
 #include <sys/resource.h>
 
 #include "triton_p.h"
@@ -49,7 +52,9 @@ static struct triton_timer_t ru_timer = {
 };
 struct triton_context_t default_ctx;
 
-static struct triton_context_t __thread *this_ctx;
+static __thread struct triton_context_t *this_ctx;
+static __thread jmp_buf jmp_env;
+static __thread void *thread_frame;
 
 #define log_debug2(fmt, ...)
 
@@ -77,6 +82,7 @@ static void __config_reload(void (*notify)(int))
 }
 
 static void ctx_thread(struct _triton_context_t *ctx);
+
 static void* triton_thread(struct _triton_thread_t *thread)
 {
 	sigset_t set;
@@ -92,21 +98,38 @@ static void* triton_thread(struct _triton_thread_t *thread)
 	sigaddset(&set, SIGUSR1);
 	sigaddset(&set, SIGQUIT);
 
+	thread_frame = __builtin_frame_address(0);
+
 	pthread_mutex_lock(&thread->sleep_lock);
 	pthread_mutex_unlock(&thread->sleep_lock);
 
 	while (1) {
 		spin_lock(&threads_lock);
-		if (!list_empty(&ctx_queue) && !need_config_reload && triton_stat.thread_active <= thread_count) {
-			thread->ctx = list_entry(ctx_queue.next, typeof(*thread->ctx), entry2);
-			log_debug2("thread: %p: dequeued ctx %p\n", thread, thread->ctx);
-			list_del(&thread->ctx->entry2);
-			spin_unlock(&threads_lock);
-			spin_lock(&thread->ctx->lock);
-			thread->ctx->thread = thread;
-			thread->ctx->queued = 0;
-			spin_unlock(&thread->ctx->lock);
-			__sync_sub_and_fetch(&triton_stat.context_pending, 1);
+		if ((!list_empty(&ctx_queue) || !list_empty(&thread->wakeup_list)) && !need_config_reload && triton_stat.thread_active <= thread_count) {
+			if (!list_empty(&thread->wakeup_list)) {
+				thread->ctx = list_entry(thread->wakeup_list.next, typeof(*thread->ctx), entry2);
+				log_debug2("thread: %p: wakeup ctx %p\n", thread, thread->ctx);
+				list_del(&thread->ctx->entry2);
+				spin_unlock(&threads_lock);
+
+				this_ctx = thread->ctx->ud;
+				if (thread->ctx->ud->before_switch)
+					thread->ctx->ud->before_switch(thread->ctx->ud, thread->ctx->bf_arg);
+
+				alloca(thread->ctx->uc->uc_stack.ss_size);
+				memcpy(thread_frame - thread->ctx->uc->uc_stack.ss_size, thread->ctx->uc->uc_stack.ss_sp, thread->ctx->uc->uc_stack.ss_size);
+				setcontext(thread->ctx->uc);
+			} else {
+				thread->ctx = list_entry(ctx_queue.next, typeof(*thread->ctx), entry2);
+				log_debug2("thread: %p: dequeued ctx %p\n", thread, thread->ctx);
+				list_del(&thread->ctx->entry2);
+				spin_unlock(&threads_lock);
+				spin_lock(&thread->ctx->lock);
+				thread->ctx->thread = thread;
+				thread->ctx->queued = 0;
+				spin_unlock(&thread->ctx->lock);
+				__sync_sub_and_fetch(&triton_stat.context_pending, 1);
+			}
 		} else {
 			if (triton_stat.thread_count > thread_count + triton_stat.context_sleeping) {
 				__sync_sub_and_fetch(&triton_stat.thread_active, 1);
@@ -118,7 +141,9 @@ static void* triton_thread(struct _triton_thread_t *thread)
 				_free(thread);
 				return NULL;
 			}
+
 			log_debug2("thread: %p: sleeping\n", thread);
+
 			if (!terminate)
 				list_add(&thread->entry2, &sleep_threads);
 
@@ -149,31 +174,33 @@ static void* triton_thread(struct _triton_thread_t *thread)
 			spin_unlock(&threads_lock);
 		}
 
-		log_debug2("thread %p: ctx=%p %p\n", thread, thread->ctx, thread->ctx ? thread->ctx->thread : NULL);
-		this_ctx = thread->ctx->ud;
-		if (thread->ctx->ud->before_switch)
-			thread->ctx->ud->before_switch(thread->ctx->ud, thread->ctx->bf_arg);
+		if (setjmp(jmp_env) == 0) {
+			log_debug2("thread %p: ctx=%p %p\n", thread, thread->ctx, thread->ctx ? thread->ctx->thread : NULL);
+			this_ctx = thread->ctx->ud;
+			if (thread->ctx->ud->before_switch)
+				thread->ctx->ud->before_switch(thread->ctx->ud, thread->ctx->bf_arg);
 
-cont:
-		log_debug2("thread %p: switch to %p\n", thread, thread->ctx);
-		ctx_thread(thread->ctx);
-		log_debug2("thread %p: switch from %p %p\n", thread, thread->ctx, thread->ctx->thread);
+			while (1) {
+				log_debug2("thread %p: switch to %p\n", thread, thread->ctx);
+				ctx_thread(thread->ctx);
+				log_debug2("thread %p: switch from %p %p\n", thread, thread->ctx, thread->ctx->thread);
 
-		spin_lock(&threads_lock);
-		if (thread->ctx->pending && !thread->ctx->need_free) {
+				spin_lock(&threads_lock);
+				if (!thread->ctx->pending || thread->ctx->need_free)
+					break;
+				spin_unlock(&threads_lock);
+			}
+
+			thread->ctx->thread = NULL;
+			need_free = thread->ctx->need_free;
 			spin_unlock(&threads_lock);
-			goto cont;
-		}
-		thread->ctx->thread = NULL;
-		need_free = thread->ctx->need_free;
-		spin_unlock(&threads_lock);
 
-		if (need_free) {
-			log_debug2("- context %p removed\n", thread->ctx);
-			triton_context_release(thread->ctx);
+			if (need_free) {
+				log_debug2("- context %p removed\n", thread->ctx);
+				triton_context_release(thread->ctx);
+			}
+			thread->ctx = NULL;
 		}
-
-		thread->ctx = NULL;
 	}
 }
 
@@ -272,6 +299,7 @@ struct _triton_thread_t *create_thread()
 	pthread_attr_setstacksize(&attr, WORKER_STACK_SIZE);
 
 	memset(thread, 0, sizeof(*thread));
+	INIT_LIST_HEAD(&thread->wakeup_list);
 	pthread_mutex_init(&thread->sleep_lock, NULL);
 	pthread_cond_init(&thread->sleep_cond, NULL);
 	pthread_mutex_lock(&thread->sleep_lock);
@@ -419,36 +447,6 @@ void __export triton_context_set_priority(struct triton_context_t *ud, int prio)
 	ctx->priority = prio > 0;
 }
 
-void __export triton_context_schedule()
-{
-	struct _triton_context_t *ctx = (struct _triton_context_t *)this_ctx->tpd;
-	struct _triton_thread_t *t = NULL;
-
-	log_debug2("ctx %p: enter schedule\n", ctx);
-	__sync_add_and_fetch(&triton_stat.context_sleeping, 1);
-	__sync_sub_and_fetch(&triton_stat.thread_active, 1);
-	pthread_mutex_lock(&ctx->thread->sleep_lock);
-	while (1) {
-		if (ctx->wakeup) {
-			ctx->wakeup = 0;
-			break;
-		} else {
-			if (!t && triton_stat.thread_count <= thread_count + triton_stat.context_sleeping) {
-				t = create_thread();
-				spin_lock(&threads_lock);
-				list_add_tail(&t->entry, &threads);
-				spin_unlock(&threads_lock);
-				pthread_mutex_unlock(&t->sleep_lock);
-			}
-			pthread_cond_wait(&ctx->thread->sleep_cond, &ctx->thread->sleep_lock);
-		}
-	}
-	pthread_mutex_unlock(&ctx->thread->sleep_lock);
-	__sync_sub_and_fetch(&triton_stat.context_sleeping, 1);
-	__sync_add_and_fetch(&triton_stat.thread_active, 1);
-	log_debug2("ctx %p: exit schedule\n", ctx);
-}
-
 struct triton_context_t __export *triton_context_self(void)
 {
 	return this_ctx;
@@ -462,6 +460,50 @@ void triton_context_print(void)
 		printf("%p\n", ctx);
 }
 
+static ucontext_t *alloc_context()
+{
+	ucontext_t *uc;
+	void *frame = __builtin_frame_address(0);
+	size_t stack_size = thread_frame - frame;
+
+	uc = malloc(sizeof(*uc) + stack_size);
+	uc->uc_stack.ss_sp = (void *)(uc + 1);
+	uc->uc_stack.ss_size = stack_size;
+	memcpy(uc->uc_stack.ss_sp, frame, stack_size);
+
+	return uc;
+}
+
+void __export triton_context_schedule()
+{
+	struct _triton_context_t *ctx = (struct _triton_context_t *)this_ctx->tpd;
+
+	log_debug2("ctx %p: enter schedule\n", ctx);
+	__sync_add_and_fetch(&triton_stat.context_sleeping, 1);
+
+	ctx->uc = alloc_context();
+
+	getcontext(ctx->uc);
+
+	barrier();
+
+	ctx = (struct _triton_context_t *)this_ctx->tpd;
+
+	spin_lock(&threads_lock);
+	if (ctx->wakeup) {
+		spin_unlock(&threads_lock);
+		ctx->wakeup = 0;
+		free(ctx->uc);
+		ctx->uc = NULL;
+		__sync_sub_and_fetch(&triton_stat.context_sleeping, 1);
+		log_debug2("ctx %p: exit schedule\n", ctx);
+	} else {
+		ctx->thread->ctx = NULL;
+		spin_unlock(&threads_lock);
+		longjmp(jmp_env, 1);
+	}
+}
+
 void __export triton_context_wakeup(struct triton_context_t *ud)
 {
 	struct _triton_context_t *ctx = (struct _triton_context_t *)ud->tpd;
@@ -481,26 +523,27 @@ void __export triton_context_wakeup(struct triton_context_t *ud)
 		return;
 	}
 
-	pthread_mutex_lock(&ctx->thread->sleep_lock);
+	spin_lock(&threads_lock);
 	ctx->wakeup = 1;
-	pthread_cond_signal(&ctx->thread->sleep_cond);
-	pthread_mutex_unlock(&ctx->thread->sleep_lock);
+	if (ctx->thread->ctx != ctx) {
+		list_add_tail(&ctx->entry2, &ctx->thread->wakeup_list);
+		r = ctx->thread->ctx == NULL;
+	}
+	spin_unlock(&threads_lock);
+
+	if (r)
+		triton_thread_wakeup(ctx->thread);
 }
 
 int __export triton_context_call(struct triton_context_t *ud, void (*func)(void *), void *arg)
 {
-	struct _triton_context_t *ctx;
+	struct _triton_context_t *ctx = ud ? (struct _triton_context_t *)ud->tpd : (struct _triton_context_t *)default_ctx.tpd;
 	struct _triton_ctx_call_t *call = mempool_alloc(call_pool);
 	int r;
 
 	if (!call)
 		return -1;
 
-	if (ud)
-		ctx = (struct _triton_context_t *)ud->tpd;
-	else
-		ctx = (struct _triton_context_t *)default_ctx.tpd;
-
 	call->func = func;
 	call->arg = arg;
 
diff --git a/accel-pppd/triton/triton.h b/accel-pppd/triton/triton.h
index b66737c..79bbee1 100644
--- a/accel-pppd/triton/triton.h
+++ b/accel-pppd/triton/triton.h
@@ -138,6 +138,8 @@ void triton_terminate(void);
 #define __export __attribute__((visibility("default")))
 #define __unused __attribute__((unused))
 
+#define barrier() asm volatile ("" ::: "memory")
+
 #undef offsetof
 #ifdef __compiler_offsetof
 #define offsetof(TYPE,MEMBER) __compiler_offsetof(TYPE,MEMBER)
diff --git a/accel-pppd/triton/triton_p.h b/accel-pppd/triton/triton_p.h
index 6c9f884..a16ef98 100644
--- a/accel-pppd/triton/triton_p.h
+++ b/accel-pppd/triton/triton_p.h
@@ -2,6 +2,7 @@
 #define TRITON_P_H
 
 #include <pthread.h>
+#include <ucontext.h>
 #include <sys/epoll.h>
 
 #include "triton.h"
@@ -18,6 +19,7 @@ struct _triton_thread_t
 	struct _triton_context_t *ctx;
 	pthread_mutex_t sleep_lock;
 	pthread_cond_t sleep_cond;
+	struct list_head wakeup_list;
 };
 
 struct _triton_context_t
@@ -43,6 +45,8 @@ struct _triton_context_t
 	int priority;
 	int refs;
 
+	ucontext_t *uc;
+
 	struct triton_context_t *ud;
 	void *bf_arg;
 };