2 files changed, 36 insertions, 9 deletions

diff --git a/include/litmus/reservations/gedf_reservation.h b/include/litmus/reservations/gedf_reservation.h
index d9694fa3ac64..578b9cad5838 100644
--- a/include/litmus/reservations/gedf_reservation.h
+++ b/include/litmus/reservations/gedf_reservation.h
@@ -11,6 +11,7 @@
 struct gedf_reservation {
         struct reservation res;
         struct gedf_cpu_entry* linked_on;
+        struct gedf_cpu_entry* scheduled_on;
         lt_t cur_budget;
         int will_remove;
         int blocked;
diff --git a/litmus/reservations/gedf_reservation.c b/litmus/reservations/gedf_reservation.c
index e08d3f9799e8..9f02617821a1 100644
--- a/litmus/reservations/gedf_reservation.c
+++ b/litmus/reservations/gedf_reservation.c
@@ -103,13 +103,28 @@ static void link_task_to_cpu(
         struct gedf_reservation* linked,
         struct gedf_cpu_entry* entry)
 {
+        struct gedf_cpu_entry* on_cpu;
+        struct gedf_reservation* tmp;
 
         if (entry->linked)
                 entry->linked->linked_on = NULL;
 
-        if (linked)
-                linked->linked_on = entry;
-
+        if (linked) {
+                on_cpu = linked->scheduled_on;
+                if (on_cpu) {
+                        BUG_ON(on_cpu->linked == linked);
+
+                        if (entry != on_cpu) {
+                                tmp = on_cpu->linked;
+                                linked->linked_on = on_cpu;
+                                on_cpu->linked = linked;
+                                update_cpu_position(on_cpu, &gedf_env->cpu_heap);
+                                linked = tmp;
+                        }
+                }
+                if (linked)
+                        linked->linked_on = entry;
+        }
         entry->linked = linked;
         update_cpu_position(entry, &gedf_env->cpu_heap);
 }
@@ -519,8 +534,12 @@ static void gedf_env_suspend(
                 unlink(gedf_env, entry->linked);
                 requeue(gedf_env, tmp);
         }
-        if (entry->scheduled && entry->scheduled->res.ops->on_preempt)
-                entry->scheduled->res.ops->on_preempt(&entry->scheduled->res, cpu);
+        if (entry->scheduled) {
+                if (entry->scheduled->res.ops->on_preempt)
+                        entry->scheduled->res.ops->on_preempt(&entry->scheduled->res, cpu);
+                if (entry->scheduled->scheduled_on == entry)
+                        entry->scheduled->scheduled_on = NULL;
+        }
         entry->scheduled = NULL;
 
         /* this essentially removes the cpu from scheduling consideration */
@@ -641,10 +660,17 @@ static struct task_struct* gedf_env_dispatch(
 
         /* if linked and scheduled differ, preempt and schedule accordingly */
         if (!np && entry->scheduled != entry->linked) {
-                if (entry->scheduled && entry->scheduled->res.ops->on_preempt)
-                        entry->scheduled->res.ops->on_preempt(&entry->scheduled->res, cpu);
-                if (entry->linked && entry->linked->res.ops->on_schedule)
-                        entry->linked->res.ops->on_schedule(&entry->linked->res, cpu);
+                if (entry->scheduled) {
+                        if (entry->scheduled->res.ops->on_preempt)
+                                entry->scheduled->res.ops->on_preempt(&entry->scheduled->res, cpu);
+                        if (entry->scheduled->scheduled_on == entry)
+                                entry->scheduled->scheduled_on = NULL;
+                }
+                if (entry->linked) {
+                        entry->linked->scheduled_on = entry;
+                        if (entry->linked->res.ops->on_schedule)
+                                entry->linked->res.ops->on_schedule(&entry->linked->res, cpu);
+                }
                 entry->scheduled = entry->linked;
         }
         raw_spin_unlock_irqrestore(&gedf_env->domain.ready_lock, flags);