Per-task budget high-resolution timers (hrtimers).

As a step towards implementing more complex budget tracking
method (ex. BWI, VXR, etc.), we need per-task budget trackers
because we may be required to drain budget from a task, even
while it is suspended or blocked.
This patch does:
1) Replaces the per-CPU hrtimers with per-task hrtimers.
2) Plugin architecture for different budget policies.

This patch creates three budget draining policies:
SIMPLE, SAWARE (suspension-aware), and SOBLIV (suspension-oblivious).
However, only SIMPLE is supported by this patch.

SIMPLE (default):
Budget drains while the task is scheduled. Budget is preserved
across self-suspensions (but not job completions, of course).
Only SIMPLE is supported in this patch. (Maintaining current Litmus
functionality.)

SAWARE:
Draining according to suspension-aware analysis. Budget should drain
whenever a task is among the top-m tasks in its cluster, where
m is the number of processors in said cluster. This draining should
happen whether or not the task is actually scheduled.

SOBLIV:
Draining according to suspension-oblivious analysis. Budget should
drain whenever the task is scheduled or suspended (but not due to
preemption). Exception: Draining should halt when we can prove that
the task is not among the top-m tasks blocked on the same lock (i.e.,
on the PQ in the OMLP-family locking protocols).

15 files changed, 566 insertions, 275 deletions

diff --git a/include/litmus/budget.h b/include/litmus/budget.h
index 763b31c0e9f6..2a3511245f7a 100644
--- a/include/litmus/budget.h
+++ b/include/litmus/budget.h
@@ -1,26 +1,14 @@
 #ifndef _LITMUS_BUDGET_H_
 #define _LITMUS_BUDGET_H_
 
-/* Update the per-processor enforcement timer (arm/reproram/cancel) for
- * the next task. */
-void update_enforcement_timer(struct task_struct* t);
+#include <linux/hrtimer.h>
+#include <linux/semaphore.h>
 
-/* Send SIG_BUDGET to a real-time task. */
-void send_sigbudget(struct task_struct* t);
+#define budget_exhausted(t) \
+        (get_exec_time(t) >= get_exec_cost(t))
 
-inline static int budget_exhausted(struct task_struct* t)
-{
-        return get_exec_time(t) >= get_exec_cost(t);
-}
-
-inline static lt_t budget_remaining(struct task_struct* t)
-{
-        if (!budget_exhausted(t))
-                return get_exec_cost(t) - get_exec_time(t);
-        else
-                /* avoid overflow */
-                return 0;
-}
+#define budget_remaining(t) \
+        ((!budget_exhausted(t)) ? (get_exec_cost(t) - get_exec_time(t)) : 0)
 
 #define budget_enforced(t) (\
         tsk_rt(t)->task_params.budget_policy != NO_ENFORCEMENT)
@@ -29,21 +17,87 @@ inline static lt_t budget_remaining(struct task_struct* t)
         tsk_rt(t)->task_params.budget_policy == PRECISE_ENFORCEMENT || \
         tsk_rt(t)->task_params.budget_signal_policy == PRECISE_SIGNALS)
 
+#define budget_quantum_tracked(t) (\
+        tsk_rt(t)->task_params.budget_policy == QUANTUM_ENFORCEMENT || \
+        tsk_rt(t)->task_params.budget_signal_policy == QUANTUM_SIGNALS)
+
 #define budget_signalled(t) (\
         tsk_rt(t)->task_params.budget_signal_policy != NO_SIGNALS)
 
 #define budget_precisely_signalled(t) (\
         tsk_rt(t)->task_params.budget_policy == PRECISE_SIGNALS)
 
-#define sigbudget_sent(t) (\
-        test_bit(RT_JOB_SIG_BUDGET_SENT, &tsk_rt(t)->job_params.flags))
+#define bt_flag_is_set(t, flag_nr) (\
+        test_bit(flag_nr, &tsk_rt(t)->budget.flags))
+
+#define bt_flag_test_and_set(t, flag_nr) (\
+        test_and_set_bit(flag_nr, &tsk_rt(t)->budget.flags))
+
+#define bt_flag_set(t, flag_nr) (\
+        set_bit(flag_nr, &tsk_rt(t)->budget.flags))
+
+#define bt_flag_clear(t, flag_nr) (\
+        clear_bit(flag_nr, &tsk_rt(t)->budget.flags))
+
+#define bt_flags_reset(t) (\
+        tsk_rt(t)->budget.flags = 0)
+
+#define requeue_preempted_job(t) \
+        (t && (!budget_exhausted(t) || !budget_enforced(t)))
+
+struct enforcement_timer
+{
+        raw_spinlock_t lock;
+        struct hrtimer timer;
+        int armed:1;
+};
+
+typedef void (*scheduled_t)(struct task_struct* t);
+typedef void (*blocked_t)(struct task_struct* t);
+typedef void (*preempt_or_sleep_t)(struct task_struct* t);
+typedef void (*exhausted_t)(struct task_struct* t);
+typedef void (*exit_t)(struct task_struct* t);
 
-static inline int requeue_preempted_job(struct task_struct* t)
+struct budget_tracker_ops
 {
-        /* Add task to ready queue only if not subject to budget enforcement or
-         * if the job has budget remaining. t may be NULL.
-         */
-        return t && (!budget_exhausted(t) || !budget_enforced(t));
-}
+        scheduled_t                     on_scheduled;   /* called from litmus_schedule(). */
+        blocked_t                       on_blocked;             /* called from plugin::schedule() */
+        preempt_or_sleep_t      on_preempt_or_sleep;    /* called from plugin::schedule() */
+
+        exit_t                          on_exit;                /* task exiting rt mode */
+
+        exhausted_t                     on_exhausted;   /* called by plugin::tick() or timer interrupt */
+};
+
+struct budget_tracker
+{
+        struct enforcement_timer timer;
+        const struct budget_tracker_ops* ops;
+        unsigned long flags;
+};
+
+/* budget tracker flags */
+enum BT_FLAGS
+{
+        BTF_BUDGET_EXHAUSTED = 0,
+        BTF_SIG_BUDGET_SENT     = 1,
+};
+
+/* Functions for simple DRAIN_SIMPLE policy common
+ * to every scheduler. Scheduler must provided
+ * implementation for simple_on_exhausted().
+ */
+void simple_on_scheduled(struct task_struct* t);
+void simple_on_blocked(struct task_struct* t);
+void simple_on_preempt_or_sleep(struct task_struct* t);
+void simple_on_exit(struct task_struct* t);
+
+
+void init_budget_tracker(struct budget_tracker* bt,
+                                const struct budget_tracker_ops* ops);
+
+
+/* Send SIG_BUDGET to a real-time task. */
+void send_sigbudget(struct task_struct* t);
 
 #endif
diff --git a/include/litmus/fdso.h b/include/litmus/fdso.h
index f7887288d8f5..2d8e6c43d908 100644
--- a/include/litmus/fdso.h
+++ b/include/litmus/fdso.h
@@ -36,7 +36,7 @@ typedef enum  {
         KFMLP_GPU_AFF_OBS                       = 12,
 
         PRIOQ_MUTEX = 13,
-        
+
         MAX_OBJ_TYPE    = 13
 } obj_type_t;
 
diff --git a/include/litmus/gpu_affinity.h b/include/litmus/gpu_affinity.h
index 47da725717b0..f610f58b1f3b 100644
--- a/include/litmus/gpu_affinity.h
+++ b/include/litmus/gpu_affinity.h
@@ -11,6 +11,7 @@ gpu_migration_dist_t gpu_migration_distance(int a, int b);
 static inline void reset_gpu_tracker(struct task_struct* t)
 {
         t->rt_param.accum_gpu_time = 0;
+        t->rt_param.gpu_time_stamp = 0;
 }
 
 static inline void start_gpu_tracker(struct task_struct* t)
@@ -22,11 +23,16 @@ static inline void stop_gpu_tracker(struct task_struct* t)
 {
         lt_t now = litmus_clock();
         t->rt_param.accum_gpu_time += (now - t->rt_param.gpu_time_stamp);
+        t->rt_param.gpu_time_stamp = 0;
 }
 
 static inline lt_t get_gpu_time(struct task_struct* t)
 {
-        return t->rt_param.accum_gpu_time;
+        lt_t accum = t->rt_param.accum_gpu_time;
+        if (t->rt_param.gpu_time_stamp != 0) {
+                accum += (litmus_clock() - t->rt_param.gpu_time_stamp);
+        }
+        return accum;
 }
 
 static inline lt_t get_gpu_estimate(struct task_struct* t, gpu_migration_dist_t dist)
diff --git a/include/litmus/litmus.h b/include/litmus/litmus.h
index 70b421d59d34..f6ea5f6e80ee 100644
--- a/include/litmus/litmus.h
+++ b/include/litmus/litmus.h
@@ -62,6 +62,7 @@ void litmus_exit_task(struct task_struct *tsk);
 #define get_priority(t)         (tsk_rt(t)->task_params.priority)
 #define get_class(t)        (tsk_rt(t)->task_params.cls)
 #define get_release_policy(t) (tsk_rt(t)->task_params.release_policy)
+#define get_drain_policy(t) (tsk_rt(t)->task_params.drain_policy)
 
 /* job_param macros */
 #define get_exec_time(t)    (tsk_rt(t)->job_params.exec_time)
@@ -69,6 +70,7 @@ void litmus_exit_task(struct task_struct *tsk);
 #define get_period(t)           (tsk_rt(t)->task_params.period)
 #define get_release(t)          (tsk_rt(t)->job_params.release)
 #define get_lateness(t)         (tsk_rt(t)->job_params.lateness)
+#define get_budget_timer(t) (tsk_rt(t)->job_params.budget_timer)
 
 #define effective_priority(t) ((!(tsk_rt(t)->inh_task)) ? t : tsk_rt(t)->inh_task)
 #define base_priority(t) (t)
diff --git a/include/litmus/rt_param.h b/include/litmus/rt_param.h
index bf0ee8dbae6e..887075b908ca 100644
--- a/include/litmus/rt_param.h
+++ b/include/litmus/rt_param.h
@@ -36,6 +36,14 @@ typedef enum {
 } budget_policy_t;
 
 typedef enum {
+        /* all drain mechanisms are ignored if budget enforcement or signalling
+           is not in use. */
+        DRAIN_SIMPLE,   /* drains while task is linked */
+        DRAIN_SAWARE,   /* drains according to suspension-aware analysis */
+        DRAIN_SOBLIV,   /* drains according to suspension-obliv analysis */
+} budget_drain_policy_t;
+
+typedef enum {
         NO_SIGNALS,                     /* job receives no signals when it exhausts its budget */
         QUANTUM_SIGNALS,        /* budget signals are only sent on quantum boundaries */
         PRECISE_SIGNALS,        /* budget signals are triggered with hrtimers */
@@ -132,6 +140,7 @@ struct rt_task {
         unsigned int    priority;
         task_class_t    cls;
         budget_policy_t budget_policy; /* ignored by pfair */
+        budget_drain_policy_t drain_policy;
         budget_signal_policy_t budget_signal_policy; /* currently ignored by pfair */
         release_policy_t release_policy;
 };
@@ -213,8 +222,9 @@ struct control_page {
 /* don't export internal data structures to user space (liblitmus) */
 #ifdef __KERNEL__
 
-#include <litmus/binheap.h>
 #include <linux/semaphore.h>
+#include <litmus/budget.h>
+#include <litmus/binheap.h>
 
 #ifdef CONFIG_LITMUS_SOFTIRQD
 #include <linux/interrupt.h>
@@ -247,15 +257,8 @@ struct rt_job {
          * Increase this sequence number when a job is released.
          */
         unsigned int    job_no;
-
-        /* bits:
-         * 0th: Set if a budget exhaustion signal has already been sent for
-         *      the current job. */
-        unsigned long   flags;
 };
 
-#define RT_JOB_SIG_BUDGET_SENT  0
-
 struct pfair_param;
 
 enum klmirqd_sem_status
@@ -278,12 +281,12 @@ typedef enum gpu_migration_dist
         MIG_LAST = MIG_NONE
 } gpu_migration_dist_t;
 
-typedef struct feedback_est{
+typedef struct feedback_est
+{
         fp_t est;
         fp_t accum_err;
 } feedback_est_t;
 
-
 #define AVG_EST_WINDOW_SIZE 20
 
 typedef int (*notify_rsrc_exit_t)(struct task_struct* tsk);
@@ -417,7 +420,6 @@ struct rt_param {
         struct binheap_node     aux_task_owner_node;
 #endif
 
-
 #ifdef CONFIG_NP_SECTION
         /* For the FMLP under PSN-EDF, it is required to make the task
          * non-preemptive from kernel space. In order not to interfere with
@@ -427,6 +429,8 @@ struct rt_param {
         unsigned int    kernel_np;
 #endif
 
+        struct budget_tracker budget;
+
         /* This field can be used by plugins to store where the task
          * is currently scheduled. It is the responsibility of the
          * plugin to avoid race conditions.
diff --git a/include/litmus/sched_plugin.h b/include/litmus/sched_plugin.h
index 78cec30866ac..6e7d6df2fb78 100644
--- a/include/litmus/sched_plugin.h
+++ b/include/litmus/sched_plugin.h
@@ -130,6 +130,7 @@ struct sched_plugin {
         schedule_t              schedule;
         finish_switch_t         finish_switch;
 
+
         /*      syscall backend         */
         complete_job_t          complete_job;
         release_at_t            release_at;
diff --git a/litmus/budget.c b/litmus/budget.c
index 2ec9d383c332..559c54709acc 100644
--- a/litmus/budget.c
+++ b/litmus/budget.c
@@ -5,102 +5,100 @@
 
 #include <litmus/litmus.h>
 #include <litmus/preempt.h>
-
+#include <litmus/sched_plugin.h>
 #include <litmus/budget.h>
 #include <litmus/signal.h>
 
-struct enforcement_timer {
-        /* The enforcement timer is used to accurately police
-         * slice budgets. */
-        struct hrtimer          timer;
-        int                     armed;
-};
-
-static DEFINE_PER_CPU(struct enforcement_timer, budget_timer);
-
-static enum hrtimer_restart on_enforcement_timeout(struct hrtimer *timer)
+inline static void cancel_enforcement_timer(struct task_struct* t)
 {
-        struct enforcement_timer* et = container_of(timer,
-                                                    struct enforcement_timer,
-                                                    timer);
+        struct enforcement_timer* et;
+        int ret;
         unsigned long flags;
 
-        local_irq_save(flags);
-        TRACE("enforcement timer fired.\n");
-        et->armed = 0;
-        /* activate scheduler */
-        litmus_reschedule_local();
-        local_irq_restore(flags);
-
-        return  HRTIMER_NORESTART;
-}
+        BUG_ON(!t);
+        BUG_ON(!is_realtime(t));
 
-/* assumes called with IRQs off */
-static void cancel_enforcement_timer(struct enforcement_timer* et)
-{
-        int ret;
+        et = &tsk_rt(t)->budget.timer;
 
         TRACE("cancelling enforcement timer.\n");
 
-        /* Since interrupts are disabled and et->armed is only
-         * modified locally, we do not need any locks.
-         */
-
         if (et->armed) {
-                ret = hrtimer_try_to_cancel(&et->timer);
-                /* Should never be inactive. */
-                BUG_ON(ret == 0);
-                /* Should never be running concurrently. */
-                BUG_ON(ret == -1);
-
-                et->armed = 0;
+                raw_spin_lock_irqsave(&et->lock, flags);
+                if (et->armed) {
+                        ret = hrtimer_try_to_cancel(&et->timer);
+                        et->armed = 0;
+                }
+                else {
+                        TRACE("timer was not armed (race).\n");
+                }
+                raw_spin_unlock_irqrestore(&et->lock, flags);
+        }
+        else {
+                TRACE("timer was not armed.\n");
         }
 }
 
-/* assumes called with IRQs off */
-static void arm_enforcement_timer(struct enforcement_timer* et,
-                                  struct task_struct* t)
+inline static void arm_enforcement_timer(struct task_struct* t)
 {
+        struct enforcement_timer* et;
         lt_t when_to_fire;
-        TRACE_TASK(t, "arming enforcement timer.\n");
+        unsigned long flags;
+
+        BUG_ON(!t);
+        BUG_ON(!is_realtime(t));
+
+        et = &tsk_rt(t)->budget.timer;
+        if (et->armed) {
+                TRACE_TASK(t, "timer already armed!\n");
+                return;
+        }
 
         /* Calling this when there is no budget left for the task
          * makes no sense, unless the task is non-preemptive. */
-        BUG_ON(budget_exhausted(t) && !is_np(t));
+        if (budget_exhausted(t)) {
+                TRACE_TASK(t, "can't arm timer because no budget remaining\n");
+                return;
+        }
+
+        if ( (!budget_enforced(t) ||
+                        (budget_enforced(t) && bt_flag_is_set(t, BTF_BUDGET_EXHAUSTED)))
+                        &&
+                (!budget_signalled(t) ||
+                        (budget_signalled(t) && bt_flag_is_set(t, BTF_SIG_BUDGET_SENT)))) {
+                TRACE_TASK(t, "trying to arm timer when budget has already been exhausted.\n");
+                return;
+        }
+
+        TRACE_TASK(t, "arming enforcement timer.\n");
 
         /* __hrtimer_start_range_ns() cancels the timer
          * anyway, so we don't have to check whether it is still armed */
+        raw_spin_lock_irqsave(&et->lock, flags);
 
-        if (likely(!is_np(t))) {
-                when_to_fire = litmus_clock() + budget_remaining(t);
-                __hrtimer_start_range_ns(&et->timer,
-                                         ns_to_ktime(when_to_fire),
-                                         0 /* delta */,
-                                         HRTIMER_MODE_ABS_PINNED,
-                                         0 /* no wakeup */);
-                et->armed = 1;
+        if (et->armed) {
+                TRACE_TASK(t, "timer already armed (race)!\n");
+                goto out;
         }
-}
 
+        when_to_fire = litmus_clock() + budget_remaining(t);
 
-/* expects to be called with IRQs off */
-void update_enforcement_timer(struct task_struct* t)
-{
-        struct enforcement_timer* et = &__get_cpu_var(budget_timer);
-
-        if (t && budget_precisely_tracked(t) && !sigbudget_sent(t)) {
-                /* Make sure we call into the scheduler when this budget
-                 * expires. */
-                arm_enforcement_timer(et, t);
-        } else if (et->armed) {
-                /* Make sure we don't cause unnecessary interrupts. */
-                cancel_enforcement_timer(et);
-        }
+        TRACE_TASK(t, "bremaining: %ld, when_to_fire: %ld\n", budget_remaining(t), when_to_fire);
+
+        __hrtimer_start_range_ns(&et->timer,
+                                 ns_to_ktime(when_to_fire),
+                                 0 /* delta */,
+                                 HRTIMER_MODE_ABS_PINNED,  // TODO: need to use non-pinned?
+                                 0 /* no wakeup */);
+        et->armed = 1;
+
+out:
+        raw_spin_unlock_irqrestore(&et->lock, flags);
 }
 
+
 void send_sigbudget(struct task_struct* t)
 {
-        if (!test_and_set_bit(RT_JOB_SIG_BUDGET_SENT, &tsk_rt(t)->job_params.flags)) {
+        if (!bt_flag_test_and_set(t, BTF_SIG_BUDGET_SENT)) {
                 /* signal has not yet been sent and we are responsible for sending
                  * since we just set the sent-bit when it was previously 0. */
 
@@ -109,17 +107,86 @@ void send_sigbudget(struct task_struct* t)
         }
 }
 
-static int __init init_budget_enforcement(void)
+
+void simple_on_scheduled(struct task_struct* t)
 {
-        int cpu;
-        struct enforcement_timer* et;
+        BUG_ON(!t);
+
+        if (budget_precisely_tracked(t) && !bt_flag_is_set(t, BTF_SIG_BUDGET_SENT)) {
+                BUG_ON(tsk_rt(t)->budget.timer.armed);
+                arm_enforcement_timer(t);
+        }
+}
+
+static void __simple_on_unscheduled(struct task_struct* t)
+{
+        BUG_ON(!t);
 
-        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;
+        if (budget_precisely_tracked(t)) {
+                cancel_enforcement_timer(t);
         }
-        return 0;
 }
 
-module_init(init_budget_enforcement);
+void simple_on_blocked(struct task_struct* t)
+{
+        __simple_on_unscheduled(t);
+}
+
+void simple_on_preempt_or_sleep(struct task_struct* t)
+{
+        __simple_on_unscheduled(t);
+}
+
+void simple_on_exit(struct task_struct* t)
+{
+        __simple_on_unscheduled(t);
+}
+
+
+
+
+static enum hrtimer_restart __on_timeout(struct hrtimer *timer)
+{
+        unsigned long flags;
+        struct budget_tracker* bt =
+                container_of(
+                        container_of(timer,
+                                                struct enforcement_timer,
+                                                timer),
+                        struct budget_tracker,
+                        timer);
+
+        struct task_struct* t =
+                container_of(
+                        container_of(bt, struct rt_param, budget),
+                        struct task_struct,
+                        rt_param);
+
+        TRACE_TASK(t, "budget timer interrupt fired at time %lu\n", litmus_clock());
+
+        raw_spin_lock_irqsave(&bt->timer.lock, flags);
+        tsk_rt(t)->budget.timer.armed = 0;
+        raw_spin_unlock_irqrestore(&bt->timer.lock, flags);
+
+        bt->ops->on_exhausted(t);
+
+        return  HRTIMER_NORESTART;
+}
+
+
+void init_budget_tracker(struct budget_tracker* bt, const struct budget_tracker_ops* ops)
+{
+        BUG_ON(!bt);
+        BUG_ON(!ops);
+
+        BUG_ON(!ops->on_scheduled);
+        BUG_ON(!ops->on_blocked);
+        BUG_ON(!ops->on_preempt_or_sleep);
+        BUG_ON(!ops->on_exhausted);
+
+        memset(bt, 0, sizeof(*bt));
+        raw_spin_lock_init(&bt->timer.lock);
+        hrtimer_init(&bt->timer.timer, CLOCK_MONOTONIC, HRTIMER_MODE_ABS);
+        bt->timer.timer.function = __on_timeout;
+        bt->ops = ops;
+}
\ No newline at end of file
diff --git a/litmus/jobs.c b/litmus/jobs.c
index 991c6e60be74..1479cddad9de 100644
--- a/litmus/jobs.c
+++ b/litmus/jobs.c
@@ -13,7 +13,9 @@ static inline void setup_release(struct task_struct *t, lt_t release)
         t->rt_param.job_params.deadline = release + get_rt_relative_deadline(t);
         t->rt_param.job_params.exec_time = 0;
 
-        clear_bit(RT_JOB_SIG_BUDGET_SENT, &t->rt_param.job_params.flags);
+        /* kludge - TODO: Move this to budget.h/.c */
+        if (t->rt_param.budget.ops)
+                bt_flags_reset(t);
 
         /* update job sequence number */
         t->rt_param.job_params.job_no++;
diff --git a/litmus/litmus.c b/litmus/litmus.c
index 97cbe0461a93..e8130e362c84 100644
--- a/litmus/litmus.c
+++ b/litmus/litmus.c
@@ -19,6 +19,8 @@
 #include <litmus/litmus_proc.h>
 #include <litmus/sched_trace.h>
 
+#include <litmus/budget.h>
+
 #ifdef CONFIG_SCHED_CPU_AFFINITY
 #include <litmus/affinity.h>
 #endif
diff --git a/litmus/sched_cedf.c b/litmus/sched_cedf.c
index 78ae5a080138..8fe646f1f0c5 100644
--- a/litmus/sched_cedf.c
+++ b/litmus/sched_cedf.c
@@ -444,6 +444,35 @@ static noinline void job_completion(struct task_struct *t, int forced)
         }
 }
 
+static void cedf_simple_on_exhausted(struct task_struct *t)
+{
+        /* Assumption: t is scheduled on the CPU executing this callback */
+
+        if (budget_signalled(t) && !bt_flag_is_set(t, BTF_SIG_BUDGET_SENT)) {
+                /* signal exhaustion */
+                send_sigbudget(t); /* will set BTF_SIG_BUDGET_SENT */
+        }
+
+        if (budget_enforced(t) && !bt_flag_test_and_set(t, BTF_BUDGET_EXHAUSTED)) {
+                if (!is_np(t)) {
+                        /* np tasks will be preempted when they become
+                         * preemptable again
+                         */
+                        litmus_reschedule_local();
+                        set_will_schedule();
+                        TRACE("cedf_scheduler_tick: "
+                                  "%d is preemptable "
+                                  " => FORCE_RESCHED\n", t->pid);
+                } else if (is_user_np(t)) {
+                        TRACE("cedf_scheduler_tick: "
+                                  "%d is non-preemptable, "
+                                  "preemption delayed.\n", t->pid);
+                        request_exit_np(t);
+                }
+        }
+}
+
+
 /* cedf_tick - this function is called for every local timer
  *                         interrupt.
  *
@@ -452,48 +481,18 @@ static noinline void job_completion(struct task_struct *t, int forced)
  */
 static void cedf_tick(struct task_struct* t)
 {
-        if (is_realtime(t) && budget_exhausted(t))
-        {
-                if (budget_signalled(t) && !sigbudget_sent(t)) {
-                        /* signal exhaustion */
-                        send_sigbudget(t);
-                }
-
-                if (budget_enforced(t)) {
-                        if (!is_np(t)) {
-                                /* np tasks will be preempted when they become
-                                 * preemptable again
-                                 */
-                                litmus_reschedule_local();
-                                set_will_schedule();
-                                TRACE("cedf_scheduler_tick: "
-                                          "%d is preemptable "
-                                          " => FORCE_RESCHED\n", t->pid);
-                        } else if (is_user_np(t)) {
-                                TRACE("cedf_scheduler_tick: "
-                                          "%d is non-preemptable, "
-                                          "preemption delayed.\n", t->pid);
-                                request_exit_np(t);
-                        }
-                }
+        if (is_realtime(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);
         }
 }
 
 
 
-
-
-
-
-
-
-
-
-
-
 #ifdef CONFIG_LITMUS_PAI_SOFTIRQD
 
-
 static void __do_lit_tasklet(struct tasklet_struct* tasklet, unsigned long flushed)
 {
         if (!atomic_read(&tasklet->count)) {
@@ -787,6 +786,7 @@ static void cedf_change_prio_pai_tasklet(struct task_struct *old_prio,
 
 #endif  // PAI
 
+
 /* Getting schedule() right is a bit tricky. schedule() may not make any
  * assumptions on the state of the current task since it may be called for a
  * number of reasons. The reasons include a scheduler_tick() determined that it
@@ -812,7 +812,7 @@ static struct task_struct* cedf_schedule(struct task_struct * prev)
 {
         cpu_entry_t* entry = &__get_cpu_var(cedf_cpu_entries);
         cedf_domain_t *cluster = entry->cluster;
-        int out_of_time, signal_budget, sleep, preempt, np, exists, blocks;
+        int out_of_time, sleep, preempt, np, exists, blocks;
         struct task_struct* next = NULL;
 
 #ifdef CONFIG_RELEASE_MASTER
@@ -838,11 +838,7 @@ static struct task_struct* cedf_schedule(struct task_struct * prev)
         blocks      = exists && !is_running(entry->scheduled);
         out_of_time = exists &&
                                   budget_enforced(entry->scheduled) &&
-                                  budget_exhausted(entry->scheduled);
-        signal_budget = exists &&
-                                        budget_signalled(entry->scheduled) &&
-                                        budget_exhausted(entry->scheduled) &&
-                                        !sigbudget_sent(entry->scheduled);
+                                  bt_flag_is_set(entry->scheduled, BTF_BUDGET_EXHAUSTED);
         np          = exists && is_np(entry->scheduled);
         sleep       = exists && is_completed(entry->scheduled);
         preempt     = entry->scheduled != entry->linked;
@@ -851,12 +847,13 @@ static struct task_struct* cedf_schedule(struct task_struct * prev)
         TRACE_TASK(prev, "invoked cedf_schedule.\n");
 #endif
 
-        if (exists)
+        if (exists) {
                 TRACE_TASK(prev,
                            "blocks:%d out_of_time:%d np:%d sleep:%d preempt:%d "
                            "state:%d sig:%d\n",
                            blocks, out_of_time, np, sleep, preempt,
                            prev->state, signal_pending(prev));
+        }
         if (entry->linked && preempt)
                 TRACE_TASK(prev, "will be preempted by %s/%d\n",
                            entry->linked->comm, entry->linked->pid);
@@ -872,9 +869,13 @@ static struct task_struct* cedf_schedule(struct task_struct * prev)
         }
 #endif
 
-        /* Send the signal that the budget has been exhausted */
-        if (signal_budget)
-                send_sigbudget(entry->scheduled);
+        /* Do budget stuff */
+        if (tsk_rt(prev)->budget.ops) {
+                if (blocks)
+                        tsk_rt(prev)->budget.ops->on_blocked(prev);
+                else if (preempt || sleep)
+                        tsk_rt(prev)->budget.ops->on_preempt_or_sleep(prev);
+        }
 
         /* If a task blocks we have no choice but to reschedule.
          */
@@ -942,6 +943,7 @@ static struct task_struct* cedf_schedule(struct task_struct * prev)
                 }
         }
 
+
 #ifdef CONFIG_REALTIME_AUX_TASKS
 out_set_state:
 #endif
@@ -1111,6 +1113,10 @@ static void cedf_task_exit(struct task_struct * t)
         /* unlink if necessary */
         raw_spin_lock_irqsave(&cluster->cluster_lock, flags);
 
+        /* disable budget enforcement */
+        if (tsk_rt(t)->budget.ops)
+                tsk_rt(t)->budget.ops->on_exit(t);
+
 #ifdef CONFIG_REALTIME_AUX_TASKS
         /* make sure we clean up on our way out */
         if (unlikely(tsk_rt(t)->is_aux_task)) {
@@ -1141,15 +1147,44 @@ static void cedf_task_exit(struct task_struct * t)
         TRACE_TASK(t, "RIP\n");
 }
 
+
+
+
+
+
+static struct budget_tracker_ops cedf_drain_simple_ops =
+{
+        .on_scheduled = simple_on_scheduled,
+        .on_blocked = simple_on_blocked,
+        .on_preempt_or_sleep = simple_on_preempt_or_sleep,
+        .on_exit = simple_on_exit,
+
+        .on_exhausted = cedf_simple_on_exhausted,
+};
+
+
 static long cedf_admit_task(struct task_struct* tsk)
 {
+        if (remote_cluster(task_cpu(tsk)) != task_cpu_cluster(tsk))
+                return -EINVAL;
+
+        if (budget_enforced(tsk) || budget_signalled(tsk)) {
+                switch(get_drain_policy(tsk)) {
+                case DRAIN_SIMPLE:
+                        init_budget_tracker(&tsk_rt(tsk)->budget, &cedf_drain_simple_ops);
+                        break;
+                default:
+                        TRACE_TASK(tsk, "Unsupported budget draining mode.\n");
+                        return -EINVAL;
+                }
+        }
+
 #ifdef CONFIG_LITMUS_NESTED_LOCKING
         INIT_BINHEAP_HANDLE(&tsk_rt(tsk)->hp_blocked_tasks,
                                                 edf_max_heap_base_priority_order);
 #endif
 
-        return (remote_cluster(task_cpu(tsk)) == task_cpu_cluster(tsk)) ?
-                                0 : -EINVAL;
+        return 0;
 }
 
 
diff --git a/litmus/sched_gsn_edf.c b/litmus/sched_gsn_edf.c
index 08cdf5c0e492..15ac94038702 100644
--- a/litmus/sched_gsn_edf.c
+++ b/litmus/sched_gsn_edf.c
@@ -421,6 +421,33 @@ static noinline void job_completion(struct task_struct *t, int forced)
                 gsnedf_job_arrival(t);
 }
 
+static void gsnedf_simple_on_exhausted(struct task_struct *t)
+{
+        /* Assumption: t is scheduled on the CPU executing this callback */
+
+        if (budget_signalled(t) && !bt_flag_is_set(t, BTF_SIG_BUDGET_SENT)) {
+                /* signal exhaustion */
+                send_sigbudget(t); /* will set BTF_SIG_BUDGET_SENT */
+        }
+
+        if (budget_enforced(t) && !bt_flag_test_and_set(t, BTF_BUDGET_EXHAUSTED)) {
+                if (!is_np(t)) {
+                        /* np tasks will be preempted when they become
+                         * preemptable again
+                         */
+                        litmus_reschedule_local();
+                        TRACE("cedf_scheduler_tick: "
+                                  "%d is preemptable "
+                                  " => FORCE_RESCHED\n", t->pid);
+                } else if (is_user_np(t)) {
+                        TRACE("cedf_scheduler_tick: "
+                                  "%d is non-preemptable, "
+                                  "preemption delayed.\n", t->pid);
+                        request_exit_np(t);
+                }
+        }
+}
+
 /* gsnedf_tick - this function is called for every local timer
  *                         interrupt.
  *
@@ -429,41 +456,16 @@ static noinline void job_completion(struct task_struct *t, int forced)
  */
 static void gsnedf_tick(struct task_struct* t)
 {
-        if (is_realtime(t) && budget_exhausted(t))
-        {
-                if (budget_signalled(t) && !sigbudget_sent(t)) {
-                        /* signal exhaustion */
-                        send_sigbudget(t);
-                }
-
-                if (budget_enforced(t)) {
-                        if (!is_np(t)) {
-                                /* np tasks will be preempted when they become
-                                 * preemptable again
-                                 */
-                                litmus_reschedule_local();
-                                TRACE("gsnedf_scheduler_tick: "
-                                          "%d is preemptable "
-                                          " => FORCE_RESCHED\n", t->pid);
-                        } else if (is_user_np(t)) {
-                                TRACE("gsnedf_scheduler_tick: "
-                                          "%d is non-preemptable, "
-                                          "preemption delayed.\n", t->pid);
-                                request_exit_np(t);
-                        }
-                }
-        }
-
-        /*
-        if(is_realtime(t)) {
-                TRACE_TASK(t, "tick %llu\n", litmus_clock());
+        if (is_realtime(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);
         }
-         */
 }
 
 
 
-
 #ifdef CONFIG_LITMUS_PAI_SOFTIRQD
 
 
@@ -797,11 +799,9 @@ static void gsnedf_change_prio_pai_tasklet(struct task_struct *old_prio,
 static struct task_struct* gsnedf_schedule(struct task_struct * prev)
 {
         cpu_entry_t* entry = &__get_cpu_var(gsnedf_cpu_entries);
-        int out_of_time, signal_budget, sleep, preempt, np, exists, blocks;
+        int out_of_time, sleep, preempt, np, exists, blocks;
         struct task_struct* next = NULL;
 
-        //int completion = 0;
-
 #ifdef CONFIG_RELEASE_MASTER
         /* Bail out early if we are the release master.
          * The release master never schedules any real-time tasks.
@@ -824,11 +824,7 @@ static struct task_struct* gsnedf_schedule(struct task_struct * prev)
         blocks      = exists && !is_running(entry->scheduled);
         out_of_time = exists &&
                 budget_enforced(entry->scheduled) &&
-                budget_exhausted(entry->scheduled);
-        signal_budget = exists &&
-                budget_signalled(entry->scheduled) &&
-                budget_exhausted(entry->scheduled) &&
-                !sigbudget_sent(entry->scheduled);
+                bt_flag_is_set(entry->scheduled, BTF_BUDGET_EXHAUSTED);
         np          = exists && is_np(entry->scheduled);
         sleep       = exists && is_completed(entry->scheduled);
         preempt     = entry->scheduled != entry->linked;
@@ -839,9 +835,9 @@ static struct task_struct* gsnedf_schedule(struct task_struct * prev)
 
         if (exists) {
                 TRACE_TASK(prev,
-                           "blocks:%d out_of_time:%d signal_budget: %d np:%d sleep:%d preempt:%d "
+                           "blocks:%d out_of_time:%d np:%d sleep:%d preempt:%d "
                            "state:%d sig:%d\n",
-                           blocks, out_of_time, signal_budget, np, sleep, preempt,
+                           blocks, out_of_time, np, sleep, preempt,
                            prev->state, signal_pending(prev));
         }
 
@@ -849,9 +845,12 @@ static struct task_struct* gsnedf_schedule(struct task_struct * prev)
                 TRACE_TASK(prev, "will be preempted by %s/%d\n",
                            entry->linked->comm, entry->linked->pid);
 
-        /* Send the signal that the budget has been exhausted */
-        if (signal_budget) {
-                send_sigbudget(entry->scheduled);
+        /* Do budget stuff */
+        if (tsk_rt(prev)->budget.ops) {
+                if (blocks)
+                        tsk_rt(prev)->budget.ops->on_blocked(prev);
+                else if (preempt || sleep)
+                        tsk_rt(prev)->budget.ops->on_preempt_or_sleep(prev);
         }
 
         /* If a task blocks we have no choice but to reschedule.
@@ -1086,6 +1085,10 @@ static void gsnedf_task_exit(struct task_struct * t)
         /* unlink if necessary */
         raw_spin_lock_irqsave(&gsnedf_lock, flags);
 
+        /* disable budget enforcement */
+        if (tsk_rt(t)->budget.ops)
+                tsk_rt(t)->budget.ops->on_exit(t);
+
 #ifdef CONFIG_REALTIME_AUX_TASKS
         /* make sure we clean up on our way out */
         if (unlikely(tsk_rt(t)->is_aux_task)) {
@@ -1115,8 +1118,29 @@ static void gsnedf_task_exit(struct task_struct * t)
 }
 
 
+static struct budget_tracker_ops gsnedf_drain_simple_ops =
+{
+        .on_scheduled = simple_on_scheduled,
+        .on_blocked = simple_on_blocked,
+        .on_preempt_or_sleep = simple_on_preempt_or_sleep,
+        .on_exit = simple_on_exit,
+
+        .on_exhausted = gsnedf_simple_on_exhausted,
+};
+
 static long gsnedf_admit_task(struct task_struct* tsk)
 {
+        if (budget_enforced(tsk) || budget_signalled(tsk)) {
+                switch(get_drain_policy(tsk)) {
+                        case DRAIN_SIMPLE:
+                                init_budget_tracker(&tsk_rt(tsk)->budget, &gsnedf_drain_simple_ops);
+                                break;
+                        default:
+                                TRACE_TASK(tsk, "Unsupported budget draining mode.\n");
+                                return -EINVAL;
+                }
+        }
+
 #ifdef CONFIG_LITMUS_NESTED_LOCKING
         INIT_BINHEAP_HANDLE(&tsk_rt(tsk)->hp_blocked_tasks,
                                                 edf_max_heap_base_priority_order);
@@ -1126,10 +1150,6 @@ static long gsnedf_admit_task(struct task_struct* tsk)
 }
 
 
-
-
-
-
 #ifdef CONFIG_LITMUS_LOCKING
 
 #include <litmus/fdso.h>
diff --git a/litmus/sched_litmus.c b/litmus/sched_litmus.c
index 9de03c95b825..60b58bb29ac4 100644
--- a/litmus/sched_litmus.c
+++ b/litmus/sched_litmus.c
@@ -51,7 +51,7 @@ litmus_schedule(struct rq *rq, struct task_struct *prev)
         lt_t _maybe_deadlock = 0;
 
         /* let the plugin schedule */
-        next = litmus->schedule(prev);
+        next = litmus->schedule(prev); /* may disable prev's budget timer */
 
         sched_state_plugin_check();
 
@@ -150,9 +150,11 @@ litmus_schedule(struct rq *rq, struct task_struct *prev)
         if (next) {
                 next->rt_param.stack_in_use = rq->cpu;
                 next->se.exec_start = rq->clock;
+
+                if (is_realtime(next) && tsk_rt(next)->budget.ops)
+                        tsk_rt(next)->budget.ops->on_scheduled(next);
         }
 
-        update_enforcement_timer(next);
         return next;
 }
 
diff --git a/litmus/sched_pfair.c b/litmus/sched_pfair.c
index d5fb3a832adc..c06326faf9ce 100644
--- a/litmus/sched_pfair.c
+++ b/litmus/sched_pfair.c
@@ -837,6 +837,16 @@ static void dump_subtasks(struct task_struct* t)
                            t->rt_param.pfair->subtasks[i].group_deadline);
 }
 
+static struct budget_tracker_ops pfair_drain_simple_ops =
+{
+        .on_scheduled = simple_on_scheduled,
+        .on_blocked = simple_on_blocked,
+        .on_preempt_or_sleep = simple_on_preempt_or_sleep,
+        .on_exit = simple_on_exit,
+
+        .on_exhausted = pfair_simple_on_exhausted,
+};
+
 static long pfair_admit_task(struct task_struct* t)
 {
         lt_t quanta;
diff --git a/litmus/sched_pfp.c b/litmus/sched_pfp.c
index 6edec830f063..4a8b8e084f6e 100644
--- a/litmus/sched_pfp.c
+++ b/litmus/sched_pfp.c
@@ -132,6 +132,33 @@ static void job_completion(struct task_struct* t, int forced)
                 sched_trace_task_release(t);
 }
 
+static void pfp_simple_on_exhausted(struct task_struct *t)
+{
+        /* Assumption: t is scheduled on the CPU executing this callback */
+
+        if (budget_signalled(t) && !bt_flag_is_set(t, BTF_SIG_BUDGET_SENT)) {
+                /* signal exhaustion */
+                send_sigbudget(t); /* will set BTF_SIG_BUDGET_SENT */
+        }
+
+        if (budget_enforced(t) && !bt_flag_test_and_set(t, BTF_BUDGET_EXHAUSTED)) {
+                if (!is_np(t)) {
+                        /* np tasks will be preempted when they become
+                         * preemptable again
+                         */
+                        litmus_reschedule_local();
+                        TRACE("cedf_scheduler_tick: "
+                                  "%d is preemptable "
+                                  " => FORCE_RESCHED\n", t->pid);
+                } else if (is_user_np(t)) {
+                        TRACE("cedf_scheduler_tick: "
+                                  "%d is non-preemptable, "
+                                  "preemption delayed.\n", t->pid);
+                        request_exit_np(t);
+                }
+        }
+}
+
 static void pfp_tick(struct task_struct *t)
 {
         pfp_domain_t *pfp = local_pfp;
@@ -142,26 +169,11 @@ static void pfp_tick(struct task_struct *t)
          */
         BUG_ON(is_realtime(t) && t != pfp->scheduled);
 
-        if (is_realtime(t) && budget_exhausted(t))
-        {
-                if (budget_signalled(t) && !sigbudget_sent(t)) {
-                        /* signal exhaustion */
-                        send_sigbudget(t);
-                }
-
-                if (budget_enforced(t)) {
-                        if (!is_np(t)) {
-                                litmus_reschedule_local();
-                                TRACE("pfp_scheduler_tick: "
-                                          "%d is preemptable "
-                                          " => FORCE_RESCHED\n", t->pid);
-                        } else if (is_user_np(t)) {
-                                TRACE("pfp_scheduler_tick: "
-                                          "%d is non-preemptable, "
-                                          "preemption delayed.\n", t->pid);
-                                request_exit_np(t);
-                        }
-                }
+        if (is_realtime(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);
         }
 }
 
@@ -170,7 +182,7 @@ static struct task_struct* pfp_schedule(struct task_struct * prev)
         pfp_domain_t*   pfp = local_pfp;
         struct task_struct*     next;
 
-        int out_of_time, signal_budget, sleep, preempt, np, exists, blocks, resched, migrate;
+        int out_of_time, sleep, preempt, np, exists, blocks, resched, migrate;
 
         raw_spin_lock(&pfp->slock);
 
@@ -186,11 +198,7 @@ static struct task_struct* pfp_schedule(struct task_struct * prev)
         blocks      = exists && !is_running(pfp->scheduled);
         out_of_time = exists &&
                                   budget_enforced(pfp->scheduled) &&
-                                  budget_exhausted(pfp->scheduled);
-        signal_budget = exists &&
-                                        budget_signalled(pfp->scheduled) &&
-                                        budget_exhausted(pfp->scheduled) &&
-                                        !sigbudget_sent(pfp->scheduled);
+                                  bt_flag_is_set(pfp->scheduled, BTF_BUDGET_EXHAUSTED);
         np          = exists && is_np(pfp->scheduled);
         sleep       = exists && is_completed(pfp->scheduled);
         migrate     = exists && get_partition(pfp->scheduled) != pfp->cpu;
@@ -202,9 +210,13 @@ static struct task_struct* pfp_schedule(struct task_struct * prev)
          */
         resched = preempt;
 
-        /* Send the signal that the budget has been exhausted */
-        if (signal_budget)
-                send_sigbudget(pfp->scheduled);
+        /* Do budget stuff */
+        if (tsk_rt(prev)->budget.ops) {
+                if (blocks)
+                        tsk_rt(prev)->budget.ops->on_blocked(prev);
+                else if (preempt || sleep)
+                        tsk_rt(prev)->budget.ops->on_preempt_or_sleep(prev);
+        }
 
         /* If a task blocks we have no choice but to reschedule.
          */
@@ -418,6 +430,11 @@ static void pfp_task_exit(struct task_struct * t)
         rt_domain_t*            dom;
 
         raw_spin_lock_irqsave(&pfp->slock, flags);
+
+        /* disable budget enforcement */
+        if (tsk_rt(t)->budget.ops)
+                tsk_rt(t)->budget.ops->on_exit(t);
+
         if (is_queued(t)) {
                 BUG(); /* This currently doesn't work. */
                 /* dequeue */
@@ -1692,17 +1709,43 @@ static long pfp_allocate_lock(struct litmus_lock **lock, int type,
 
 #endif
 
+static struct budget_tracker_ops pfp_drain_simple_ops =
+{
+        .on_scheduled = simple_on_scheduled,
+        .on_blocked = simple_on_blocked,
+        .on_preempt_or_sleep = simple_on_preempt_or_sleep,
+        .on_exit = simple_on_exit,
+
+        .on_exhausted = pfp_simple_on_exhausted,
+};
+
 static long pfp_admit_task(struct task_struct* tsk)
 {
+        long ret = 0;
+
         if (task_cpu(tsk) == tsk->rt_param.task_params.cpu &&
 #ifdef CONFIG_RELEASE_MASTER
             /* don't allow tasks on release master CPU */
             task_cpu(tsk) != remote_dom(task_cpu(tsk))->release_master &&
 #endif
-            litmus_is_valid_fixed_prio(get_priority(tsk)))
-                return 0;
+            litmus_is_valid_fixed_prio(get_priority(tsk))) {
+
+                if (budget_enforced(tsk) || budget_signalled(tsk)) {
+                        switch(get_drain_policy(tsk)) {
+                                case DRAIN_SIMPLE:
+                                        init_budget_tracker(&tsk_rt(tsk)->budget, &pfp_drain_simple_ops);
+                                        break;
+                                default:
+                                        TRACE_TASK(tsk, "Unsupported budget draining mode.\n");
+                                        ret = -EINVAL;
+                                        break;
+                        }
+                }
+        }
         else
-                return -EINVAL;
+                ret = -EINVAL;
+
+        return ret;
 }
 
 static long pfp_activate_plugin(void)
diff --git a/litmus/sched_psn_edf.c b/litmus/sched_psn_edf.c
index d1177cab152d..3b3edfe908ff 100644
--- a/litmus/sched_psn_edf.c
+++ b/litmus/sched_psn_edf.c
@@ -164,6 +164,33 @@ static void job_completion(struct task_struct* t, int forced)
         prepare_for_next_period(t);
 }
 
+static void psnedf_simple_on_exhausted(struct task_struct *t)
+{
+        /* Assumption: t is scheduled on the CPU executing this callback */
+
+        if (budget_signalled(t) && !bt_flag_is_set(t, BTF_SIG_BUDGET_SENT)) {
+                /* signal exhaustion */
+                send_sigbudget(t); /* will set BTF_SIG_BUDGET_SENT */
+        }
+
+        if (budget_enforced(t) && !bt_flag_test_and_set(t, BTF_BUDGET_EXHAUSTED)) {
+                if (!is_np(t)) {
+                        /* np tasks will be preempted when they become
+                         * preemptable again
+                         */
+                        litmus_reschedule_local();
+                        TRACE("cedf_scheduler_tick: "
+                                  "%d is preemptable "
+                                  " => FORCE_RESCHED\n", t->pid);
+                } else if (is_user_np(t)) {
+                        TRACE("cedf_scheduler_tick: "
+                                  "%d is non-preemptable, "
+                                  "preemption delayed.\n", t->pid);
+                        request_exit_np(t);
+                }
+        }
+}
+
 static void psnedf_tick(struct task_struct *t)
 {
         psnedf_domain_t *pedf = local_pedf;
@@ -174,26 +201,11 @@ static void psnedf_tick(struct task_struct *t)
          */
         BUG_ON(is_realtime(t) && t != pedf->scheduled);
 
-        if (is_realtime(t) && budget_exhausted(t))
-        {
-                if (budget_signalled(t) && !sigbudget_sent(t)) {
-                        /* signal exhaustion */
-                        send_sigbudget(t);
-                }
-
-                if (budget_enforced(t)) {
-                        if (!is_np(t)) {
-                                litmus_reschedule_local();
-                                TRACE("psnedf_scheduler_tick: "
-                                          "%d is preemptable "
-                                          " => FORCE_RESCHED\n", t->pid);
-                        } else if (is_user_np(t)) {
-                                TRACE("psnedf_scheduler_tick: "
-                                          "%d is non-preemptable, "
-                                          "preemption delayed.\n", t->pid);
-                                request_exit_np(t);
-                        }
-                }
+        if (is_realtime(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);
         }
 }
 
@@ -203,7 +215,7 @@ static struct task_struct* psnedf_schedule(struct task_struct * prev)
         rt_domain_t*            edf  = &pedf->domain;
         struct task_struct*     next;
 
-        int out_of_time, signal_budget, sleep, preempt, np, exists, blocks, resched;
+        int out_of_time, sleep, preempt, np, exists, blocks, resched;
 
         raw_spin_lock(&pedf->slock);
 
@@ -219,11 +231,7 @@ static struct task_struct* psnedf_schedule(struct task_struct * prev)
         blocks      = exists && !is_running(pedf->scheduled);
         out_of_time = exists &&
                                   budget_enforced(pedf->scheduled) &&
-                                  budget_exhausted(pedf->scheduled);
-        signal_budget = exists &&
-                                        budget_signalled(pedf->scheduled) &&
-                                        budget_exhausted(pedf->scheduled) &&
-                                        !sigbudget_sent(pedf->scheduled);
+                                  bt_flag_is_set(pedf->scheduled, BTF_BUDGET_EXHAUSTED);
         np          = exists && is_np(pedf->scheduled);
         sleep       = exists && is_completed(pedf->scheduled);
         preempt     = edf_preemption_needed(edf, prev);
@@ -234,9 +242,13 @@ static struct task_struct* psnedf_schedule(struct task_struct * prev)
          */
         resched = preempt;
 
-        /* Send the signal that the budget has been exhausted */
-        if (signal_budget)
-                send_sigbudget(pedf->scheduled);
+        /* Do budget stuff */
+        if (tsk_rt(prev)->budget.ops) {
+                if (blocks)
+                        tsk_rt(prev)->budget.ops->on_blocked(prev);
+                else if (preempt || sleep)
+                        tsk_rt(prev)->budget.ops->on_preempt_or_sleep(prev);
+        }
 
         /* If a task blocks we have no choice but to reschedule.
          */
@@ -380,6 +392,11 @@ static void psnedf_task_exit(struct task_struct * t)
         rt_domain_t*            edf;
 
         raw_spin_lock_irqsave(&pedf->slock, flags);
+
+        /* disable budget enforcement */
+        if (tsk_rt(t)->budget.ops)
+                tsk_rt(t)->budget.ops->on_exit(t);
+
         if (is_queued(t)) {
                 /* dequeue */
                 edf  = task_edf(t);
@@ -626,17 +643,43 @@ static long psnedf_activate_plugin(void)
         return 0;
 }
 
+static struct budget_tracker_ops psnedf_drain_simple_ops =
+{
+        .on_scheduled = simple_on_scheduled,
+        .on_blocked = simple_on_blocked,
+        .on_preempt_or_sleep = simple_on_preempt_or_sleep,
+        .on_exit = simple_on_exit,
+
+        .on_exhausted = psnedf_simple_on_exhausted,
+};
+
 static long psnedf_admit_task(struct task_struct* tsk)
 {
+        long ret = 0;
+
         if (task_cpu(tsk) == tsk->rt_param.task_params.cpu
 #ifdef CONFIG_RELEASE_MASTER
             /* don't allow tasks on release master CPU */
              && task_cpu(tsk) != remote_edf(task_cpu(tsk))->release_master
 #endif
-                )
-                return 0;
+                ) {
+
+                if (budget_enforced(tsk) || budget_signalled(tsk)) {
+                        switch(get_drain_policy(tsk)) {
+                                case DRAIN_SIMPLE:
+                                        init_budget_tracker(&tsk_rt(tsk)->budget, &psnedf_drain_simple_ops);
+                                        break;
+                                default:
+                                        TRACE_TASK(tsk, "Unsupported budget draining mode.\n");
+                                        ret = -EINVAL;
+                                        break;
+                        }
+                }
+        }
         else
-                return -EINVAL;
+                ret = -EINVAL;
+
+        return ret;
 }
 
 /*      Plugin object   */