slew of bug fixes: mostly races and sched-atomic

19 files changed, 594 insertions, 118 deletions

diff --git a/include/litmus/budget.h b/include/litmus/budget.h
index d5d4d61eb06a..0ae9c9f30023 100644
--- a/include/litmus/budget.h
+++ b/include/litmus/budget.h
@@ -11,14 +11,20 @@ struct enforcement_timer
         raw_spinlock_t lock;
         struct hrtimer timer;
         int armed:1;
+        unsigned int job_when_armed;
 };
 
+int cancel_enforcement_timer(struct task_struct* t);
+
 typedef void (*scheduled_t)(struct task_struct* t);
 typedef void (*blocked_t)(struct task_struct* t);
 typedef void (*preempt_t)(struct task_struct* t);
 typedef void (*sleep_t)(struct task_struct* t);
 typedef void (*wakeup_t)(struct task_struct* t);
-typedef enum hrtimer_restart (*exhausted_t)(struct task_struct* t);
+
+#define IN_SCHEDULE 1
+
+typedef enum hrtimer_restart (*exhausted_t)(struct task_struct* t, int in_schedule);
 typedef void (*exit_t)(struct task_struct* t);
 typedef void (*inherit_t)(struct task_struct* t, struct task_struct* prio_inh);
 typedef void (*disinherit_t)(struct task_struct* t, struct task_struct* prio_inh);
@@ -111,7 +117,15 @@ void reevaluate_inheritance(struct task_struct* t);
                 } \
         }while(0)
 
-#define budget_state_machine2(a, b, evt) \
+#define budget_state_machine2(t, evt, param) \
+        do { \
+                if (get_budget_timer(t).ops && \
+                        get_budget_timer(t).ops->evt != NULL) { \
+                        get_budget_timer(t).ops->evt(t, param); \
+                } \
+        }while(0)
+
+#define budget_state_machine_chgprio(a, b, evt) \
         do { \
                 if (get_budget_timer(a).ops && \
                         get_budget_timer(b).ops && \
diff --git a/include/litmus/jobs.h b/include/litmus/jobs.h
index ec73cc6ae7f8..797e1d470a69 100644
--- a/include/litmus/jobs.h
+++ b/include/litmus/jobs.h
@@ -3,6 +3,7 @@
 
 void prepare_for_next_period(struct task_struct *t);
 void release_at(struct task_struct *t, lt_t start);
+void setup_release(struct task_struct *t, lt_t start);
 long complete_job(void);
 
 #if defined(CONFIG_REALTIME_AUX_TASKS) || defined(CONFIG_LITMUS_NVIDIA)
diff --git a/include/litmus/litmus.h b/include/litmus/litmus.h
index 7efd146c2198..cb87dcb5f5f7 100644
--- a/include/litmus/litmus.h
+++ b/include/litmus/litmus.h
@@ -142,7 +142,7 @@ clear_bit(flag_nr, &tsk_rt(t)->budget.flags))
 #define bt_flags_reset(t) (\
 tsk_rt(t)->budget.flags = 0)
 
-#define requeue_preempted_job(t) \
+#define should_requeue_preempted_job(t) \
 (t && !is_completed(t) && (!budget_exhausted(t) || !budget_enforced(t)))
 
 
diff --git a/include/litmus/rt_param.h b/include/litmus/rt_param.h
index c480d95690f8..6b6e19752a81 100644
--- a/include/litmus/rt_param.h
+++ b/include/litmus/rt_param.h
@@ -335,6 +335,8 @@ struct klmirqd_info
         struct list_head                worklist;
 
         struct list_head                klmirqd_reg;
+
+        struct completion*              exited;
 };
 #endif
 
diff --git a/kernel/sched.c b/kernel/sched.c
index 51fb72b5af79..7d27a3851a4d 100644
--- a/kernel/sched.c
+++ b/kernel/sched.c
@@ -91,6 +91,7 @@
 #endif
 
 static void litmus_tick(struct rq*, struct task_struct*);
+static void litmus_handle_budget_exhaustion(struct task_struct*);
 
 /*
  * Convert user-nice values [ -20 ... 0 ... 19 ]
@@ -4432,6 +4433,11 @@ litmus_need_resched_nonpreemptible:
 
         post_schedule(rq);
 
+        if (is_realtime(current) &&
+                unlikely(budget_enforced(current) && budget_exhausted(current))) {
+                litmus_handle_budget_exhaustion(current);
+        }
+
         if (sched_state_validate_switch()) {
                 TS_SCHED2_END(prev);
                 goto litmus_need_resched_nonpreemptible;
diff --git a/litmus/budget.c b/litmus/budget.c
index 9ae530ad3b2b..4f38b3d8d45c 100644
--- a/litmus/budget.c
+++ b/litmus/budget.c
@@ -9,10 +9,10 @@
 #include <litmus/budget.h>
 #include <litmus/signal.h>
 
-inline static void cancel_enforcement_timer(struct task_struct* t)
+int cancel_enforcement_timer(struct task_struct* t)
 {
         struct enforcement_timer* et;
-        int ret;
+        int ret = 0;
         unsigned long flags;
 
         BUG_ON(!t);
@@ -20,22 +20,28 @@ inline static void cancel_enforcement_timer(struct task_struct* t)
 
         et = &tsk_rt(t)->budget.timer;
 
-        TRACE("cancelling enforcement timer.\n");
+        TRACE_TASK(t, "canceling enforcement timer.\n");
 
         if (et->armed) {
                 raw_spin_lock_irqsave(&et->lock, flags);
                 if (et->armed) {
                         ret = hrtimer_try_to_cancel(&et->timer);
-                        et->armed = 0;
-                }
-                else {
-                        TRACE("timer was not armed (race).\n");
+                        if (ret < 0)
+                                TRACE_TASK(t, "timer already running. failed to cancel.\n");
+                        else {
+                                TRACE_TASK(t, "canceled timer with %lld ns remaining.\n",
+                                        ktime_to_ns(hrtimer_expires_remaining(&et->timer)));
+                                et->armed = 0;
+                        }
                 }
+                else
+                        TRACE_TASK(t, "timer was not armed (race).\n");
                 raw_spin_unlock_irqrestore(&et->lock, flags);
         }
-        else {
-                TRACE("timer was not armed.\n");
-        }
+        else
+                TRACE_TASK(t, "timer was not armed.\n");
+
+        return ret;
 }
 
 
@@ -283,15 +289,14 @@ void sobliv_on_wakeup(struct task_struct* t)
 void sobliv_on_inherit(struct task_struct* t, struct task_struct* prio_inh)
 {
         /* TODO: Budget credit accounting. */
-
         BUG_ON(!prio_inh);
-        TRACE_TASK(t, "called %s\n", __FUNCTION__);
+//      TRACE_TASK(t, "called %s\n", __FUNCTION__);
 }
 
 void sobliv_on_disinherit(struct task_struct* t, struct task_struct* prio_inh)
 {
         /* TODO: Budget credit accounting. */
-        TRACE_TASK(t, "called %s\n", __FUNCTION__);
+//      TRACE_TASK(t, "called %s\n", __FUNCTION__);
 }
 
 void sobliv_on_enter_top_m(struct task_struct* t)
@@ -397,7 +402,7 @@ void reevaluate_inheritance(struct task_struct* t)
 
 static enum hrtimer_restart __on_timeout(struct hrtimer *timer)
 {
-        enum hrtimer_restart restart;
+        enum hrtimer_restart restart = HRTIMER_NORESTART;
         unsigned long flags;
 
         struct budget_tracker* bt =
@@ -420,12 +425,18 @@ static enum hrtimer_restart __on_timeout(struct hrtimer *timer)
         tsk_rt(t)->budget.timer.armed = 0;
         raw_spin_unlock_irqrestore(&bt->timer.lock, flags);
 
-        restart = bt->ops->on_exhausted(t);
+        if (unlikely(bt_flag_is_set(t, BTF_WAITING_FOR_RELEASE))) {
+                TRACE_TASK(t, "spurious exhastion while waiting for release. dropping.\n");
+                goto out;
+        }
+
+        restart = bt->ops->on_exhausted(t,!IN_SCHEDULE);
 
         raw_spin_lock_irqsave(&bt->timer.lock, flags);
         tsk_rt(t)->budget.timer.armed = (restart == HRTIMER_RESTART);
         raw_spin_unlock_irqrestore(&bt->timer.lock, flags);
 
+out:
         return restart;
 }
 
diff --git a/litmus/edf_common.c b/litmus/edf_common.c
index 7f0bdaceab5b..5519e05aaf83 100644
--- a/litmus/edf_common.c
+++ b/litmus/edf_common.c
@@ -63,7 +63,7 @@ int edf_higher_prio(struct task_struct* first, struct task_struct* second)
         /* There is no point in comparing a task to itself. */
         if (first && first == second) {
                 TRACE_CUR("WARNING: pointless edf priority comparison: %s/%d\n", first->comm, first->pid);
-                WARN_ON(1);
+//              WARN_ON(1);
                 return 0;
         }
 
diff --git a/litmus/ikglp_lock.c b/litmus/ikglp_lock.c
index 4963820eaee2..fb0080352a9f 100644
--- a/litmus/ikglp_lock.c
+++ b/litmus/ikglp_lock.c
@@ -5,6 +5,8 @@
 #include <litmus/sched_plugin.h>
 #include <litmus/fdso.h>
 
+#include <litmus/litmus_proc.h>
+
 #if defined(CONFIG_LITMUS_AFFINITY_LOCKING) && defined(CONFIG_LITMUS_NVIDIA)
 #include <litmus/gpu_affinity.h>
 #include <litmus/nvidia_info.h>
@@ -326,6 +328,159 @@ static void print_donors(struct binheap_node *n, int depth)
 }
 #endif
 
+struct ikglp_proc_print_heap_args
+{
+        struct ikglp_semaphore *sem;
+        int *size;
+        char **next;
+};
+
+static void __ikglp_pq_to_proc(struct binheap_node *n, void *args)
+{
+        struct ikglp_proc_print_heap_args *hargs;
+        ikglp_heap_node_t *request;
+        int w;
+
+        if (!n)
+                return;
+
+        hargs = (struct ikglp_proc_print_heap_args*) args;
+        request = binheap_entry(n, ikglp_heap_node_t, node);
+
+        w = scnprintf(*(hargs->next), *(hargs->size), "\t%s/%d\n",
+                                  request->task->comm, request->task->pid);
+        *(hargs->size) -= w;
+        *(hargs->next) += w;
+}
+
+static void __ikglp_donor_to_proc(struct binheap_node *n, void *args)
+{
+        struct ikglp_proc_print_heap_args *hargs;
+        ikglp_wait_state_t *donor_node;
+        int w;
+
+        if (!n)
+                return;
+
+        hargs = (struct ikglp_proc_print_heap_args*) args;
+        donor_node = binheap_entry(n, ikglp_wait_state_t, node);
+
+        w = scnprintf(*(hargs->next), *(hargs->size), "\t%s/%d (donee: %s/%d)\n",
+                                  donor_node->task->comm,
+                                  donor_node->task->pid,
+                                  donor_node->donee_info->task->comm,
+                                  donor_node->donee_info->task->pid);
+        *(hargs->size) -= w;
+        *(hargs->next) += w;
+}
+
+
+static int ikglp_proc_print(char *page, char **start, off_t off, int count, int *eof, void *data)
+{
+        struct ikglp_semaphore *sem = ikglp_from_lock((struct litmus_lock*)data);
+
+        int attempts = 0;
+        const int max_attempts = 10;
+        int locked = 0;
+        unsigned long flags;
+
+        int size = count;
+        char *next = page;
+
+        struct ikglp_proc_print_heap_args heap_args = {sem, &size, &next};
+
+        int w;
+        int i;
+
+        while(attempts < max_attempts)
+        {
+                locked = raw_spin_trylock_irqsave(&sem->real_lock, flags);
+
+                if (unlikely(!locked)) {
+                        ++attempts;
+                        cpu_relax();
+                }
+                else {
+                        break;
+                }
+        }
+
+        if (!locked) {
+                w = scnprintf(next, size, "%s is busy.\n", sem->litmus_lock.name);
+                size -= w;
+                next += w;
+                return count - size;
+        }
+
+        w = scnprintf(next, size, "nr_replicas: %u\n", sem->nr_replicas); size -= w; next += w;
+        w = scnprintf(next, size, "max_fifo_len: %u\n", sem->max_fifo_len); size -= w; next += w;
+        w = scnprintf(next, size, "max_in_fifos: %u\n", sem->max_in_fifos); size -= w; next += w;
+        w = scnprintf(next, size, "current nr_in_fifos: %u\n", sem->nr_in_fifos); size -= w; next += w;
+        w = scnprintf(next, size, "nr in top-m: %u\n\n", sem->top_m_size); size -= w; next += w;
+
+        for (i = 0; i < sem->nr_replicas; ++i)
+        {
+                struct fifo_queue *fq = &sem->fifo_queues[i];
+                w = scnprintf(next, size, "replica %d: owner = %s/%d (Virtually Unlocked = %u), hp waiter = %s/%d, length = %u\n",
+                                          i,
+                                          (fq->owner) ? fq->owner->comm : "null",
+                                          (fq->owner) ? fq->owner->pid : 0,
+                                          fq->is_vunlocked,
+                                          (fq->hp_waiter) ? fq->hp_waiter->comm : "null",
+                                          (fq->hp_waiter) ? fq->hp_waiter->pid : 0,
+                                          fq->count);
+                size -= w; next += w;
+
+
+                if (waitqueue_active(&fq->wait)) {
+                        struct list_head *pos;
+                        list_for_each(pos, &fq->wait.task_list) {
+                                wait_queue_t *q = list_entry(pos, wait_queue_t, task_list);
+                                struct task_struct *blocked_task = (struct task_struct*) q->private;
+                                w = scnprintf(next, size,
+                                                          "\t%s/%d (inh: %s/%d)\n",
+                                                          blocked_task->comm, blocked_task->pid,
+                                                          (tsk_rt(blocked_task)->inh_task) ?
+                                                                tsk_rt(blocked_task)->inh_task->comm : "null",
+                                                          (tsk_rt(blocked_task)->inh_task) ?
+                                                                tsk_rt(blocked_task)->inh_task->pid : 0);
+                                size -= w;
+                                next += w;
+                        }
+                }
+                else {
+                        w = scnprintf(next, size, "\t<NONE>\n");
+                        size -= w;
+                        next += w;
+                }
+        }
+
+        if (binheap_empty(&sem->priority_queue)) {
+                w = scnprintf(next, size, "pq:\n\t<NONE>\n");
+                size -= w;
+                next += w;
+        }
+        else {
+                w = scnprintf(next, size, "donors:\n"); size -= w; next += w;
+                binheap_for_each(&sem->priority_queue, __ikglp_pq_to_proc, &heap_args);
+        }
+
+        if (binheap_empty(&sem->donors)) {
+                w = scnprintf(next, size, "donors:\n\t<NONE>\n");
+                size -= w;
+                next += w;
+        }
+        else {
+                w = scnprintf(next, size, "donors:\n"); size -= w; next += w;
+                binheap_for_each(&sem->donors, __ikglp_donor_to_proc, &heap_args);
+        }
+
+        raw_spin_unlock_irqrestore(&sem->real_lock, flags);
+
+        return count - size;
+}
+
+
 static void ikglp_add_global_list(struct ikglp_semaphore *sem,
                                                                   struct task_struct *t,
                                                                   ikglp_heap_node_t *node)
@@ -2011,6 +2166,32 @@ void ikglp_free(struct litmus_lock* l)
 
 
 
+static void ikglp_proc_add(struct litmus_lock *l)
+{
+        if (!l->name)
+                l->name = kmalloc(LOCK_NAME_LEN*sizeof(char), GFP_KERNEL);
+        snprintf(l->name, LOCK_NAME_LEN, "ikglp-%d", l->ident);
+        litmus_add_proc_lock(l, ikglp_proc_print);
+}
+
+static void ikglp_proc_remove(struct litmus_lock *l)
+{
+        if (l->name) {
+                litmus_remove_proc_lock(l);
+
+                kfree(l->name);
+                l->name = NULL;
+        }
+}
+
+
+static struct litmus_lock_proc_ops ikglp_proc_ops =
+{
+        .add = ikglp_proc_add,
+        .remove = ikglp_proc_remove
+};
+
+
 struct litmus_lock* ikglp_new(unsigned int m,
                                                           struct litmus_lock_ops* ops,
                                                           void* __user uarg)
@@ -2067,6 +2248,7 @@ struct litmus_lock* ikglp_new(unsigned int m,
         }
 
         sem->litmus_lock.ops = ops;
+        sem->litmus_lock.proc = &ikglp_proc_ops;
 
 #ifdef CONFIG_DEBUG_SPINLOCK
         {
diff --git a/litmus/jobs.c b/litmus/jobs.c
index 018b4982afb5..ecfaf4cd477a 100644
--- a/litmus/jobs.c
+++ b/litmus/jobs.c
@@ -6,7 +6,8 @@
 #include <litmus/litmus.h>
 #include <litmus/jobs.h>
 
-static inline void setup_release(struct task_struct *t, lt_t release)
+//static inline void setup_release(struct task_struct *t, lt_t release)
+void setup_release(struct task_struct *t, lt_t release)
 {
         /* prepare next release */
         t->rt_param.job_params.release = release;
diff --git a/litmus/litmus.c b/litmus/litmus.c
index 8bb05ab0e021..112e269a6348 100644
--- a/litmus/litmus.c
+++ b/litmus/litmus.c
@@ -654,6 +654,30 @@ out:
         return ret;
 }
 
+
+
+#ifdef CONFIG_LITMUS_SOFTIRQD
+struct kill_workers_data
+{
+        struct work_struct work;
+};
+
+static void __kill_workers(struct work_struct *w)
+{
+        struct kill_workers_data *work =
+                container_of(w, struct kill_workers_data, work);
+
+        TRACE_TASK(current, "is processing litmus-worker-shutdown request.\n");
+
+#ifdef CONFIG_LITMUS_NVIDIA
+        shutdown_nvidia_info();
+#endif
+        if (!klmirqd_is_dead())
+                kill_klmirqd();
+        kfree(work);
+}
+#endif
+
 /* Switching a plugin in use is tricky.
  * We must watch out that no real-time tasks exists
  * (and that none is created in parallel) and that the plugin is not
@@ -661,18 +685,59 @@ out:
  */
 int switch_sched_plugin(struct sched_plugin* plugin)
 {
+        int ret;
+        lt_t maybe_deadlock;
+        lt_t spin_time = 100000000; /* 100 ms */
+
         BUG_ON(!plugin);
 
 #ifdef CONFIG_LITMUS_SOFTIRQD
         if (!klmirqd_is_dead()) {
-                kill_klmirqd();
+                struct kill_workers_data *wq_job =
+                                kmalloc(sizeof(struct kill_workers_data), GFP_ATOMIC);
+                INIT_WORK(&wq_job->work, __kill_workers);
+
+                schedule_work(&wq_job->work);
+
+                /* we're atomic, so try to spin until workers are shut down */
+                maybe_deadlock = litmus_clock();
+                while (!klmirqd_is_dead()) {
+                        cpu_relax();
+                        mb();
+
+                        if (lt_before(maybe_deadlock + spin_time, litmus_clock())) {
+                                printk("Could not kill klmirqd! Try again if running "
+                                           "on uniprocessor.\n");
+                                ret = -EBUSY;
+                                goto out;
+                        }
+                }
+
+                TRACE("klmirqd dead! task count = %d\n", atomic_read(&rt_task_count));
         }
+#else
+#ifdef CONFIG_LITMUS_NVIDIA
+        /* nvidia handling is not threaded */
+        shutdown_nvidia_info();
+#endif
 #endif
 
-        if (atomic_read(&rt_task_count) == 0)
-                return stop_machine(do_plugin_switch, plugin, NULL);
-        else
-                return -EBUSY;
+        maybe_deadlock = litmus_clock();
+        while (atomic_read(&rt_task_count) != 0) {
+                cpu_relax();
+                mb();
+
+                if (lt_before(maybe_deadlock + spin_time, litmus_clock())) {
+                        ret = -EBUSY;
+                        goto out;
+                }
+        }
+
+        ret = stop_machine(do_plugin_switch, plugin, NULL);
+        TRACE("stop_machine returned: %d\n", ret);
+
+out:
+        return ret;
 }
 
 /* Called upon fork.
diff --git a/litmus/litmus_softirq.c b/litmus/litmus_softirq.c
index 8bee9697dabd..fbc6ab7cc5f5 100644
--- a/litmus/litmus_softirq.c
+++ b/litmus/litmus_softirq.c
@@ -6,6 +6,7 @@
 #include <linux/smp.h>
 #include <linux/slab.h>
 #include <linux/mutex.h>
+#include <linux/completion.h>
 
 #include <linux/sched.h>
 #include <linux/cpuset.h>
@@ -29,9 +30,10 @@ struct klmirqd_registration
 {
         raw_spinlock_t lock;
         u32 nr_threads;
+        struct list_head threads;
+
         unsigned int initialized:1;
         unsigned int shuttingdown:1;
-        struct list_head threads;
 };
 
 static atomic_t klmirqd_id_gen = ATOMIC_INIT(-1);
@@ -92,13 +94,18 @@ void kill_klmirqd(void)
 
                         if(info->terminating != 1)
                         {
+                                struct completion exit;
+                                init_completion(&exit);
+
                                 info->terminating = 1;
+                                info->exited = &exit;
                                 mb(); /* just to be sure? */
                                 flush_pending(info->klmirqd);
 
                                 /* signal termination */
                                 raw_spin_unlock_irqrestore(&klmirqd_state.lock, flags);
                                 kthread_stop(info->klmirqd);
+                                wait_for_completion(&exit); /* gets completed when task exits rt-mode */
                                 raw_spin_lock_irqsave(&klmirqd_state.lock, flags);
                         }
                 }
@@ -126,14 +133,21 @@ void kill_klmirqd_thread(struct task_struct* klmirqd_thread)
         info = tsk_rt(klmirqd_thread)->klmirqd_info;
 
         if(info->terminating != 1) {
+                struct completion exit;
+                init_completion(&exit);
                 info->terminating = 1;
+                info->exited = &exit;
                 mb();
 
                 flush_pending(klmirqd_thread);
+                raw_spin_unlock_irqrestore(&klmirqd_state.lock, flags);
+
                 kthread_stop(klmirqd_thread);
+                wait_for_completion(&exit); /* gets completed when task exits rt-mode */
+        }
+        else {
+                raw_spin_unlock_irqrestore(&klmirqd_state.lock, flags);
         }
-
-        raw_spin_unlock_irqrestore(&klmirqd_state.lock, flags);
 }
 
 struct klmirqd_launch_data
@@ -157,8 +171,6 @@ static void __launch_klmirqd_thread(struct work_struct *work)
 
     TRACE("Creating klmirqd thread\n");
 
-
-
         if (launch_data->cpu_affinity != -1) {
                 if (launch_data->name[0] == '\0') {
                         id = atomic_inc_return(&klmirqd_id_gen);
@@ -295,8 +307,9 @@ static int become_litmus_daemon(struct task_struct* tsk)
 static int become_normal_daemon(struct task_struct* tsk)
 {
         int ret = 0;
-
         struct sched_param param = { .sched_priority = 0};
+
+        TRACE_TASK(tsk, "exiting real-time mode\n");
         sched_setscheduler_nocheck(tsk, SCHED_NORMAL, &param);
 
         return ret;
@@ -315,22 +328,25 @@ static int register_klmirqd(struct task_struct* tsk)
                 goto out;
         }
 
+        /* allocate and initialize klmirqd data for the thread */
+        info = kmalloc(sizeof(struct klmirqd_info), GFP_KERNEL);
+        if (!info) {
+                TRACE("Failed to allocate klmirqd_info struct!\n");
+                retval = -1; /* todo: pick better code */
+                goto out;
+        }
+
         raw_spin_lock_irqsave(&klmirqd_state.lock, flags);
 
         if (!__klmirqd_is_ready()) {
                 TRACE("klmirqd is not ready! Did you forget to initialize it?\n");
+                kfree(info);
                 WARN_ON(1);
                 retval = -1;
                 goto out_unlock;
         }
 
         /* allocate and initialize klmirqd data for the thread */
-        info = kmalloc(sizeof(struct klmirqd_info), GFP_KERNEL);
-        if (!info) {
-                TRACE("Failed to allocate klmirqd_info struct!\n");
-                retval = -1; /* todo: pick better code */
-                goto out_unlock;
-        }
         memset(info, 0, sizeof(struct klmirqd_info));
         info->klmirqd = tsk;
         info->pending_tasklets_hi.tail = &info->pending_tasklets_hi.head;
@@ -338,7 +354,7 @@ static int register_klmirqd(struct task_struct* tsk)
         INIT_LIST_HEAD(&info->worklist);
         INIT_LIST_HEAD(&info->klmirqd_reg);
         raw_spin_lock_init(&info->lock);
-
+        info->exited = NULL;
 
         /* now register with klmirqd */
         list_add_tail(&info->klmirqd_reg, &klmirqd_state.threads);
@@ -360,8 +376,9 @@ static int unregister_klmirqd(struct task_struct* tsk)
         unsigned long flags;
         struct klmirqd_info *info = tsk_rt(tsk)->klmirqd_info;
 
+        TRACE_CUR("unregistering.\n");
+
         if (!tsk_rt(tsk)->is_interrupt_thread || !info) {
-                TRACE("%s/%d is not a klmirqd thread!\n", tsk->comm, tsk->pid);
                 WARN_ON(1);
                 retval = -1;
                 goto out;
@@ -734,6 +751,7 @@ static int run_klmirqd(void* callback)
 {
     int retval = 0;
         struct klmirqd_info* info = NULL;
+        struct completion* exit = NULL;
         klmirqd_callback_t* cb = (klmirqd_callback_t*)(callback);
 
         retval = become_litmus_daemon(current);
@@ -828,11 +846,18 @@ static int run_klmirqd(void* callback)
 
 failed_unregister:
         /* remove our registration from klmirqd */
+        exit = info->exited;
+        mb();
         unregister_klmirqd(current);
 
 failed_sched_normal:
         become_normal_daemon(current);
 
+        if (exit) {
+                TRACE_TASK(current, "signalling exit\n");
+                complete(exit);
+        }
+
 failed:
         return retval;
 }
@@ -1008,7 +1033,7 @@ int __litmus_tasklet_schedule(struct tasklet_struct *t, struct task_struct* klmi
         if (unlikely(!is_realtime(klmirqd_thread) ||
                 !tsk_rt(klmirqd_thread)->is_interrupt_thread ||
                 !tsk_rt(klmirqd_thread)->klmirqd_info)) {
-                TRACE("%s: %s/%d can't handle tasklets\n", klmirqd_thread->comm, klmirqd_thread->pid);
+                TRACE("can't handle tasklets\n");
                 return ret;
         }
 
diff --git a/litmus/nvidia_info.c b/litmus/nvidia_info.c
index 3951fc92ad03..0050ce65e521 100644
--- a/litmus/nvidia_info.c
+++ b/litmus/nvidia_info.c
@@ -259,47 +259,67 @@ static struct notifier_block nvidia_going = {
         .priority = 1,
 };
 
-int init_nvidia_info(void)
+
+struct init_nvinfo_wq_data
+{
+        struct work_struct work;
+};
+
+static void __init_nvidia_info(struct work_struct *w)
 {
+        struct init_nvinfo_wq_data *work =
+                        container_of(w, struct init_nvinfo_wq_data, work);
+        struct module* mod;
+
         mutex_lock(&module_mutex);
-        nvidia_mod = find_module("nvidia");
+        mod = find_module("nvidia");
         mutex_unlock(&module_mutex);
-        if(nvidia_mod != NULL)
-        {
+
+        if(mod != NULL) {
                 TRACE("%s : Found NVIDIA module. Core Code: %p to %p\n", __FUNCTION__,
-                          (void*)(nvidia_mod->module_core),
-                          (void*)(nvidia_mod->module_core) + nvidia_mod->core_size);
-                init_nv_device_reg();
+                          (void*)(mod->module_core),
+                          (void*)(mod->module_core) + mod->core_size);
 
+                init_nv_device_reg();
+                nvidia_mod = mod; /* make module visible to others */
                 register_module_notifier(&nvidia_going);
-
-                return(0);
         }
-        else
-        {
+        else {
                 TRACE("%s : Could not find NVIDIA module!  Loaded?\n", __FUNCTION__);
-
                 init_nv_device_reg();
-                return(0);
         }
+
+        kfree(work);
+}
+
+int init_nvidia_info(void)
+{
+        struct init_nvinfo_wq_data *wq_job =
+                kmalloc(sizeof(struct init_nvinfo_wq_data), GFP_ATOMIC);
+        INIT_WORK(&wq_job->work, __init_nvidia_info);
+        schedule_work(&wq_job->work);
+        return 0;
 }
 
 void shutdown_nvidia_info(void)
 {
-        nvidia_mod = NULL;
-        mb();
+        if (nvidia_mod) {
+                nvidia_mod = NULL;
+                mb();
 
-        unregister_module_notifier(&nvidia_going);
-        shutdown_nv_device_reg();
+                unregister_module_notifier(&nvidia_going);
+                shutdown_nv_device_reg();
+        }
 }
 
 /* works with pointers to static data inside the module too. */
 int is_nvidia_func(void* func_addr)
 {
         int ret = 0;
-        if(nvidia_mod)
+        struct module* mod = nvidia_mod;
+        if(mod)
         {
-                ret = within_module_core((long unsigned int)func_addr, nvidia_mod);
+                ret = within_module_core((long unsigned int)func_addr, mod);
                 /*
                 if(ret)
                 {
@@ -485,7 +505,6 @@ static int init_nv_device_reg(void)
         memset(NV_DEVICE_REG, 0, sizeof(NV_DEVICE_REG));
         mb();
 
-
         for(i = 0; i < num_online_gpus(); ++i) {
                 raw_spin_lock_init(&NV_DEVICE_REG[i].lock);
                 INIT_BINHEAP_HANDLE(&NV_DEVICE_REG[i].owners, gpu_owner_max_priority_order);
@@ -558,18 +577,20 @@ static int shutdown_nv_device_reg(void)
                            )
                         {
                                 raw_spin_lock_irqsave(&reg->lock, flags);
-
                                 if (reg->interrupt_thread && reg->interrupt_ready) {
                                         struct task_struct* th = reg->interrupt_thread;
                                         reg->interrupt_thread = NULL;
                                         mb();
                                         reg->interrupt_ready = 0;
                                         mb();
-
+                                        raw_spin_unlock_irqrestore(&reg->lock, flags);
                                         kill_klmirqd_thread(th);
                                 }
+                                else
+                                        raw_spin_unlock_irqrestore(&reg->lock, flags);
 
 #ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+                                raw_spin_lock_irqsave(&reg->lock, flags);
                                 if (reg->workq_thread && reg->workq_ready) {
                                         struct task_struct* th = reg->workq_thread;
                                         reg->workq_thread = NULL;
@@ -577,10 +598,12 @@ static int shutdown_nv_device_reg(void)
                                         reg->workq_ready = 0;
                                         mb();
 
+                                        raw_spin_unlock_irqrestore(&reg->lock, flags);
                                         kill_klmirqd_thread(th);
                                 }
+                                else
+                                        raw_spin_unlock_irqrestore(&reg->lock, flags);
 #endif
-                                raw_spin_unlock_irqrestore(&reg->lock, flags);
                         }
 
                         while (!binheap_empty(&reg->owners)) {
@@ -642,6 +665,8 @@ static struct task_struct* __get_klm_thread(nv_device_registry_t* reg, nvklmtype
                         klmirqd = reg->workq_thread;
                 break;
 #endif
+        default:
+                break;
         }
 
         return klmirqd;
@@ -671,9 +696,22 @@ static void __unlock_klm_thread(nv_device_registry_t* reg, unsigned long* flags,
 struct task_struct* get_and_lock_nvklmirqd_thread(u32 target_device_id, unsigned long* flags)
 {
         nv_device_registry_t *reg;
+        struct task_struct *th;
         BUG_ON(target_device_id >= NV_DEVICE_NUM);
+
+        if (unlikely(nvidia_mod == NULL))
+                return NULL;
+
         reg = &NV_DEVICE_REG[target_device_id];
-        return __get_and_lock_klm_thread(reg, flags, INTERRUPT_TH);
+        th = __get_and_lock_klm_thread(reg, flags, INTERRUPT_TH);
+
+        barrier();
+        if (unlikely(nvidia_mod == NULL)) {
+                th = NULL;
+                __unlock_klm_thread(reg, flags, INTERRUPT_TH);
+        }
+
+        return th;
 }
 
 void unlock_nvklmirqd_thread(u32 target_device_id, unsigned long* flags)
@@ -700,9 +738,22 @@ struct task_struct* get_nvklmirqd_thread(u32 target_device_id)
 struct task_struct* get_and_lock_nvklmworkqd_thread(u32 target_device_id, unsigned long* flags)
 {
         nv_device_registry_t *reg;
+        struct task_struct *th;
         BUG_ON(target_device_id >= NV_DEVICE_NUM);
+
+        if (unlikely(nvidia_mod == NULL))
+                return NULL;
+
         reg = &NV_DEVICE_REG[target_device_id];
-        return __get_and_lock_klm_thread(reg, flags, WORKQ_TH);
+        th = __get_and_lock_klm_thread(reg, flags, WORKQ_TH);
+
+        barrier();
+        if (unlikely(nvidia_mod == NULL)) {
+                th = NULL;
+                __unlock_klm_thread(reg, flags, WORKQ_TH);
+        }
+
+        return th;
 }
 
 void unlock_nvklmworkqd_thread(u32 target_device_id, unsigned long* flags)
diff --git a/litmus/sched_cedf.c b/litmus/sched_cedf.c
index 50eb4d446303..1a7f0ffdae33 100644
--- a/litmus/sched_cedf.c
+++ b/litmus/sched_cedf.c
@@ -527,11 +527,11 @@ static void check_for_preemptions(cedf_domain_t *cluster)
                                                 &per_cpu(cedf_cpu_entries, task_cpu(task)));
                         if(affinity)
                                 last = affinity;
-                        else if(requeue_preempted_job(last->linked))
+                        else if(should_requeue_preempted_job(last->linked))
                                 requeue(last->linked);
                 }
 #else
-                if (requeue_preempted_job(last->linked))
+                if (should_requeue_preempted_job(last->linked))
                         requeue(last->linked);
 #endif
                 link_task_to_cpu(task, last);
@@ -674,39 +674,56 @@ static noinline void job_completion(struct task_struct *t, int forced)
         }
 }
 
-static enum hrtimer_restart cedf_simple_on_exhausted(struct task_struct *t)
+static enum hrtimer_restart cedf_simple_on_exhausted(struct task_struct *t, int in_schedule)
 {
         /* Assumption: t is scheduled on the CPU executing this callback */
 
+        if (in_schedule) {
+                BUG_ON(tsk_rt(t)->scheduled_on != smp_processor_id());
+                if (budget_precisely_tracked(t) && cancel_enforcement_timer(t) < 0) {
+                        TRACE_TASK(t, "raced with timer. deffering to timer.\n");
+                        goto out;
+                }
+        }
+
         if (budget_signalled(t) && !bt_flag_is_set(t, BTF_SIG_BUDGET_SENT)) {
                 /* signal exhaustion */
                 send_sigbudget(t); /* will set BTF_SIG_BUDGET_SENT */
         }
 
         if (budget_enforced(t) && !bt_flag_test_and_set(t, BTF_BUDGET_EXHAUSTED)) {
-                if (!is_np(t)) {
+                if (likely(!is_np(t))) {
                         /* np tasks will be preempted when they become
                          * preemptable again
                          */
-                        TRACE_TASK(t, "is preemptable => FORCE_RESCHED\n");
-
-                        litmus_reschedule_local();
-                        set_will_schedule();
+                        if (!in_schedule) {
+                                TRACE_TASK(t, "is preemptable => FORCE_RESCHED\n");
+                                litmus_reschedule_local();
+                                set_will_schedule();
+                        }
                 } else if (is_user_np(t)) {
                         TRACE_TASK(t, "is non-preemptable, preemption delayed.\n");
-
                         request_exit_np(t);
                 }
         }
 
+out:
         return HRTIMER_NORESTART;
 }
 
 
-static enum hrtimer_restart cedf_simple_io_on_exhausted(struct task_struct *t)
+static enum hrtimer_restart cedf_simple_io_on_exhausted(struct task_struct *t, int in_schedule)
 {
         enum hrtimer_restart restart = HRTIMER_NORESTART;
 
+        if (in_schedule) {
+                BUG_ON(tsk_rt(t)->scheduled_on != smp_processor_id());
+                if (budget_precisely_tracked(t) && cancel_enforcement_timer(t) == -1) {
+                        TRACE_TASK(t, "raced with timer. deffering to timer.\n");
+                        goto out;
+                }
+        }
+
         /* t may or may not be scheduled */
 
         if (budget_signalled(t) && !bt_flag_is_set(t, BTF_SIG_BUDGET_SENT)) {
@@ -720,24 +737,28 @@ static enum hrtimer_restart cedf_simple_io_on_exhausted(struct task_struct *t)
 
         if (budget_enforced(t) && !bt_flag_is_set(t, BTF_BUDGET_EXHAUSTED)) {
                 int cpu = (tsk_rt(t)->linked_on != NO_CPU) ?
-                tsk_rt(t)->linked_on : tsk_rt(t)->scheduled_on;
+                        tsk_rt(t)->linked_on : tsk_rt(t)->scheduled_on;
 
                 if (is_np(t) && is_user_np(t)) {
-                        TRACE_TASK(t, "is non-preemptable, preemption delayed.\n");
                         bt_flag_set(t, BTF_BUDGET_EXHAUSTED);
+                        TRACE_TASK(t, "is non-preemptable, preemption delayed.\n");
                         request_exit_np(t);
                 }
                 /* where do we need to call resched? */
                 else if (cpu == smp_processor_id()) {
-                        TRACE_TASK(t, "is preemptable => FORCE_RESCHED\n");
                         bt_flag_set(t, BTF_BUDGET_EXHAUSTED);
-                        litmus_reschedule_local();
-                        set_will_schedule();
+                        if (!in_schedule) {
+                                TRACE_TASK(t, "is preemptable => FORCE_RESCHED\n");
+                                litmus_reschedule_local();
+                                set_will_schedule();
+                        }
                 }
                 else if (cpu != NO_CPU) {
-                        TRACE_TASK(t, "is preemptable on remote cpu (%d) => FORCE_RESCHED\n", cpu);
                         bt_flag_set(t, BTF_BUDGET_EXHAUSTED);
-                        litmus_reschedule(cpu);
+                        if (!in_schedule) {
+                                TRACE_TASK(t, "is preemptable on remote cpu (%d) => FORCE_RESCHED\n", cpu);
+                                litmus_reschedule(cpu);
+                        }
                 }
                 else if (unlikely(tsk_rt(t)->blocked_lock)) {
                         /* we shouldn't be draining while waiting for litmus lock, but we
@@ -749,6 +770,8 @@ static enum hrtimer_restart cedf_simple_io_on_exhausted(struct task_struct *t)
                         cedf_domain_t *cluster;
                         unsigned long flags;
 
+                        BUG_ON(in_schedule);
+
                         cluster = task_cpu_cluster(t);
 
                         // 1) refresh budget through job completion
@@ -843,6 +866,7 @@ static enum hrtimer_restart cedf_simple_io_on_exhausted(struct task_struct *t)
                 }
         }
 
+out:
         return restart;
 }
 
@@ -875,10 +899,18 @@ static void cedf_trigger_vunlock(struct task_struct *t)
 }
 #endif
 
-static enum hrtimer_restart cedf_sobliv_on_exhausted(struct task_struct *t)
+static enum hrtimer_restart cedf_sobliv_on_exhausted(struct task_struct *t, int in_schedule)
 {
         enum hrtimer_restart restart = HRTIMER_NORESTART;
 
+        if (in_schedule) {
+                BUG_ON(tsk_rt(t)->scheduled_on != smp_processor_id());
+                if (budget_precisely_tracked(t) && cancel_enforcement_timer(t) == -1) {
+                        TRACE_TASK(t, "raced with timer. deffering to timer.\n");
+                        goto out;
+                }
+        }
+
         /* t may or may not be scheduled */
 
         if (budget_signalled(t) && !bt_flag_is_set(t, BTF_SIG_BUDGET_SENT)) {
@@ -911,21 +943,27 @@ static enum hrtimer_restart cedf_sobliv_on_exhausted(struct task_struct *t)
                 }
                 /* where do we need to call resched? */
                 else if (cpu == smp_processor_id()) {
-                        TRACE_TASK(t, "is preemptable => FORCE_RESCHED\n");
                         bt_flag_set(t, BTF_BUDGET_EXHAUSTED);
-                        litmus_reschedule_local();
-                        set_will_schedule();
+                        if (!in_schedule) {
+                                TRACE_TASK(t, "is preemptable => FORCE_RESCHED\n");
+                                litmus_reschedule_local();
+                                set_will_schedule();
+                        }
                 }
                 else if (cpu != NO_CPU) {
-                        TRACE_TASK(t, "is preemptable on remote cpu (%d) => FORCE_RESCHED\n", cpu);
                         bt_flag_set(t, BTF_BUDGET_EXHAUSTED);
-                        litmus_reschedule(cpu);
+                        if (!in_schedule) {
+                                litmus_reschedule(cpu);
+                                TRACE_TASK(t, "is preemptable on remote cpu (%d) => FORCE_RESCHED\n", cpu);
+                        }
                 }
                 else {
                         lt_t remaining;
                         cedf_domain_t *cluster;
                         unsigned long flags;
 
+                        BUG_ON(in_schedule);
+
                         cluster = task_cpu_cluster(t);
 
                         // 1) refresh budget through job completion
@@ -1025,6 +1063,7 @@ static enum hrtimer_restart cedf_sobliv_on_exhausted(struct task_struct *t)
                 }
         }
 
+out:
         return restart;
 }
 
@@ -1041,7 +1080,7 @@ static void cedf_tick(struct task_struct* t)
                 tsk_rt(t)->budget.ops && budget_quantum_tracked(t) &&
                 budget_exhausted(t)) {
                 TRACE_TASK(t, "budget exhausted\n");
-                budget_state_machine(t,on_exhausted);
+                budget_state_machine2(t,on_exhausted,!IN_SCHEDULE);
         }
 }
 
@@ -1388,17 +1427,28 @@ static struct task_struct* cedf_schedule(struct task_struct * prev)
         }
 #endif
 
+        /* Detect and handle budget exhaustion if it hasn't already been done.
+         * Do this before acquring any locks. */
+        if (prev && is_realtime(prev)   &&
+                budget_exhausted(prev)          &&
+                !is_completed(prev)                     && /* don't bother with jobs on their way out */
+                ((budget_enforced(prev)  && !bt_flag_is_set(prev, BTF_BUDGET_EXHAUSTED)) ||
+                 (budget_signalled(prev) && !bt_flag_is_set(prev, BTF_SIG_BUDGET_SENT))) ) {
+                TRACE_TASK(prev, "handling exhaustion in schedule() at %llu\n", litmus_clock());
+                budget_state_machine2(prev,on_exhausted,IN_SCHEDULE);
+        }
+
 #ifdef CONFIG_LITMUS_NESTED_LOCKING
         /* prevent updates to inheritance relations while we work with 'prev' */
         /* recheck inheritance if the task holds locks, is running, and will
          * have its deadline pushed out by job_completion() */
         recheck_inheritance =
-                        (prev != NULL) &&
-                        is_realtime(prev) &&
-                        holds_locks(prev) &&
-                        !is_np(prev) &&
+                        prev                            &&
+                        is_realtime(prev)       &&
+                        holds_locks(prev)       &&
+                        !is_np(prev)            &&
                         !is_completed(prev) &&
-                        is_running(prev) &&
+                        is_running(prev)        &&
                         budget_enforced(prev) &&
                         bt_flag_is_set(prev, BTF_BUDGET_EXHAUSTED);
         if (recheck_inheritance) {
@@ -1726,6 +1776,16 @@ static void cedf_task_exit(struct task_struct * t)
         cedf_change_prio_pai_tasklet(t, NULL);
 #endif
 
+        /*
+BUG: t is forced to exit by another task.
+meanwhile, the scheduler selects to migrate to be scheduled
+
+-- this triggers BAD BAD BAD
+
+if (current != t) and t is linked (but not scheduled?), do something.
+
+*/
+
         /* unlink if necessary */
         raw_readyq_lock_irqsave(&cluster->cluster_lock, flags);
 
@@ -1936,7 +1996,7 @@ static int __increase_priority_inheritance(struct task_struct* t,
 
                 /* clear out old inheritance relation */
                 if (old_prio_inh) {
-                        budget_state_machine2(t,old_prio_inh,on_disinherit);
+                        budget_state_machine_chgprio(t,old_prio_inh,on_disinherit);
                         clear_inh_task_linkback(t, old_prio_inh);
                 }
 
@@ -1946,7 +2006,7 @@ static int __increase_priority_inheritance(struct task_struct* t,
 
                 /* update inheritance relation */
                 if (prio_inh)
-                        budget_state_machine2(t,prio_inh,on_inherit);
+                        budget_state_machine_chgprio(t,prio_inh,on_inherit);
 
                 linked_on  = tsk_rt(t)->linked_on;
 
@@ -2133,7 +2193,7 @@ static int __decrease_priority_inheritance(struct task_struct* t,
 
                 /* clear out old inheritance relation */
                 if (old_prio_inh) {
-                        budget_state_machine2(t,old_prio_inh,on_disinherit);
+                        budget_state_machine_chgprio(t,old_prio_inh,on_disinherit);
                         clear_inh_task_linkback(t, old_prio_inh);
                 }
 
@@ -2149,7 +2209,7 @@ static int __decrease_priority_inheritance(struct task_struct* t,
                 tsk_rt(t)->inh_task = prio_inh;
 
                 if (prio_inh)
-                        budget_state_machine2(t,prio_inh,on_inherit);
+                        budget_state_machine_chgprio(t,prio_inh,on_inherit);
 
                 if(tsk_rt(t)->scheduled_on != NO_CPU) {
                         TRACE_TASK(t, "is scheduled.\n");
diff --git a/litmus/sched_gsn_edf.c b/litmus/sched_gsn_edf.c
index 0f96fb7ddb25..4aeb7a0db3bd 100644
--- a/litmus/sched_gsn_edf.c
+++ b/litmus/sched_gsn_edf.c
@@ -362,11 +362,11 @@ static void check_for_preemptions(void)
                                                 &per_cpu(gsnedf_cpu_entries, task_cpu(task)));
                         if (affinity)
                                 last = affinity;
-                        else if (requeue_preempted_job(last->linked))
+                        else if (should_requeue_preempted_job(last->linked))
                                 requeue(last->linked);
                 }
 #else
-                if (requeue_preempted_job(last->linked))
+                if (should_requeue_preempted_job(last->linked))
                         requeue(last->linked);
 #endif
 
@@ -420,7 +420,7 @@ static noinline void job_completion(struct task_struct *t, int forced)
                 gsnedf_job_arrival(t);
 }
 
-static enum hrtimer_restart gsnedf_simple_on_exhausted(struct task_struct *t)
+static enum hrtimer_restart gsnedf_simple_on_exhausted(struct task_struct *t, int in_schedule)
 {
         /* Assumption: t is scheduled on the CPU executing this callback */
 
@@ -461,7 +461,7 @@ static void gsnedf_tick(struct task_struct* t)
                 tsk_rt(t)->budget.ops && budget_quantum_tracked(t) &&
                 budget_exhausted(t)) {
                 TRACE_TASK(t, "budget exhausted\n");
-                budget_state_machine(t,on_exhausted);
+                budget_state_machine2(t,on_exhausted,!IN_SCHEDULE);
         }
 }
 
@@ -1171,14 +1171,14 @@ static int __increase_priority_inheritance(struct task_struct* t,
         if(__edf_higher_prio(prio_inh, BASE, t, EFFECTIVE)) {
 #endif
                 if (tsk_rt(t)->inh_task)
-                        budget_state_machine2(t,tsk_rt(t)->inh_task,on_disinherit);
+                        budget_state_machine_chgprio(t,tsk_rt(t)->inh_task,on_disinherit);
 
                 TRACE_TASK(t, "inherits priority from %s/%d\n",
                                    prio_inh->comm, prio_inh->pid);
                 tsk_rt(t)->inh_task = prio_inh;
 
                 if (prio_inh)
-                        budget_state_machine2(t,prio_inh,on_inherit);
+                        budget_state_machine_chgprio(t,prio_inh,on_inherit);
 
                 linked_on  = tsk_rt(t)->linked_on;
 
@@ -1309,7 +1309,7 @@ static int __decrease_priority_inheritance(struct task_struct* t,
 #endif
 
                 if (tsk_rt(t)->inh_task)
-                        budget_state_machine2(t,tsk_rt(t)->inh_task,on_disinherit);
+                        budget_state_machine_chgprio(t,tsk_rt(t)->inh_task,on_disinherit);
 
                 /* A job only stops inheriting a priority when it releases a
                  * resource. Thus we can make the following assumption.*/
@@ -1322,7 +1322,7 @@ static int __decrease_priority_inheritance(struct task_struct* t,
                 tsk_rt(t)->inh_task = prio_inh;
 
                 if (prio_inh)
-                        budget_state_machine2(t,prio_inh,on_inherit);
+                        budget_state_machine_chgprio(t,prio_inh,on_inherit);
 
                 if(tsk_rt(t)->scheduled_on != NO_CPU) {
                         TRACE_TASK(t, "is scheduled.\n");
diff --git a/litmus/sched_litmus.c b/litmus/sched_litmus.c
index eadd4fb8e5a4..3a594cd51a44 100644
--- a/litmus/sched_litmus.c
+++ b/litmus/sched_litmus.c
@@ -126,7 +126,7 @@ litmus_schedule(struct rq *rq, struct task_struct *prev)
                         }
                 }
 
-                set_task_cpu(next, smp_processor_id());
+//              set_task_cpu(next, smp_processor_id());
 
                 /* DEBUG: now that we have the lock we need to make sure a
                  *  couple of things still hold:
@@ -135,7 +135,7 @@ litmus_schedule(struct rq *rq, struct task_struct *prev)
                  * If either is violated, then the active plugin is
                  * doing something wrong.
                  */
-                if (!is_realtime(next) || !is_running(next)) {
+                if (unlikely(!is_realtime(next) || !is_running(next))) {
                         /* BAD BAD BAD */
                         TRACE_TASK(next,"BAD: migration invariant FAILED: "
                                    "rt=%d running=%d\n",
@@ -144,6 +144,10 @@ litmus_schedule(struct rq *rq, struct task_struct *prev)
                         /* drop the task */
                         next = NULL;
                 }
+                else {
+                        set_task_cpu(next, smp_processor_id());
+                }
+
                 /* release the other CPU's runqueue, but keep ours */
                 raw_spin_unlock(&other_rq->lock);
         }
@@ -230,6 +234,52 @@ static void pre_schedule_litmus(struct rq *rq, struct task_struct *prev)
                 tsk_rt(prev)->present = 0;
 }
 
+static void litmus_handle_budget_exhaustion(struct task_struct *t)
+{
+        /* We're unlikely to pick a task that has an exhausted budget, so this
+         * provides a failsafe. */
+
+        /* BUG: Virtual unlock of OMLP-family locking protocols is not triggered.
+         *
+         * TODO-FIX: Add a new virtual-unlock call to budget state machine and do
+         * the virtual unlock in plugin::schedule(), instead of in budget
+         * timer handler. This bug should only be raised EXTREMELY infrequently.
+         */
+
+        int handle_exhausion = 1;
+
+        BUG_ON(current != t);
+
+        if (is_np(t) && bt_flag_is_set(t, BTF_BUDGET_EXHAUSTED)) {
+                /* ignore. will handle exhausion in the future. */
+                TRACE_TASK(t, "task is np and already flagged as exhausted. allow scheduling.\n");
+                return;
+        }
+
+        if (unlikely(bt_flag_is_set(t, BTF_WAITING_FOR_RELEASE))) {
+                TRACE_TASK(t, "waiting for release. skipping exhaustion.\n");
+                return;
+        }
+
+        if (budget_precisely_tracked(t)) {
+                if (cancel_enforcement_timer(t) < 0) {
+                        TRACE_TASK(t, "schedule() raced with timer. deffering to timer.\n");
+                        handle_exhausion = 0;
+                }
+        }
+
+        if (handle_exhausion) {
+                if (likely(!is_np(t))) {
+                        TRACE_TASK(t, "picked task without budget => FORCE_RESCHED.\n");
+                        litmus_reschedule_local();
+                }
+                else if (is_user_np(t)) {
+                        TRACE_TASK(t, "is non-preemptable, preemption delayed.\n");
+                        request_exit_np(t);
+                }
+        }
+}
+
 /* pick_next_task_litmus() - litmus_schedule() function
  *
  * return the next task to be scheduled
diff --git a/litmus/sched_pfp.c b/litmus/sched_pfp.c
index 891469054ab4..d9d293ec98f9 100644
--- a/litmus/sched_pfp.c
+++ b/litmus/sched_pfp.c
@@ -132,7 +132,7 @@ static void job_completion(struct task_struct* t, int forced)
                 sched_trace_task_release(t);
 }
 
-static enum hrtimer_restart pfp_simple_on_exhausted(struct task_struct *t)
+static enum hrtimer_restart pfp_simple_on_exhausted(struct task_struct *t, int in_schedule)
 {
         /* Assumption: t is scheduled on the CPU executing this callback */
 
@@ -175,7 +175,7 @@ static void pfp_tick(struct task_struct *t)
                 tsk_rt(t)->budget.ops && budget_quantum_tracked(t) &&
                 budget_exhausted(t)) {
                 TRACE_TASK(t, "budget exhausted\n");
-                budget_state_machine(t,on_exhausted);
+                budget_state_machine2(t,on_exhausted,IN_SCHEDULE);
         }
 }
 
@@ -485,12 +485,12 @@ static void fp_set_prio_inh(pfp_domain_t* pfp, struct task_struct* t,
                 fp_dequeue(pfp, t);
 
         if (t->rt_param.inh_task)
-                budget_state_machine2(t,t->rt_param.inh_task,on_disinherit);
+                budget_state_machine_chgprio(t,t->rt_param.inh_task,on_disinherit);
 
         t->rt_param.inh_task = prio_inh;
 
         if (prio_inh)
-                budget_state_machine2(t,prio_inh,on_inherit);
+                budget_state_machine_chgprio(t,prio_inh,on_inherit);
 
         if (requeue)
                 /* add again to the right queue */
diff --git a/litmus/sched_plugin.c b/litmus/sched_plugin.c
index cfc572584fca..35a98eb806a2 100644
--- a/litmus/sched_plugin.c
+++ b/litmus/sched_plugin.c
@@ -122,14 +122,14 @@ static long litmus_dummy_complete_job(void)
 
 static long litmus_dummy_activate_plugin(void)
 {
-#ifdef CONFIG_LITMUS_NVIDIA
-        shutdown_nvidia_info();
-#endif
         return 0;
 }
 
 static long litmus_dummy_deactivate_plugin(void)
 {
+//#ifdef CONFIG_LITMUS_NVIDIA
+//      shutdown_nvidia_info();
+//#endif
         return 0;
 }
 
diff --git a/litmus/sched_psn_edf.c b/litmus/sched_psn_edf.c
index b8246acacaa9..a51b444b30e1 100644
--- a/litmus/sched_psn_edf.c
+++ b/litmus/sched_psn_edf.c
@@ -164,7 +164,7 @@ static void job_completion(struct task_struct* t, int forced)
         prepare_for_next_period(t);
 }
 
-static enum hrtimer_restart psnedf_simple_on_exhausted(struct task_struct *t)
+static enum hrtimer_restart psnedf_simple_on_exhausted(struct task_struct *t, int in_schedule)
 {
         /* Assumption: t is scheduled on the CPU executing this callback */
 
@@ -205,7 +205,7 @@ static void psnedf_tick(struct task_struct *t)
                 tsk_rt(t)->budget.ops && budget_quantum_tracked(t) &&
                 budget_exhausted(t)) {
                 TRACE_TASK(t, "budget exhausted\n");
-                budget_state_machine(t,on_exhausted);
+                budget_state_machine2(t,on_exhausted,!IN_SCHEDULE);
         }
 }
 
diff --git a/litmus/sync.c b/litmus/sync.c
index dfd9e40ab28d..89a28f8a1cc1 100644
--- a/litmus/sync.c
+++ b/litmus/sync.c
@@ -60,11 +60,18 @@ static long do_wait_for_ts_release(struct timespec *wake)
 #endif
                 bt_flag_set(t, BTF_WAITING_FOR_RELEASE);
                 mb();
+                budget_state_machine(t,on_exit); // do this here and not in schedule()?
         }
 
+        TRACE_TASK(t, "waiting for ts release.\n");
+        if (is_rt)
+                BUG_ON(!bt_flag_is_set(t, BTF_WAITING_FOR_RELEASE));
+
         /* We are enqueued, now we wait for someone to wake us up. */
         ret = wait_for_completion_interruptible(&wait.completion);
 
+        TRACE_TASK(t, "released by ts release!\n");
+
         if (is_rt) {
                 bt_flag_clear(t, BTF_WAITING_FOR_RELEASE);
 #if defined(CONFIG_REALTIME_AUX_TASKS) || defined(CONFIG_LITMUS_NVIDIA)
@@ -92,7 +99,8 @@ static long do_wait_for_ts_release(struct timespec *wake)
                         /* completion succeeded, set up release. subtract off
                          * period because schedule()->job_completion() will
                          * advances us to the correct time */
-                        litmus->release_at(t, phasedRelease - t->rt_param.task_params.period);
+                        //litmus->release_at(t, phasedRelease - t->rt_param.task_params.period);
+                        setup_release(t, phasedRelease - t->rt_param.task_params.period);  // breaks pfair
                         schedule();
                 }
                 else {