2 files changed, 11 insertions, 0 deletions

diff --git a/include/litmus.h b/include/litmus.h
index 0d0b31c..0371a6c 100644
--- a/include/litmus.h
+++ b/include/litmus.h
@@ -105,6 +105,7 @@ task_class_t str2class(const char* str);
 /* non-preemptive section support */
 void enter_np(void);
 void exit_np(void);
+void exit_np_trace(void);
 
 /* task system support */
 int wait_for_ts_release(void);
diff --git a/src/kernel_iface.c b/src/kernel_iface.c
index 69e306b..34767ae 100644
--- a/src/kernel_iface.c
+++ b/src/kernel_iface.c
@@ -68,6 +68,15 @@ void exit_np(void)
         if (likely(ctrl_page != NULL) && --ctrl_page->sched.np.flag) {
                 /* became preemptive, let's check for delayed preemptions */
                 __sync_synchronize();
+                if (ctrl_page->sched.np.preempt)
+                        sched_yield();
+        }
+}
+
+void exit_np_trace(void)
+{
+        if (likely(ctrl_page != NULL) && --ctrl_page->sched.np.flag) {
+                __sync_synchronize();
                 if (ctrl_page->sched.np.preempt) {
                         ctrl_page->ts_syscall_start = (uint64_t) get_cycles();
                         sched_yield();
@@ -75,6 +84,7 @@ void exit_np(void)
         }
 }
 
+
 /* init and return a ptr to the control page for
  * preemption and migration overhead analysis
  *