1 files changed, 11 insertions, 20 deletions

diff --git a/litmus/ce_domain.c b/litmus/ce_domain.c
index 5b4fd1cb438f..d6f713fbf789 100644
--- a/litmus/ce_domain.c
+++ b/litmus/ce_domain.c
@@ -7,6 +7,7 @@
 #include <litmus/rt_param.h>
 #include <litmus/domain.h>
 #include <litmus/sched_mc.h>
+#include <litmus/ce_domain.h>
 
 /*
  * Called for:
@@ -41,24 +42,13 @@ void ce_requeue(domain_t *dom, struct task_struct *ts)
         }
 }
 
+void mc_ce_task_exit(struct task_struct*);
 /*
  * Called when a task exits the system.
  */
-void ce_exit(domain_t *dom, struct task_struct *ts)
+void ce_task_exit(domain_t *dom, struct task_struct *ts)
 {
-        struct ce_dom_data *ce_data = dom->data;
-        const int lvl_a_id = tsk_mc_data(ts)->mc_task.lvl_a_id;
-        struct pid *pid;
-
-        BUG_ON(task_cpu(ts) != get_partition(ts));
-        BUG_ON(CRIT_LEVEL_A != tsk_mc_crit(ts));
-        BUG_ON(lvl_a_id >= ce_data->num_pid_entries);
-        pid = ce_data->pid_entries[lvl_a_id].pid;
-        BUG_ON(!pid);
-        put_pid(pid);
-        ce_data->pid_entries[lvl_a_id].pid = NULL;
-        if (ce_data->should_schedule == ts)
-                ce_data->should_schedule = NULL;
+        mc_ce_task_exit(ts);
 }
 
 /*
@@ -80,11 +70,12 @@ struct task_struct* ce_peek_and_take_ready(domain_t *dom)
         return ret;
 }
 
-int ce_higher_prio(domain_t *dom, struct task_struct *_a,
+int ce_higher_prio(struct task_struct *_a,
                 struct task_struct *_b)
 {
         const struct task_struct *a = _a;
-        struct ce_dom_data *ce_data = dom->data;
+        const domain_t *dom = get_task_domain(a);
+        const struct ce_dom_data *ce_data = dom->data;
         return (a == ce_data->should_schedule);
 }
 
@@ -102,18 +93,18 @@ void ce_start(struct task_struct *ts, lt_t start)
         mc_ce_release_at(ts, start);
 }
 
+domain_data_t *ce_domain_for(int);
 long mc_ce_activate_plugin(void);
-domain_t *ce_domain_for(int);
 long ce_activate_plugin(void)
 {
-        domain_t *dom;
+        domain_data_t *dom_data;
         struct ce_dom_data *ce_data;
         int cpu;
 
         /* first change the timer callback function */
         for_each_online_cpu(cpu) {
-                dom = ce_domain_for(cpu);
-                ce_data = dom->data;
+                dom_data = ce_domain_for(cpu);
+                ce_data = dom_data->domain.data;
                 ce_data->timer.function = ce_timer_function;
         }
         /* then run the regular CE activate plugin */