1 files changed, 16 insertions, 11 deletions

diff --git a/litmus/sched_pfp.c b/litmus/sched_pfp.c
index ad2b8d81f29f..891469054ab4 100644
--- a/litmus/sched_pfp.c
+++ b/litmus/sched_pfp.c
@@ -175,7 +175,7 @@ static void pfp_tick(struct task_struct *t)
                 tsk_rt(t)->budget.ops && budget_quantum_tracked(t) &&
                 budget_exhausted(t)) {
                 TRACE_TASK(t, "budget exhausted\n");
-                tsk_rt(t)->budget.ops->on_exhausted(t);
+                budget_state_machine(t,on_exhausted);
         }
 }
 
@@ -213,14 +213,12 @@ static struct task_struct* pfp_schedule(struct task_struct * prev)
         resched = preempt;
 
         /* Do budget stuff */
-        if (tsk_rt(prev)->budget.ops) {
-                if (blocks)
-                        tsk_rt(prev)->budget.ops->on_blocked(prev);
-                else if (sleep)
-                        tsk_rt(prev)->budget.ops->on_sleep(prev);
-                else if (preempt)
-                        tsk_rt(prev)->budget.ops->on_preempt(prev);
-        }
+        if (blocks)
+                budget_state_machine(prev,on_blocked);
+        else if (sleep)
+                budget_state_machine(prev,on_sleep);
+        else if (preempt)
+                budget_state_machine(prev,on_preempt);
 
         /* If a task blocks we have no choice but to reschedule.
          */
@@ -389,6 +387,8 @@ static void pfp_task_wake_up(struct task_struct *task)
                 sched_trace_task_release(task);
         }
 
+        budget_state_machine(task,on_wakeup);
+
         /* Only add to ready queue if it is not the currently-scheduled
          * task. This could be the case if a task was woken up concurrently
          * on a remote CPU before the executing CPU got around to actually
@@ -436,8 +436,7 @@ static void pfp_task_exit(struct task_struct * t)
         raw_readyq_lock_irqsave(&pfp->slock, flags);
 
         /* disable budget enforcement */
-        if (tsk_rt(t)->budget.ops)
-                tsk_rt(t)->budget.ops->on_exit(t);
+        budget_state_machine(t,on_exit);
 
         if (is_queued(t)) {
                 BUG(); /* This currently doesn't work. */
@@ -485,8 +484,14 @@ static void fp_set_prio_inh(pfp_domain_t* pfp, struct task_struct* t,
                 /* first remove */
                 fp_dequeue(pfp, t);
 
+        if (t->rt_param.inh_task)
+                budget_state_machine2(t,t->rt_param.inh_task,on_disinherit);
+
         t->rt_param.inh_task = prio_inh;
 
+        if (prio_inh)
+                budget_state_machine2(t,prio_inh,on_inherit);
+
         if (requeue)
                 /* add again to the right queue */
                 fp_prio_add(&pfp->ready_queue, t, priority_index(t));