Cleanup priority tracking for budget enforcement.wip-gpusync-merge

2 files changed, 43 insertions, 28 deletions

diff --git a/include/litmus/sbinheap.h b/include/litmus/sbinheap.h
index 3f952beca63b..a108d9e92288 100644
--- a/include/litmus/sbinheap.h
+++ b/include/litmus/sbinheap.h
@@ -150,6 +150,12 @@ static inline int sbinheap_empty(struct sbinheap *heap)
         return(heap->size == 0);
 }
 
+/* Returns true if the sbinheap is full. */
+static inline int sbinheap_full(struct sbinheap *heap)
+{
+        return(heap->size == heap->max_size);
+}
+
 /* Get the maximum size of the heap */
 static inline idx_t sbinheap_max_size(struct sbinheap *heap)
 {
diff --git a/litmus/sched_cedf.c b/litmus/sched_cedf.c
index 61fef55f38a5..760fb4269e00 100644
--- a/litmus/sched_cedf.c
+++ b/litmus/sched_cedf.c
@@ -258,44 +258,53 @@ static void cedf_untrack_in_top_m(struct task_struct *t)
 {
         /* cluster lock must be held */
         cedf_domain_t *cluster = task_cpu_cluster(t);
+        int exited_top_m = 0;
 
-        if (!sbinheap_is_in_heap(&tsk_rt(t)->budget.top_m_node) &&
-                !binheap_is_in_heap(&tsk_rt(t)->budget.not_top_m_node))
-                return;
+        BUG_ON(sbinheap_is_in_heap(&tsk_rt(t)->budget.top_m_node) &&
+               binheap_is_in_heap(&tsk_rt(t)->budget.not_top_m_node));
+
+        if (sbinheap_is_in_heap(&tsk_rt(t)->budget.top_m_node)) {
+                exited_top_m = bt_flag_is_set(t, BTF_IS_TOP_M);
+                WARN_ON_ONCE(!exited_top_m);
 
-        if (bt_flag_is_set(t, BTF_IS_TOP_M)) {
-                /* delete t's entry */
                 sbinheap_delete(&tsk_rt(t)->budget.top_m_node, &cluster->top_m);
                 INIT_SBINHEAP_NODE(&tsk_rt(t)->budget.top_m_node);
-                budget_state_machine(t,on_exit_top_m);
-                bt_flag_clear(t, BTF_IS_TOP_M);
 
-                /* move a task over from the overflow heap */
-                if(!binheap_empty(&cluster->not_top_m)) {
-                        struct budget_tracker *bt =
-                                binheap_top_entry(&cluster->not_top_m, struct budget_tracker,
-                                        not_top_m_node);
-                        struct task_struct *to_move =
-                                container_of(
-                                         container_of(bt, struct rt_param, budget),
-                                                 struct task_struct,
-                                                 rt_param);
-
-                        binheap_delete_root(&cluster->not_top_m, struct budget_tracker,
-                                                not_top_m_node);
-                        INIT_BINHEAP_NODE(&tsk_rt(to_move)->budget.not_top_m_node);
-
-                        sbinheap_add(&tsk_rt(to_move)->budget.top_m_node,
-                                                &cluster->top_m,
-                                                struct budget_tracker, top_m_node);
-                        bt_flag_set(to_move, BTF_IS_TOP_M);
-                        budget_state_machine(to_move,on_enter_top_m);
+                if (likely(exited_top_m)) {
+                        budget_state_machine(t,on_exit_top_m);
+                        bt_flag_clear(t, BTF_IS_TOP_M);
                 }
         }
-        else {
+        else if(likely(binheap_is_in_heap(&tsk_rt(t)->budget.not_top_m_node))) {
                 binheap_delete(&tsk_rt(t)->budget.not_top_m_node, &cluster->not_top_m);
                 INIT_BINHEAP_NODE(&tsk_rt(t)->budget.not_top_m_node);
         }
+
+        /* move a task over from the overflow heap */
+        if(!sbinheap_full(&cluster->top_m) &&
+           !binheap_empty(&cluster->not_top_m)) {
+                struct budget_tracker *bt =
+                        binheap_top_entry(&cluster->not_top_m, struct budget_tracker,
+                                not_top_m_node);
+                struct task_struct *to_move =
+                        container_of(
+                                 container_of(bt, struct rt_param, budget),
+                                         struct task_struct,
+                                         rt_param);
+
+                /* task should have already been moved into the top m */
+                WARN_ON_ONCE(!exited_top_m);
+
+                binheap_delete_root(&cluster->not_top_m, struct budget_tracker,
+                                        not_top_m_node);
+                INIT_BINHEAP_NODE(&tsk_rt(to_move)->budget.not_top_m_node);
+
+                sbinheap_add(&tsk_rt(to_move)->budget.top_m_node,
+                                        &cluster->top_m,
+                                        struct budget_tracker, top_m_node);
+                bt_flag_set(to_move, BTF_IS_TOP_M);
+                budget_state_machine(to_move,on_enter_top_m);
+        }
 }