2 files changed, 12 insertions, 31 deletions

diff --git a/litmus/rt_domain.c b/litmus/rt_domain.c
index 81a5ac16f164..fd0e6c235e1b 100644
--- a/litmus/rt_domain.c
+++ b/litmus/rt_domain.c
@@ -304,6 +304,7 @@ void __add_ready(rt_domain_t* rt, struct task_struct *new)
 
         BUG_ON(bheap_node_in_heap(tsk_rt(new)->heap_node));
 
+        new->rt_param.domain = rt;
         bheap_insert(rt->order, &rt->ready_queue, tsk_rt(new)->heap_node);
         rt->check_resched(rt);
 }
diff --git a/litmus/sched_mc.c b/litmus/sched_mc.c
index 0bad4516ba50..457bd2a94a29 100644
--- a/litmus/sched_mc.c
+++ b/litmus/sched_mc.c
@@ -238,33 +238,6 @@ static void mc_edf_domain_init(rt_domain_t* rt, check_resched_needed_t resched,
         rt_domain_init(rt, mc_edf_ready_order, resched, release);
 }
 
-
-/* END clone of edf_common.c*/
-
-/* Return the domain of a task */
-static rt_domain_t* domain_of(struct task_struct* task)
-{
-        switch (tsk_mc_crit(task))
-        {
-                case CRIT_LEVEL_A:
-                        return remote_a_queue(get_partition(task));
-                        break;
-                case CRIT_LEVEL_B:
-                        return remote_b_queue(get_partition(task));
-                        break;
-                case CRIT_LEVEL_C:
-                        return &crit_c;
-                        break;
-                case CRIT_LEVEL_D:
-                        return &crit_d;
-                        break;
-                case NUM_CRIT_LEVELS:
-                default:
-                        /*Should never get here*/
-                        BUG();
-        }
-}
-
 #define WANT_ALL_SCHED_EVENTS
 
 /* Called by update_cpu_position and lowest_prio_cpu in bheap operations
@@ -586,7 +559,7 @@ static noinline void unlink(struct task_struct* t)
                  * case.
                  */
                 TRACE("Weird is_queued situation happened\n");
-                remove(domain_of(t), t);
+                remove(tsk_rt(t)->domain, t);
         }
 }
 
@@ -612,10 +585,10 @@ static noinline void requeue(struct task_struct* task)
         BUG_ON(is_queued(task));
 
         if (is_released(task, litmus_clock()))
-                __add_ready(domain_of(task), task);
+                __add_ready(tsk_rt(task)->domain, task);
         else {
                 /* it has got to wait */
-                add_release(domain_of(task), task);
+                add_release(tsk_rt(task)->domain, task);
         }
 }
 
@@ -1078,15 +1051,22 @@ static void mc_finish_switch(struct task_struct *prev)
  *      No lock required by caller
  *      Obtains lock and calls mc_job_arrival before releasing lock
  */
-static void mc_task_new(struct task_struct * t, int on_rq, int running)
+static void mc_task_new(struct task_struct *t, int on_rq, int running)
 {
         unsigned long           flags;
         cpu_entry_t*            entry;
+        enum crit_level         lvl;
 
         TRACE("mixed crit: task new %d\n", t->pid);
 
         raw_spin_lock_irqsave(&global_lock, flags);
 
+        lvl = tsk_mc_crit(t);
+        t->rt_param.domain =
+                (lvl == CRIT_LEVEL_A) ? remote_a_queue(get_partition(t)) :
+                (lvl == CRIT_LEVEL_B) ? remote_b_queue(get_partition(t)) :
+                (lvl == CRIT_LEVEL_C) ? &crit_c : &crit_d;
+
         /* setup job params */
         release_at(t, litmus_clock());
         tsk_mc_data(t)->mc_job.ghost_budget = 0;