1 files changed, 10 insertions, 0 deletions

diff --git a/litmus/sched_psn_edf.c b/litmus/sched_psn_edf.c
index 4e117be9546..b26a5ff439c 100644
--- a/litmus/sched_psn_edf.c
+++ b/litmus/sched_psn_edf.c
@@ -23,6 +23,8 @@
 #include <litmus/sched_trace.h>
 #include <litmus/trace.h>
 
+#include <litmus/color.h>
+
 typedef struct {
         rt_domain_t             domain;
         int                     cpu;
@@ -271,6 +273,14 @@ static struct task_struct* psnedf_schedule(struct task_struct * prev)
                 TRACE("becoming idle at %llu\n", litmus_clock());
         }
 
+        if (prev && next != prev && is_realtime(prev) && get_rt_job(prev) > 1) {
+                color_sched_out_task(prev);
+        }
+
+        if (next && next != prev && is_realtime(next)) {
+                color_sched_in_task(next);
+        }
+
         pedf->scheduled = next;
         sched_state_task_picked();
         raw_spin_unlock(&pedf->slock);