2 files changed, 10 insertions, 4 deletions

diff --git a/include/litmus.h b/include/litmus.h
index 7e5092b..2e48e36 100644
--- a/include/litmus.h
+++ b/include/litmus.h
@@ -105,7 +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);
+int exit_np_trace(void);
 int  requested_to_preempt(void);
 
 /* task system support */
diff --git a/src/kernel_iface.c b/src/kernel_iface.c
index 8ab42de..4ded334 100644
--- a/src/kernel_iface.c
+++ b/src/kernel_iface.c
@@ -65,7 +65,9 @@ void enter_np(void)
 
 void exit_np(void)
 {
-        if (likely(ctrl_page != NULL) && --ctrl_page->sched.np.flag) {
+        if (likely(ctrl_page != NULL) &&
+            ctrl_page->sched.np.flag &&
+            !(--ctrl_page->sched.np.flag)) {
                 /* became preemptive, let's check for delayed preemptions */
                 __sync_synchronize();
                 if (ctrl_page->sched.np.preempt)
@@ -73,15 +75,19 @@ void exit_np(void)
         }
 }
 
-void exit_np_trace(void)
+int exit_np_trace(void)
 {
-        if (likely(ctrl_page != NULL) && --ctrl_page->sched.np.flag) {
+        if (likely(ctrl_page != NULL) &&
+            ctrl_page->sched.np.flag &&
+            !(--ctrl_page->sched.np.flag)) {
                 __sync_synchronize();
                 if (ctrl_page->sched.np.preempt) {
                         ctrl_page->ts_syscall_start = (uint64_t) get_cycles();
                         sched_yield();
+                        return 1;
                 }
         }
+        return 0;
 }
 
 int requested_to_preempt(void)