Implement proper remote preemption support

To date, Litmus has just hooked into the smp_send_reschedule() IPI
handler and marked tasks as having to reschedule to implement remote
preemptions. This was never particularly clean, but so far we got away
with it. However, changes in the underlying Linux, and peculartities
of the ARM code (interrupts enabled before context switch) break this
naive approach. This patch introduces new state-machine based remote
preemption support. By examining the local state before calling
set_tsk_need_resched(), we avoid confusing the underlying Linux
scheduler. Further, this patch avoids sending unncessary IPIs.

12 files changed, 341 insertions, 51 deletions

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 708523b06920..c9ac4fc837ba 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -96,6 +96,7 @@ struct sched_param {
 #include <asm/processor.h>
 
 #include <litmus/rt_param.h>
+#include <litmus/preempt.h>
 
 struct exec_domain;
 struct futex_pi_state;
@@ -2301,6 +2302,7 @@ static inline int test_tsk_thread_flag(struct task_struct *tsk, int flag)
 static inline void set_tsk_need_resched(struct task_struct *tsk)
 {
         set_tsk_thread_flag(tsk,TIF_NEED_RESCHED);
+        sched_state_will_schedule(tsk);
 }
 
 static inline void clear_tsk_need_resched(struct task_struct *tsk)
diff --git a/include/litmus/preempt.h b/include/litmus/preempt.h
new file mode 100644
index 000000000000..260c6fe17986
--- /dev/null
+++ b/include/litmus/preempt.h
@@ -0,0 +1,164 @@
+#ifndef LITMUS_PREEMPT_H
+#define LITMUS_PREEMPT_H
+
+#include <linux/types.h>
+#include <linux/cache.h>
+#include <linux/percpu.h>
+#include <asm/atomic.h>
+
+#include <litmus/debug_trace.h>
+
+extern DEFINE_PER_CPU_SHARED_ALIGNED(atomic_t, resched_state);
+
+#ifdef CONFIG_DEBUG_KERNEL
+const char* sched_state_name(int s);
+#define TRACE_STATE(fmt, args...) TRACE("SCHED_STATE " fmt, args)
+#else
+#define TRACE_STATE(fmt, args...) /* ignore */
+#endif
+
+#define VERIFY_SCHED_STATE(x)                                           \
+        do { int __s = get_sched_state();                               \
+                if ((__s & (x)) == 0)                                   \
+                        TRACE_STATE("INVALID s=0x%x (%s) not "          \
+                                    "in 0x%x (%s) [%s]\n",              \
+                                    __s, sched_state_name(__s),         \
+                                    (x), #x, __FUNCTION__);             \
+        } while (0);
+
+#define TRACE_SCHED_STATE_CHANGE(x, y, cpu)                             \
+        TRACE_STATE("[P%d] 0x%x (%s) -> 0x%x (%s)\n",                   \
+                    cpu,  (x), sched_state_name(x),                     \
+                    (y), sched_state_name(y))
+
+
+typedef enum scheduling_state {
+        TASK_SCHEDULED    = (1 << 0),  /* The currently scheduled task is the one that
+                                        * should be scheduled, and the processor does not
+                                        * plan to invoke schedule(). */
+        SHOULD_SCHEDULE   = (1 << 1),  /* A remote processor has determined that the
+                                        * processor should reschedule, but this has not
+                                        * been communicated yet (IPI still pending). */
+        WILL_SCHEDULE     = (1 << 2),  /* The processor has noticed that it has to
+                                        * reschedule and will do so shortly. */
+        TASK_PICKED       = (1 << 3),  /* The processor is currently executing schedule(),
+                                        * has selected a new task to schedule, but has not
+                                        * yet performed the actual context switch. */
+        PICKED_WRONG_TASK = (1 << 4),  /* The processor has not yet performed the context
+                                        * switch, but a remote processor has already
+                                        * determined that a higher-priority task became
+                                        * eligible after the task was picked. */
+} sched_state_t;
+
+static inline sched_state_t get_sched_state_on(int cpu)
+{
+        return atomic_read(&per_cpu(resched_state, cpu));
+}
+
+static inline sched_state_t get_sched_state(void)
+{
+        return atomic_read(&__get_cpu_var(resched_state));
+}
+
+static inline int is_in_sched_state(int possible_states)
+{
+        return get_sched_state() & possible_states;
+}
+
+static inline int cpu_is_in_sched_state(int cpu, int possible_states)
+{
+        return get_sched_state_on(cpu) & possible_states;
+}
+
+static inline void set_sched_state(sched_state_t s)
+{
+        TRACE_SCHED_STATE_CHANGE(get_sched_state(), s, smp_processor_id());
+        atomic_set(&__get_cpu_var(resched_state), s);
+}
+
+static inline int sched_state_transition(sched_state_t from, sched_state_t to)
+{
+        sched_state_t old_state;
+
+        old_state = atomic_cmpxchg(&__get_cpu_var(resched_state), from, to);
+        if (old_state == from) {
+                TRACE_SCHED_STATE_CHANGE(from, to, smp_processor_id());
+                return 1;
+        } else
+                return 0;
+}
+
+static inline int sched_state_transition_on(int cpu,
+                                            sched_state_t from,
+                                            sched_state_t to)
+{
+        sched_state_t old_state;
+
+        old_state = atomic_cmpxchg(&per_cpu(resched_state, cpu), from, to);
+        if (old_state == from) {
+                TRACE_SCHED_STATE_CHANGE(from, to, cpu);
+                return 1;
+        } else
+                return 0;
+}
+
+/* Plugins must call this function after they have decided which job to
+ * schedule next.  IMPORTANT: this function must be called while still holding
+ * the lock that is used to serialize scheduling decisions.
+ *
+ * (Ideally, we would like to use runqueue locks for this purpose, but that
+ * would lead to deadlocks with the migration code.)
+ */
+static inline void sched_state_task_picked(void)
+{
+        VERIFY_SCHED_STATE(WILL_SCHEDULE);
+
+        /* WILL_SCHEDULE has only a local tansition => simple store is ok */
+        set_sched_state(TASK_PICKED);
+}
+
+static inline void sched_state_entered_schedule(void)
+{
+        /* Update state for the case that we entered schedule() not due to
+         * set_tsk_need_resched() */
+        set_sched_state(WILL_SCHEDULE);
+}
+
+/* Called by schedule() to check if the scheduling decision is still valid
+ * after a context switch. Returns 1 if the CPU needs to reschdule. */
+static inline int sched_state_validate_switch(void)
+{
+        int left_state_ok = 0;
+
+        VERIFY_SCHED_STATE(PICKED_WRONG_TASK | TASK_PICKED);
+
+        if (is_in_sched_state(TASK_PICKED)) {
+                /* Might be good; let's try to transition out of this
+                 * state. This must be done atomically since remote processors
+                 * may try to change the state, too. */
+                left_state_ok = sched_state_transition(TASK_PICKED, TASK_SCHEDULED);
+        }
+
+        if (!left_state_ok) {
+                /* We raced with a higher-priority task arrival => not
+                 * valid. The CPU needs to reschedule. */
+                set_sched_state(WILL_SCHEDULE);
+                return 1;
+        } else
+                return 0;
+}
+
+/* State transition events. See litmus/preempt.c for details. */
+void sched_state_will_schedule(struct task_struct* tsk);
+void sched_state_ipi(void);
+/* Cause a CPU (remote or local) to reschedule. */
+void litmus_reschedule(int cpu);
+void litmus_reschedule_local(void);
+
+#ifdef CONFIG_DEBUG_KERNEL
+void sched_state_plugin_check(void);
+#else
+#define sched_state_plugin_check() /* no check */
+#endif
+
+#endif
diff --git a/kernel/sched.c b/kernel/sched.c
index a9dd6f96c731..60fbae0c747c 100644
--- a/kernel/sched.c
+++ b/kernel/sched.c
@@ -3794,6 +3794,7 @@ asmlinkage void __sched schedule(void)
 
 need_resched:
         preempt_disable();
+        sched_state_entered_schedule();
         cpu = smp_processor_id();
         rq = cpu_rq(cpu);
         rcu_note_context_switch(cpu);
@@ -3872,7 +3873,7 @@ need_resched_nonpreemptible:
 
         post_schedule(rq);
 
-        if (unlikely(reacquire_kernel_lock(prev)))
+        if (sched_state_validate_switch() || unlikely(reacquire_kernel_lock(prev)))
                 goto need_resched_nonpreemptible;
 
         preempt_enable_no_resched();
diff --git a/litmus/Makefile b/litmus/Makefile
index 7bd1abdcb84a..b7366b530749 100644
--- a/litmus/Makefile
+++ b/litmus/Makefile
@@ -3,6 +3,7 @@
 #
 
 obj-y     = sched_plugin.o litmus.o \
+            preempt.o \
             litmus_proc.o \
             budget.o \
             jobs.o \
diff --git a/litmus/budget.c b/litmus/budget.c
index f6f5ca81c9d6..310e9a3d4172 100644
--- a/litmus/budget.c
+++ b/litmus/budget.c
@@ -3,6 +3,7 @@
 #include <linux/hrtimer.h>
 
 #include <litmus/litmus.h>
+#include <litmus/preempt.h>
 
 struct enforcement_timer {
         /* The enforcement timer is used to accurately police
@@ -24,7 +25,7 @@ static enum hrtimer_restart on_enforcement_timeout(struct hrtimer *timer)
         TRACE("enforcement timer fired.\n");
         et->armed = 0;
         /* activate scheduler */
-        set_tsk_need_resched(current);
+        litmus_reschedule_local();
         local_irq_restore(flags);
 
         return  HRTIMER_NORESTART;
diff --git a/litmus/preempt.c b/litmus/preempt.c
new file mode 100644
index 000000000000..ebe2e3461895
--- /dev/null
+++ b/litmus/preempt.c
@@ -0,0 +1,131 @@
+#include <linux/sched.h>
+
+#include <litmus/litmus.h>
+#include <litmus/preempt.h>
+
+/* The rescheduling state of each processor.
+ */
+DEFINE_PER_CPU_SHARED_ALIGNED(atomic_t, resched_state);
+
+void sched_state_will_schedule(struct task_struct* tsk)
+{
+        /* Litmus hack: we only care about processor-local invocations of
+         * set_tsk_need_resched(). We can't reliably set the flag remotely
+         * since it might race with other updates to the scheduling state.  We
+         * can't rely on the runqueue lock protecting updates to the sched
+         * state since processors do not acquire the runqueue locks for all
+         * updates to the sched state (to avoid acquiring two runqueue locks at
+         * the same time). Further, if tsk is residing on a remote processor,
+         * then that processor doesn't actually know yet that it is going to
+         * reschedule; it still must receive an IPI (unless a local invocation
+         * races).
+         */
+        if (likely(task_cpu(tsk) == smp_processor_id())) {
+                VERIFY_SCHED_STATE(TASK_SCHEDULED | SHOULD_SCHEDULE | TASK_PICKED | WILL_SCHEDULE);
+                if (is_in_sched_state(TASK_PICKED | PICKED_WRONG_TASK))
+                        set_sched_state(PICKED_WRONG_TASK);
+                else
+                        set_sched_state(WILL_SCHEDULE);
+        } else
+                /* Litmus tasks should never be subject to a remote
+                 * set_tsk_need_resched(). */
+                BUG_ON(is_realtime(tsk));
+        TRACE_TASK(tsk, "set_tsk_need_resched() ret:%p\n",
+                   __builtin_return_address(0));
+}
+
+/* Called by the IPI handler after another CPU called smp_send_resched(). */
+void sched_state_ipi(void)
+{
+        /* If the IPI was slow, we might be in any state right now. The IPI is
+         * only meaningful if we are in SHOULD_SCHEDULE. */
+        if (is_in_sched_state(SHOULD_SCHEDULE)) {
+                /* Cause scheduler to be invoked.
+                 * This will cause a transition to WILL_SCHEDULE. */
+                set_tsk_need_resched(current);
+                TRACE_STATE("IPI -> set_tsk_need_resched(%s/%d)\n",
+                            current->comm, current->pid);
+        } else {
+                /* ignore */
+                TRACE_STATE("ignoring IPI in state %x (%s)\n",
+                            get_sched_state(),
+                            sched_state_name(get_sched_state()));
+        }
+}
+
+/* Called by plugins to cause a CPU to reschedule. IMPORTANT: the caller must
+ * hold the lock that is used to serialize scheduling decisions. */
+void litmus_reschedule(int cpu)
+{
+        int picked_transition_ok = 0;
+        int scheduled_transition_ok = 0;
+
+        /* The (remote) CPU could be in any state. */
+
+        /* The critical states are TASK_PICKED and TASK_SCHEDULED, as the CPU
+         * is not aware of the need to reschedule at this point. */
+
+        /* is a context switch in progress? */
+        if (cpu_is_in_sched_state(cpu, TASK_PICKED))
+                picked_transition_ok = sched_state_transition_on(
+                        cpu, TASK_PICKED, PICKED_WRONG_TASK);
+
+        if (!picked_transition_ok &&
+            cpu_is_in_sched_state(cpu, TASK_SCHEDULED)) {
+                /* We either raced with the end of the context switch, or the
+                 * CPU was in TASK_SCHEDULED anyway. */
+                scheduled_transition_ok = sched_state_transition_on(
+                        cpu, TASK_SCHEDULED, SHOULD_SCHEDULE);
+        }
+
+        /* If the CPU was in state TASK_SCHEDULED, then we need to cause the
+         * scheduler to be invoked. */
+        if (scheduled_transition_ok) {
+                if (smp_processor_id() == cpu)
+                        set_tsk_need_resched(current);
+                else
+                        smp_send_reschedule(cpu);
+        }
+
+        TRACE_STATE("%s picked-ok:%d sched-ok:%d\n",
+                    __FUNCTION__,
+                    picked_transition_ok,
+                    scheduled_transition_ok);
+}
+
+void litmus_reschedule_local(void)
+{
+        if (is_in_sched_state(TASK_PICKED))
+                set_sched_state(PICKED_WRONG_TASK);
+        else if (is_in_sched_state(TASK_SCHEDULED | SHOULD_SCHEDULE)) {
+                set_sched_state(WILL_SCHEDULE);
+                set_tsk_need_resched(current);
+        }
+}
+
+#ifdef CONFIG_DEBUG_KERNEL
+
+void sched_state_plugin_check(void)
+{
+        if (!is_in_sched_state(TASK_PICKED | PICKED_WRONG_TASK)) {
+                TRACE("!!!! plugin did not call sched_state_task_picked()!"
+                      "Calling sched_state_task_picked() is mandatory---fix this.\n");
+                set_sched_state(TASK_PICKED);
+        }
+}
+
+#define NAME_CHECK(x) case x:  return #x
+const char* sched_state_name(int s)
+{
+        switch (s) {
+                NAME_CHECK(TASK_SCHEDULED);
+                NAME_CHECK(SHOULD_SCHEDULE);
+                NAME_CHECK(WILL_SCHEDULE);
+                NAME_CHECK(TASK_PICKED);
+                NAME_CHECK(PICKED_WRONG_TASK);
+        default:
+                return "UNKNOWN";
+        };
+}
+
+#endif
diff --git a/litmus/sched_cedf.c b/litmus/sched_cedf.c
index a729d97535e9..615560f21d60 100644
--- a/litmus/sched_cedf.c
+++ b/litmus/sched_cedf.c
@@ -34,6 +34,7 @@
 
 #include <litmus/litmus.h>
 #include <litmus/jobs.h>
+#include <litmus/preempt.h>
 #include <litmus/sched_plugin.h>
 #include <litmus/edf_common.h>
 #include <litmus/sched_trace.h>
@@ -341,7 +342,7 @@ static void cedf_tick(struct task_struct* t)
                         /* np tasks will be preempted when they become
                          * preemptable again
                          */
-                        set_tsk_need_resched(t);
+                        litmus_reschedule_local();
                         set_will_schedule();
                         TRACE("cedf_scheduler_tick: "
                               "%d is preemptable "
@@ -466,6 +467,7 @@ static struct task_struct* cedf_schedule(struct task_struct * prev)
                 if (exists)
                         next = prev;
 
+        sched_state_task_picked();
         raw_spin_unlock(&cluster->lock);
 
 #ifdef WANT_ALL_SCHED_EVENTS
diff --git a/litmus/sched_gsn_edf.c b/litmus/sched_gsn_edf.c
index e101768740ad..f0337cfd8631 100644
--- a/litmus/sched_gsn_edf.c
+++ b/litmus/sched_gsn_edf.c
@@ -18,6 +18,8 @@
 #include <litmus/edf_common.h>
 #include <litmus/sched_trace.h>
 
+#include <litmus/preempt.h>
+
 #include <litmus/bheap.h>
 
 #include <linux/module.h>
@@ -95,21 +97,12 @@ typedef struct  {
         int                     cpu;
         struct task_struct*     linked;         /* only RT tasks */
         struct task_struct*     scheduled;      /* only RT tasks */
-        atomic_t                will_schedule;  /* prevent unneeded IPIs */
         struct bheap_node*      hn;
 } cpu_entry_t;
 DEFINE_PER_CPU(cpu_entry_t, gsnedf_cpu_entries);
 
 cpu_entry_t* gsnedf_cpus[NR_CPUS];
 
-#define set_will_schedule() \
-        (atomic_set(&__get_cpu_var(gsnedf_cpu_entries).will_schedule, 1))
-#define clear_will_schedule() \
-        (atomic_set(&__get_cpu_var(gsnedf_cpu_entries).will_schedule, 0))
-#define test_will_schedule(cpu) \
-        (atomic_read(&per_cpu(gsnedf_cpu_entries, cpu).will_schedule))
-
-
 /* the cpus queue themselves according to priority in here */
 static struct bheap_node gsnedf_heap_node[NR_CPUS];
 static struct bheap      gsnedf_cpu_heap;
@@ -341,8 +334,7 @@ static void gsnedf_tick(struct task_struct* t)
                         /* np tasks will be preempted when they become
                          * preemptable again
                          */
-                        set_tsk_need_resched(t);
-                        set_will_schedule();
+                        litmus_reschedule_local();
                         TRACE("gsnedf_scheduler_tick: "
                               "%d is preemptable "
                               " => FORCE_RESCHED\n", t->pid);
@@ -391,7 +383,6 @@ static struct task_struct* gsnedf_schedule(struct task_struct * prev)
 #endif
 
         raw_spin_lock(&gsnedf_lock);
-        clear_will_schedule();
 
         /* sanity checking */
         BUG_ON(entry->scheduled && entry->scheduled != prev);
@@ -473,6 +464,8 @@ static struct task_struct* gsnedf_schedule(struct task_struct * prev)
                 if (exists)
                         next = prev;
 
+        sched_state_task_picked();
+
         raw_spin_unlock(&gsnedf_lock);
 
 #ifdef WANT_ALL_SCHED_EVENTS
@@ -780,7 +773,6 @@ static long gsnedf_activate_plugin(void)
         for_each_online_cpu(cpu) {
                 entry = &per_cpu(gsnedf_cpu_entries, cpu);
                 bheap_node_init(&entry->hn, entry);
-                atomic_set(&entry->will_schedule, 0);
                 entry->linked    = NULL;
                 entry->scheduled = NULL;
 #ifdef CONFIG_RELEASE_MASTER
@@ -829,7 +821,6 @@ static int __init init_gsn_edf(void)
         for (cpu = 0; cpu < NR_CPUS; cpu++)  {
                 entry = &per_cpu(gsnedf_cpu_entries, cpu);
                 gsnedf_cpus[cpu] = entry;
-                atomic_set(&entry->will_schedule, 0);
                 entry->cpu       = cpu;
                 entry->hn        = &gsnedf_heap_node[cpu];
                 bheap_node_init(&entry->hn, entry);
diff --git a/litmus/sched_litmus.c b/litmus/sched_litmus.c
index 65873152e68f..e6952896dc4b 100644
--- a/litmus/sched_litmus.c
+++ b/litmus/sched_litmus.c
@@ -3,6 +3,7 @@
 #include <litmus/litmus.h>
 #include <litmus/budget.h>
 #include <litmus/sched_plugin.h>
+#include <litmus/preempt.h>
 
 static void update_time_litmus(struct rq *rq, struct task_struct *p)
 {
@@ -52,6 +53,8 @@ litmus_schedule(struct rq *rq, struct task_struct *prev)
         /* let the plugin schedule */
         next = litmus->schedule(prev);
 
+        sched_state_plugin_check();
+
         /* check if a global plugin pulled a task from a different RQ */
         if (next && task_rq(next) != rq) {
                 /* we need to migrate the task */
@@ -198,7 +201,7 @@ static void yield_task_litmus(struct rq *rq)
          * then determine if a preemption is still required.
          */
         clear_exit_np(current);
-        set_tsk_need_resched(current);
+        litmus_reschedule_local();
 }
 
 /* Plugins are responsible for this.
diff --git a/litmus/sched_pfair.c b/litmus/sched_pfair.c
index ea77d3295290..c7d5cf7aa2b3 100644
--- a/litmus/sched_pfair.c
+++ b/litmus/sched_pfair.c
@@ -16,6 +16,7 @@
 
 #include <litmus/litmus.h>
 #include <litmus/jobs.h>
+#include <litmus/preempt.h>
 #include <litmus/rt_domain.h>
 #include <litmus/sched_plugin.h>
 #include <litmus/sched_trace.h>
@@ -241,11 +242,7 @@ static void check_preempt(struct task_struct* t)
                 PTRACE_TASK(t, "linked_on:%d, scheduled_on:%d\n",
                            tsk_rt(t)->linked_on, tsk_rt(t)->scheduled_on);
                 /* preempt */
-                if (cpu == smp_processor_id())
-                        set_tsk_need_resched(current);
-                else {
-                        smp_send_reschedule(cpu);
-                }
+                litmus_reschedule(cpu);
         }
 }
 
@@ -545,7 +542,7 @@ static void pfair_tick(struct task_struct* t)
 
         if (state->local != current
             && (is_realtime(current) || is_present(state->local)))
-                set_tsk_need_resched(current);
+                litmus_reschedule_local();
 }
 
 static int safe_to_schedule(struct task_struct* t, int cpu)
@@ -577,7 +574,7 @@ static struct task_struct* pfair_schedule(struct task_struct * prev)
                 if (next)
                         tsk_rt(next)->scheduled_on = state->cpu;
         }
-
+        sched_state_task_picked();
         raw_spin_unlock(&pfair_lock);
 
         if (next)
diff --git a/litmus/sched_plugin.c b/litmus/sched_plugin.c
index ec04454a0cf9..d912a6494d20 100644
--- a/litmus/sched_plugin.c
+++ b/litmus/sched_plugin.c
@@ -10,7 +10,7 @@
 
 #include <litmus/litmus.h>
 #include <litmus/sched_plugin.h>
-
+#include <litmus/preempt.h>
 #include <litmus/jobs.h>
 
 /*
@@ -19,36 +19,30 @@
  * non-preemptive section aware and does not invoke the scheduler / send
  * IPIs if the to-be-preempted task is actually non-preemptive.
  */
-void preempt_if_preemptable(struct task_struct* t, int on_cpu)
+void preempt_if_preemptable(struct task_struct* t, int cpu)
 {
         /* t is the real-time task executing on CPU on_cpu If t is NULL, then
          * on_cpu is currently scheduling background work.
          */
 
-        int send_ipi;
+        int reschedule = 0;
 
-        if (smp_processor_id() == on_cpu) {
-                /* local CPU case */
-                if (t) {
+        if (!t)
+                /* move non-real-time task out of the way */
+                reschedule = 1;
+        else {
+                if (smp_processor_id() == cpu) {
+                        /* local CPU case */
                         /* check if we need to poke userspace */
                         if (is_user_np(t))
                                 /* yes, poke it */
                                 request_exit_np(t);
-                        else
-                                /* no, see if we are allowed to preempt the
+                        else if (!is_kernel_np(t))
+                                /* only if we are allowed to preempt the
                                  * currently-executing task */
-                                if (!is_kernel_np(t))
-                                        set_tsk_need_resched(t);
-                } else
-                        /* move non-real-time task out of the way */
-                        set_tsk_need_resched(current);
-        } else {
-                /* remote CPU case */
-                if (!t)
-                        /* currently schedules non-real-time work */
-                        send_ipi = 1;
-                else {
-                        /* currently schedules real-time work */
+                                reschedule = 1;
+                } else {
+                        /* remote CPU case */
                         if (is_user_np(t)) {
                                 /* need to notify user space of delayed
                                  * preemption */
@@ -60,14 +54,14 @@ void preempt_if_preemptable(struct task_struct* t, int on_cpu)
                                 mb();
                         }
                         /* Only send an ipi if remote task might have raced our
-                         * request, i.e., send an IPI to make sure if it exited
-                         * its critical section.
+                         * request, i.e., send an IPI to make sure in case it
+                         * exited its critical section.
                          */
-                        send_ipi = !is_np(t) && !is_kernel_np(t);
+                        reschedule = !is_np(t) && !is_kernel_np(t);
                 }
-                if (likely(send_ipi))
-                        smp_send_reschedule(on_cpu);
         }
+        if (likely(reschedule))
+                litmus_reschedule(cpu);
 }
 
 
@@ -81,6 +75,7 @@ static void litmus_dummy_finish_switch(struct task_struct * prev)
 
 static struct task_struct* litmus_dummy_schedule(struct task_struct * prev)
 {
+        sched_state_task_picked();
         return NULL;
 }
 
diff --git a/litmus/sched_psn_edf.c b/litmus/sched_psn_edf.c
index aa567f2b91b9..b89823d5c026 100644
--- a/litmus/sched_psn_edf.c
+++ b/litmus/sched_psn_edf.c
@@ -16,6 +16,7 @@
 
 #include <litmus/litmus.h>
 #include <litmus/jobs.h>
+#include <litmus/preempt.h>
 #include <litmus/sched_plugin.h>
 #include <litmus/edf_common.h>
 #include <litmus/sched_trace.h>
@@ -108,7 +109,7 @@ static void psnedf_tick(struct task_struct *t)
 
         if (is_realtime(t) && budget_enforced(t) && budget_exhausted(t)) {
                 if (!is_np(t)) {
-                        set_tsk_need_resched(t);
+                        litmus_reschedule_local();
                         TRACE("psnedf_scheduler_tick: "
                               "%d is preemptable "
                               " => FORCE_RESCHED\n", t->pid);
@@ -204,6 +205,7 @@ static struct task_struct* psnedf_schedule(struct task_struct * prev)
         }
 
         pedf->scheduled = next;
+        sched_state_task_picked();
         raw_spin_unlock(&pedf->slock);
 
         return next;