1 files changed, 3 insertions, 1 deletions

diff --git a/litmus/budget.c b/litmus/budget.c
index b99177a26313..310e9a3d4172 100644
--- a/litmus/budget.c
+++ b/litmus/budget.c
@@ -1,7 +1,9 @@
+#include <linux/sched.h>
 #include <linux/percpu.h>
 #include <linux/hrtimer.h>
 
 #include <litmus/litmus.h>
+#include <litmus/preempt.h>
 
 struct enforcement_timer {
         /* The enforcement timer is used to accurately police
@@ -23,7 +25,7 @@ static enum hrtimer_restart on_enforcement_timeout(struct hrtimer *timer)
         TRACE("enforcement timer fired.\n");
         et->armed = 0;
         /* activate scheduler */
-        set_tsk_need_resched(current);
+        litmus_reschedule_local();
         local_irq_restore(flags);
 
         return  HRTIMER_NORESTART;