2 files changed, 72 insertions, 16 deletions

diff --git a/kern/thread.c b/kern/thread.c
index 61a2d5b0..959208ab 100644
--- a/kern/thread.c
+++ b/kern/thread.c
@@ -28,15 +28,7 @@
 #include <machine/cpu.h>
 #include <machine/tcb.h>
 
-/*
- * Per processor run queue.
- */
-struct thread_runq {
-    struct thread *current;
-    struct list threads;
-} __aligned(CPU_L1_SIZE);
-
-static struct thread_runq thread_runqs[MAX_CPUS];
+struct thread_runq thread_runqs[MAX_CPUS];
 
 /*
  * Caches for allocated threads and their stacks.
@@ -72,12 +64,6 @@ thread_runq_dequeue(struct thread_runq *runq)
     return thread;
 }
 
-static inline struct thread_runq *
-thread_runq_local(void)
-{
-    return &thread_runqs[cpu_id()];
-}
-
 void __init
 thread_setup(void)
 {
@@ -138,6 +124,7 @@ thread_create(struct thread **threadp, const char *name, struct task *task,
         name = task->name;
 
     thread->flags = 0;
+    thread->preempt = 0;
     thread->task = task;
     thread->stack = stack;
     strlcpy(thread->name, name, sizeof(thread->name));
diff --git a/kern/thread.h b/kern/thread.h
index 52c0f813..3c6395c8 100644
--- a/kern/thread.h
+++ b/kern/thread.h
@@ -18,9 +18,12 @@
 #ifndef _KERN_THREAD_H
 #define _KERN_THREAD_H
 
+#include <kern/assert.h>
 #include <kern/list.h>
 #include <kern/macros.h>
+#include <kern/param.h>
 #include <kern/task.h>
+#include <machine/cpu.h>
 #include <machine/tcb.h>
 
 /*
@@ -38,7 +41,8 @@
  */
 struct thread {
     struct tcb tcb;
-    int flags;
+    short flags;
+    unsigned short preempt;
     struct list runq_node;
     struct list task_node;
     struct task *task;
@@ -49,6 +53,16 @@ struct thread {
 };
 
 /*
+ * Per processor run queue.
+ */
+struct thread_runq {
+    struct thread *current;
+    struct list threads;
+} __aligned(CPU_L1_SIZE);
+
+extern struct thread_runq thread_runqs[MAX_CPUS];
+
+/*
  * Initialize the thread module.
  */
 void thread_setup(void);
@@ -88,4 +102,59 @@ void thread_reschedule(void);
  */
 void thread_tick(void);
 
+static inline struct thread_runq *
+thread_runq_local(void)
+{
+    return &thread_runqs[cpu_id()];
+}
+
+static inline struct thread *
+thread_current(void)
+{
+    return thread_runq_local()->current;
+}
+
+/*
+ * Preemption control functions.
+ */
+
+static inline int
+thread_preempt_enabled(void)
+{
+    return (thread_current()->preempt == 0);
+}
+
+static inline void
+thread_preempt_enable_no_resched(void)
+{
+    struct thread *thread;
+
+    thread = thread_current();
+    assert(thread->preempt != 0);
+    thread->preempt--;
+}
+
+static inline void
+thread_preempt_enable(void)
+{
+    struct thread *thread;
+
+    thread = thread_current();
+    assert(thread->preempt != 0);
+    thread->preempt--;
+
+    if ((thread->preempt == 0) && (thread->flags & THREAD_RESCHEDULE))
+        thread_schedule();
+}
+
+static inline void
+thread_preempt_disable(void)
+{
+    struct thread *thread;
+
+    thread = thread_current();
+    thread->preempt++;
+    assert(thread->preempt != 0);
+}
+
 #endif /* _KERN_THREAD_H */