1 files changed, 15 insertions, 14 deletions

diff --git a/litmus/budget.c b/litmus/budget.c
index 310e9a3d4172..84f3f22770b1 100644
--- a/litmus/budget.c
+++ b/litmus/budget.c
@@ -4,13 +4,8 @@
 
 #include <litmus/litmus.h>
 #include <litmus/preempt.h>
-
-struct enforcement_timer {
-        /* The enforcement timer is used to accurately police
-         * slice budgets. */
-        struct hrtimer          timer;
-        int                     armed;
-};
+#include <litmus/budget.h>
+#include <litmus/sched_trace.h>
 
 DEFINE_PER_CPU(struct enforcement_timer, budget_timer);
 
@@ -32,7 +27,7 @@ static enum hrtimer_restart on_enforcement_timeout(struct hrtimer *timer)
 }
 
 /* assumes called with IRQs off */
-static void cancel_enforcement_timer(struct enforcement_timer* et)
+void cancel_enforcement_timer(struct enforcement_timer* et)
 {
         int ret;
 
@@ -54,11 +49,10 @@ static void cancel_enforcement_timer(struct enforcement_timer* et)
 }
 
 /* assumes called with IRQs off */
-static void arm_enforcement_timer(struct enforcement_timer* et,
-                                  struct task_struct* t)
+void arm_enforcement_timer(struct enforcement_timer* et,
+                           struct task_struct* t)
 {
         lt_t when_to_fire;
-        TRACE_TASK(t, "arming enforcement timer.\n");
 
         /* Calling this when there is no budget left for the task
          * makes no sense, unless the task is non-preemptive. */
@@ -67,8 +61,11 @@ static void arm_enforcement_timer(struct enforcement_timer* et,
         /* __hrtimer_start_range_ns() cancels the timer
          * anyway, so we don't have to check whether it is still armed */
 
-        if (likely(!is_np(t))) {
+        if (likely(!is_user_np(t))) {
                 when_to_fire = litmus_clock() + budget_remaining(t);
+                TRACE_TASK(t, "arming enforcement timer for %llu.\n",
+                           when_to_fire);
+
                 __hrtimer_start_range_ns(&et->timer,
                                          ns_to_ktime(when_to_fire),
                                          0 /* delta */,
@@ -94,6 +91,11 @@ void update_enforcement_timer(struct task_struct* t)
         }
 }
 
+void init_enforcement_timer(struct enforcement_timer *et)
+{
+        hrtimer_init(&et->timer, CLOCK_MONOTONIC, HRTIMER_MODE_ABS);
+        et->timer.function = on_enforcement_timeout;
+}
 
 static int __init init_budget_enforcement(void)
 {
@@ -102,8 +104,7 @@ static int __init init_budget_enforcement(void)
 
         for (cpu = 0; cpu < NR_CPUS; cpu++)  {
                 et = &per_cpu(budget_timer, cpu);
-                hrtimer_init(&et->timer, CLOCK_MONOTONIC, HRTIMER_MODE_ABS);
-                et->timer.function = on_enforcement_timeout;
+                init_enforcement_timer(et);
         }
         return 0;
 }