1 files changed, 15 insertions, 7 deletions

diff --git a/litmus/sched_psn_edf.c b/litmus/sched_psn_edf.c
index eaaec38f43da..4e117be9546b 100644
--- a/litmus/sched_psn_edf.c
+++ b/litmus/sched_psn_edf.c
@@ -17,6 +17,7 @@
 #include <litmus/litmus.h>
 #include <litmus/jobs.h>
 #include <litmus/preempt.h>
+#include <litmus/budget.h>
 #include <litmus/sched_plugin.h>
 #include <litmus/edf_common.h>
 #include <litmus/sched_trace.h>
@@ -132,6 +133,15 @@ static void unboost_priority(struct task_struct* t)
 
 #endif
 
+static int psnedf_preempt_check(psnedf_domain_t *pedf)
+{
+        if (edf_preemption_needed(&pedf->domain, pedf->scheduled)) {
+                preempt(pedf);
+                return 1;
+        } else
+                return 0;
+}
+
 /* This check is trivial in partioned systems as we only have to consider
  * the CPU of the partition.
  */
@@ -142,11 +152,7 @@ static int psnedf_check_resched(rt_domain_t *edf)
         /* because this is a callback from rt_domain_t we already hold
          * the necessary lock for the ready queue
          */
-        if (edf_preemption_needed(edf, pedf->scheduled)) {
-                preempt(pedf);
-                return 1;
-        } else
-                return 0;
+        return psnedf_preempt_check(pedf);
 }
 
 static void job_completion(struct task_struct* t, int forced)
@@ -301,7 +307,7 @@ static void psnedf_task_new(struct task_struct * t, int on_rq, int running)
         } else {
                 requeue(t, edf);
                 /* maybe we have to reschedule */
-                preempt(pedf);
+                psnedf_preempt_check(pedf);
         }
         raw_spin_unlock_irqrestore(&pedf->slock, flags);
 }
@@ -337,8 +343,10 @@ static void psnedf_task_wake_up(struct task_struct *task)
          * de-scheduling the task, i.e., wake_up() raced with schedule()
          * and won.
          */
-        if (pedf->scheduled != task)
+        if (pedf->scheduled != task) {
                 requeue(task, edf);
+                psnedf_preempt_check(pedf);
+        }
 
         raw_spin_unlock_irqrestore(&pedf->slock, flags);
         TRACE_TASK(task, "wake up done\n");