4 files changed, 9 insertions, 5 deletions

diff --git a/include/litmus/litmus.h b/include/litmus/litmus.h
index a263bc68a7b5..28fb7b331cf3 100644
--- a/include/litmus/litmus.h
+++ b/include/litmus/litmus.h
@@ -281,4 +281,8 @@ static inline quanta_t time2quanta(lt_t time, enum round round)
 /* By how much is cpu staggered behind CPU 0? */
 u64 cpu_stagger_offset(int cpu);
 
+#define TS_SYSCALL_IN_START                                             \
+        if (has_control_page(current))                                  \
+                __TS_SYSCALL_IN_START(&get_control_page(current)->ts_syscall_start);
+
 #endif
diff --git a/include/litmus/trace.h b/include/litmus/trace.h
index 4246260865fa..d6829c416912 100644
--- a/include/litmus/trace.h
+++ b/include/litmus/trace.h
@@ -76,7 +76,7 @@ feather_callback void save_timestamp_time(unsigned long event, unsigned long tim
  * always the next number after the start time event id.
  */
 
-#define TS_SYSCALL_IN_START(p)          TIMESTAMP_TIME(10, p)
+#define __TS_SYSCALL_IN_START(p)        TIMESTAMP_TIME(10, p)
 #define TS_SYSCALL_IN_END               TIMESTAMP_PID(11)
 
 #define TS_SYSCALL_OUT_START            TIMESTAMP_PID(20)
diff --git a/litmus/locking.c b/litmus/locking.c
index 2db2c2813c0b..84a1d8309699 100644
--- a/litmus/locking.c
+++ b/litmus/locking.c
@@ -72,8 +72,7 @@ asmlinkage long sys_litmus_lock(int lock_od)
         struct od_table_entry* entry;
         struct litmus_lock* l;
 
-        if (has_control_page(current))
-                TS_SYSCALL_IN_START(&get_control_page(current)->ts_syscall_start);
+        TS_SYSCALL_IN_START;
 
         TS_SYSCALL_IN_END;
 
@@ -101,8 +100,7 @@ asmlinkage long sys_litmus_unlock(int lock_od)
         struct od_table_entry* entry;
         struct litmus_lock* l;
 
-        if (has_control_page(current))
-                TS_SYSCALL_IN_START(&get_control_page(current)->ts_syscall_start);
+        TS_SYSCALL_IN_START;
 
         TS_SYSCALL_IN_END;
 
diff --git a/litmus/sched_litmus.c b/litmus/sched_litmus.c
index e6952896dc4b..231c86daa1f9 100644
--- a/litmus/sched_litmus.c
+++ b/litmus/sched_litmus.c
@@ -194,6 +194,8 @@ static void dequeue_task_litmus(struct rq *rq, struct task_struct *p,
 
 static void yield_task_litmus(struct rq *rq)
 {
+        TS_SYSCALL_IN_START;
+
         BUG_ON(rq->curr != current);
         /* sched_yield() is called to trigger delayed preemptions.
          * Thus, mark the current task as needing to be rescheduled.