1 files changed, 41 insertions, 5 deletions

diff --git a/include/litmus/budget.h b/include/litmus/budget.h
index 8e426a71f03d..08d5e0970d1d 100644
--- a/include/litmus/budget.h
+++ b/include/litmus/budget.h
@@ -4,6 +4,8 @@
 #include <linux/hrtimer.h>
 #include <linux/semaphore.h>
 
+#include <litmus/binheap.h>
+
 struct enforcement_timer
 {
         raw_spinlock_t lock;
@@ -15,17 +17,22 @@ typedef void (*scheduled_t)(struct task_struct* t);
 typedef void (*blocked_t)(struct task_struct* t);
 typedef void (*preempt_t)(struct task_struct* t);
 typedef void (*sleep_t)(struct task_struct* t);
+typedef void (*wakeup_t)(struct task_struct* t);
 typedef enum hrtimer_restart (*exhausted_t)(struct task_struct* t);
 typedef void (*exit_t)(struct task_struct* t);
 typedef void (*inherit_t)(struct task_struct* t, struct task_struct* prio_inh);
 typedef void (*disinherit_t)(struct task_struct* t, struct task_struct* prio_inh);
 
+typedef void (*enter_top_m_t)(struct task_struct* t);
+typedef void (*exit_top_m_t)(struct task_struct* t);
+
 struct budget_tracker_ops
 {
         scheduled_t                     on_scheduled;   /* called from litmus_schedule(). */
         blocked_t                       on_blocked;             /* called from plugin::schedule() */
         preempt_t                       on_preempt;             /* called from plugin::schedule() */
         sleep_t                         on_sleep;               /* called from plugin::schedule() */
+        wakeup_t                        on_wakeup;
 
         exit_t                          on_exit;                /* task exiting rt mode */
 
@@ -33,6 +40,9 @@ struct budget_tracker_ops
 
         inherit_t                       on_inherit;
         disinherit_t            on_disinherit;
+
+        enter_top_m_t           on_enter_top_m;
+        exit_top_m_t            on_exit_top_m;
 };
 
 struct budget_tracker
@@ -40,13 +50,17 @@ struct budget_tracker
         struct enforcement_timer timer;
         const struct budget_tracker_ops* ops;
         unsigned long flags;
+
+        struct binheap_node top_m_node;
+        lt_t suspend_timestamp;
 };
 
 /* budget tracker flags */
 enum BT_FLAGS
 {
-        BTF_BUDGET_EXHAUSTED = 0,
-        BTF_SIG_BUDGET_SENT     = 1,
+        BTF_BUDGET_EXHAUSTED    = 0,
+        BTF_SIG_BUDGET_SENT             = 1,
+        BTF_IS_TOP_M                    = 2,
 };
 
 /* Functions for simple DRAIN_SIMPLE policy common
@@ -66,16 +80,38 @@ void simple_on_exit(struct task_struct* t);
  *
  * Limitation: Quantum budget tracking is unsupported.
  */
-void sobliv_on_scheduled(struct task_struct* t);
+//void sobliv_on_scheduled(struct task_struct* t);
 void sobliv_on_blocked(struct task_struct* t);
-void sobliv_on_sleep(struct task_struct* t);
+void sobliv_on_wakeup(struct task_struct* t);
+//void sobliv_on_sleep(struct task_struct* t);
+//void sobliv_on_preempt(struct task_struct* t);
 /* Use the DRAIN_SIMPLE implementations */
-#define sobliv_on_preempt       simple_on_preempt
 #define sobliv_on_exit  simple_on_exit
 void sobliv_on_inherit(struct task_struct* t, struct task_struct* prio_inh);
 void sobliv_on_disinherit(struct task_struct* t, struct task_struct* prio_inh);
+void sobliv_on_enter_top_m(struct task_struct* t);
+void sobliv_on_exit_top_m(struct task_struct* t);
+
 void sobliv_revaluate_task(struct task_struct* t);
 
+#define budget_state_machine(t, evt) \
+        do { \
+                if (get_budget_timer(t).ops && \
+                        get_budget_timer(t).ops->evt != NULL) { \
+                        get_budget_timer(t).ops->evt(t); \
+                } \
+        }while(0)
+
+#define budget_state_machine2(a, b, evt) \
+        do { \
+                if (get_budget_timer(a).ops && \
+                        get_budget_timer(b).ops && \
+                        get_budget_timer(a).ops->evt != NULL && \
+                        get_budget_timer(b).ops->evt != NULL) {\
+                        get_budget_timer(a).ops->evt(a, b); \
+                } \
+        }while(0)
+
 
 void init_budget_tracker(struct budget_tracker* bt,
                                 const struct budget_tracker_ops* ops);