lkl: host ops: fix jump buffer API and implementation

arch/lkl/include/asm/sched.h

@@ -0,0 +1,27 @@
+#ifndef _ASM_LKL_SCHED_H
+#define _ASM_LKL_SCHED_H
+
+#include <linux/sched.h>
+
+static inline void thread_sched_jb(void)
+{
+	set_ti_thread_flag(current_thread_info(), TIF_SCHED_JB);
+
+	if (test_ti_thread_flag(current_thread_info(), TIF_HOST_THREAD)) {
+		set_current_state(TASK_UNINTERRUPTIBLE);
+		lkl_ops->jmp_buf_set(&current_thread_info()->sched_jb,
+				     schedule);
+	} else  {
+		lkl_ops->jmp_buf_set(&current_thread_info()->sched_jb,
+				     lkl_idle_tail_schedule);
+	}
+}
+
+static inline void thread_set_sched_exit(void)
+{
+	set_ti_thread_flag(current_thread_info(), TIF_SCHED_EXIT);
+}
+
+void switch_to_host_task(struct task_struct *);
+
+#endif /*  _ASM_LKL_SCHED_H */

arch/lkl/include/asm/thread_info.h

@@ -63,21 +63,6 @@ void threads_cnt_dec(void);
 #define TIF_HOST_THREAD			9
 #define TIF_IDLE			10
 
-static inline void set_ti_thread_flag(struct thread_info *ti, int flag);
-
-static inline int thread_set_sched_jmp(void)
-{
-	set_ti_thread_flag(current_thread_info(), TIF_SCHED_JB);
-	return lkl_ops->jmp_buf_set(&current_thread_info()->sched_jb);
-}
-
-static inline void thread_set_sched_exit(void)
-{
-	set_ti_thread_flag(current_thread_info(), TIF_SCHED_EXIT);
-}
-
-void switch_to_host_task(struct task_struct *);
-
 #define __HAVE_THREAD_FUNCTIONS
 
 #define task_thread_info(task)	((struct thread_info *)(task)->stack)

arch/lkl/include/uapi/asm/host_ops.h

@@ -70,6 +70,17 @@ struct lkl_jmp_buf {
  *
  * @gettid - returns the host thread id of the caller, which need not
  * be the same as the handle returned by thread_create
+ *
+ * @jmp_buf_set - runs the give function and setups a jump back point by saving
+ * the context in the jump buffer; jmp_buf_longjmp can be called from the give
+ * function or any callee in that function to return back to the jump back
+ * point
+ *
+ * NOTE: we can't return from jmp_buf_set before calling jmp_buf_longjmp or
+ * otherwise the saved context (stack) is not going to be valid, so we must pass
+ * the function that will eventually call longjmp here
+ *
+ * @jmp_buf_longjmp - perform a jump back to the saved jump buffer
  */
 struct lkl_host_operations {
 	const char *virtio_devices;
@@ -114,7 +125,7 @@ struct lkl_host_operations {
 
 	long (*gettid)(void);
 
-	int (*jmp_buf_set)(struct lkl_jmp_buf *jmpb);
+	void (*jmp_buf_set)(struct lkl_jmp_buf *jmpb, void (*f)(void));
 	void (*jmp_buf_longjmp)(struct lkl_jmp_buf *jmpb, int val);
 };
 

arch/lkl/kernel/cpu.c

@@ -7,6 +7,7 @@
 #include <asm/cpu.h>
 #include <asm/thread_info.h>
 #include <asm/unistd.h>
+#include <asm/sched.h>
 #include <asm/syscalls.h>
 
 
@@ -137,14 +138,7 @@ void lkl_cpu_put(void)
 		if (in_interrupt())
 			lkl_bug("%s: in interrupt\n", __func__);
 		lkl_ops->mutex_unlock(cpu.lock);
-		if (test_thread_flag(TIF_HOST_THREAD)) {
-			set_current_state(TASK_UNINTERRUPTIBLE);
-			if (!thread_set_sched_jmp())
-				schedule();
-		} else {
-			if (!thread_set_sched_jmp())
-				lkl_idle_tail_schedule();
-		}
+		thread_sched_jb();
 		return;
 	}
 
@@ -242,10 +236,8 @@ void arch_cpu_idle_prepare(void)
 	 * We hijack the idle loop here so that we can let the idle thread
 	 * jump back to the beginning.
 	 */
-	while (1) {
-		if (!lkl_ops->jmp_buf_set(&cpu.idle_jb))
-			cpu_idle_loop();
-	}
+	while (1)
+		lkl_ops->jmp_buf_set(&cpu.idle_jb, cpu_idle_loop);
 }
 
 void lkl_cpu_wakeup_idle(void)

arch/lkl/kernel/syscalls.c

@@ -14,6 +14,7 @@
 #include <asm/syscalls.h>
 #include <asm/syscalls_32.h>
 #include <asm/cpu.h>
+#include <asm/sched.h>
 
 static asmlinkage long sys_virtio_mmio_device_add(long base, long size,
 						  unsigned int irq);
@@ -114,9 +115,7 @@ long lkl_syscall(long no, long *params)
 	ret = run_syscall(no, params);
 
 	if (no == __NR_reboot) {
-		set_current_state(TASK_UNINTERRUPTIBLE);
-		if (!thread_set_sched_jmp())
-			schedule();
+		thread_sched_jb();
 		return ret;
 	}
 

arch/lkl/kernel/threads.c

@@ -3,6 +3,7 @@
 #include <linux/sched.h>
 #include <asm/host_ops.h>
 #include <asm/cpu.h>
+#include <asm/sched.h>
 
 static volatile int threads_counter;
 
@@ -140,15 +141,7 @@ void switch_to_host_task(struct task_struct *task)
 	task_thread_info(task)->tid = lkl_ops->thread_self();
 
 	wake_up_process(task);
-	if (test_thread_flag(TIF_HOST_THREAD)) {
-		set_current_state(TASK_UNINTERRUPTIBLE);
-		if (!thread_set_sched_jmp())
-			schedule();
-	} else {
-		if (!thread_set_sched_jmp())
-			lkl_idle_tail_schedule();
-	}
-
+	thread_sched_jb();
 	lkl_ops->sem_down(task_thread_info(task)->sched_sem);
 	schedule_tail(abs_prev);
 }

tools/lkl/lib/jmp_buf.c

@@ -1,9 +1,10 @@
 #include <setjmp.h>
 #include <lkl_host.h>
 
-int jmp_buf_set(struct lkl_jmp_buf *jmpb)
+void jmp_buf_set(struct lkl_jmp_buf *jmpb, void (*f)(void))
 {
-	return setjmp(*((jmp_buf *)jmpb->buf));
+	if (!setjmp(*((jmp_buf *)jmpb->buf)))
+		f();
 }
 
 void jmp_buf_longjmp(struct lkl_jmp_buf *jmpb, int val)

tools/lkl/lib/jmp_buf.h

@@ -1,7 +1,7 @@
 #ifndef _LKL_LIB_JMP_BUF_H
 #define _LKL_LIB_JMP_BUF_H
 
-int jmp_buf_set(struct lkl_jmp_buf *jmpb);
+void jmp_buf_set(struct lkl_jmp_buf *jmpb, void (*f)(void));
 void jmp_buf_longjmp(struct lkl_jmp_buf *jmpb, int val);
 
 #endif