[LibOS] Use atomic operations for shim_context.preempt

enable_preempt(), disable_preempt(), and other functions operated on
shim_context.preempt using non-atomic operations. This commit replaces
the old broken implementation with atomic operations.

LibOS/shim/include/shim_internal.h

@@ -200,16 +200,16 @@ void syscall_wrapper_after_syscalldb(void);
 
 void check_stack_hook (void);
 
-static inline uint64_t get_cur_preempt (void) {
+static inline int64_t get_cur_preempt (void) {
     shim_tcb_t* tcb = shim_get_tls();
     assert(tcb);
-    return tcb->context.preempt;
+    return atomic_read(&tcb->context.preempt);
 }
 
 #define BEGIN_SHIM(name, args ...)                          \
     SHIM_ARG_TYPE __shim_##name(args) {                     \
         SHIM_ARG_TYPE ret = 0;                              \
-        uint64_t preempt = get_cur_preempt();               \
+        int64_t preempt = get_cur_preempt();                \
         /* handle_signal(); */                              \
         /* check_stack_hook(); */                           \
         BEGIN_SYSCALL_PROFILE();
@@ -467,13 +467,14 @@ static inline PAL_HANDLE thread_create (void * func, void * arg)
     return DkThreadCreate(func, arg);
 }
 
-static inline void __disable_preempt (shim_tcb_t * tcb)
+static inline int64_t __disable_preempt (shim_tcb_t * tcb)
 {
     //tcb->context.syscall_nr += SYSCALL_NR_PREEMPT_INC;
-    tcb->context.preempt++;
+    int64_t preempt = atomic_inc_return(&tcb->context.preempt);
     /* Assert if this counter overflows */
-    assert(tcb->context.preempt != 0);
-    //debug("disable preempt: %d\n", tcb->context.preempt);
+    assert(preempt != 0);
+    //debug("disable preempt: %d\n", preempt);
+    return preempt;
 }
 
 static inline void disable_preempt (shim_tcb_t * tcb)
@@ -486,11 +487,10 @@ static inline void disable_preempt (shim_tcb_t * tcb)
 
 static inline void __enable_preempt (shim_tcb_t * tcb)
 {
-    //tcb->context.syscall_nr -= SYSCALL_NR_PREEMPT_INC;
+    int64_t preempt = atomic_add_return(-1, &tcb->context.preempt);
     /* Assert if this counter underflows */
-    assert(tcb->context.preempt > 0);
-    tcb->context.preempt--;
-    //debug("enable preempt: %d\n", tcb->context.preempt);
+    assert(preempt >= 0);
+    //debug("enable preempt: %d\n", preempt);
 }
 
 void __handle_signal (shim_tcb_t * tcb, int sig);
@@ -500,10 +500,11 @@ static inline void enable_preempt (shim_tcb_t * tcb)
     if (!tcb && !(tcb = shim_get_tls()))
         return;
 
-    if (!tcb->context.preempt)
+    int64_t preempt = atomic_read(&tcb->context.preempt);
+    if (!preempt)
         return;
 
-    if (tcb->context.preempt == 1)
+    if (preempt == 1)
         __handle_signal(tcb, 0);
 
     __enable_preempt(tcb);

LibOS/shim/include/shim_tls.h

@@ -3,6 +3,22 @@
 
 #ifndef __ASSEMBLER__
 
+#ifdef IN_SHIM
+
+#include <shim_defs.h>
+#include <atomic.h>
+
+#else  /* !IN_SHIM */
+/* workaround to make glibc build
+ * the following structure must match to the one defined in pal/lib/atomic.h
+ */
+#ifdef __x86_64__
+struct atomic_int {
+    volatile int64_t counter;
+};
+#endif
+#endif /* IN_SHIM */
+
 #define SHIM_TLS_CANARY 0xdeadbeef
 
 struct lock_record {
@@ -40,15 +56,9 @@ struct shim_context {
     struct shim_regs *      regs;
     struct shim_context *   next;
     uint64_t                enter_time;
-    uint64_t                preempt;
+    struct atomic_int       preempt;
 };
 
-#ifdef IN_SHIM
-
-#include <shim_defs.h>
-
-#endif /* IN_SHIM */
-
 struct debug_buf;
 
 typedef struct shim_tcb shim_tcb_t;

LibOS/shim/src/bookkeep/shim_signal.c

@@ -163,7 +163,7 @@ void deliver_signal (siginfo_t * info, PAL_CONTEXT * context)
     struct shim_thread * cur_thread = (struct shim_thread *) tcb->tp;
     int sig = info->si_signo;
 
-    __disable_preempt(tcb);
+    int64_t preempt = __disable_preempt(tcb);
 
     struct shim_signal * signal = __alloca(sizeof(struct shim_signal));
     /* save in signal */
@@ -172,7 +172,7 @@ void deliver_signal (siginfo_t * info, PAL_CONTEXT * context)
     __store_context(tcb, context, signal);
     signal->pal_context = context;
 
-    if (tcb->context.preempt > 1 ||
+    if (preempt > 1 ||
         __sigismember(&cur_thread->signal_mask, sig)) {
         struct shim_signal ** signal_log = NULL;
         if ((signal = malloc_copy(signal,sizeof(struct shim_signal))) &&
@@ -561,8 +561,8 @@ static void resume_upcall (PAL_PTR event, PAL_NUM arg, PAL_CONTEXT * context)
         return;
 
     if (!is_internal_tid(get_cur_tid())) {
-        __disable_preempt(tcb);
-        if (tcb->context.preempt <= 1)
+        int64_t preempt = __disable_preempt(tcb);
+        if (preempt <= 1)
             __handle_signal(tcb, 0);
         __enable_preempt(tcb);
     }
@@ -717,10 +717,10 @@ void handle_signal (void)
     if (!thread || !thread->has_signal.counter)
         return;
 
-    __disable_preempt(tcb);
+    int64_t preempt = __disable_preempt(tcb);
 
-    if (tcb->context.preempt > 1)
-        debug("signal delayed (%ld)\n", tcb->context.preempt);
+    if (preempt > 1)
+        debug("signal delayed (%ld)\n", preempt);
     else
         __handle_signal(tcb, 0);