[LibOS] Remove the data race on thread::is_alive

Previously, there was a data race on thread::is_alive between one thread
checking whether it is the last thread alive via check_last_thread() and
another thread exiting via thread_exit(). The former checks if is_alive
is true, the latter sets it to false. However, the exiting thread will
truly exit only after it called DkThreadExit(), thus the race on is_alive
led to scenarios where two threads believe to be the last threads alive
and compete on terminating Async Helper/IPC threads and exiting the whole
process. This commit introduces cleanup_thread() called by Async Helper
to set is_alive to false and delete the thread, freeing its resources.
The data race is thus removed, and shim_thread object leak is prevented.

LibOS/shim/include/shim_internal.h

@@ -41,6 +41,8 @@
 #include <atomic.h>
 #include <shim_tcb.h>
 
+noreturn void shim_clean_and_exit(int exit_code);
+
 /* important macros and static inline functions */
 static inline unsigned int get_cur_tid(void)
 {
@@ -144,8 +146,6 @@ static inline PAL_HANDLE __open_shim_stdio (void)
     return shim_stdio;
 }
 
-noreturn void shim_terminate (int err);
-
 #define USE_PAUSE       0
 
 static inline void do_pause (void);
@@ -160,7 +160,7 @@ static inline void do_pause (void);
     do {                                                                    \
         __SYS_PRINTF("BUG() " __FILE__ ":%d\n", __LINE__);                  \
         PAUSE();                                                            \
-        shim_terminate(-ENOTRECOVERABLE);                                   \
+        shim_clean_and_exit(-ENOTRECOVERABLE);                              \
     } while (0)
 
 #define DEBUG_HERE() \
@@ -769,9 +769,6 @@ static inline bool memory_migrated(void * mem)
 extern void * __load_address, * __load_address_end;
 extern void * __code_address, * __code_address_end;
 
-/* cleanup and terminate process, preserve exit code if err == 0 */
-noreturn void shim_clean_and_exit(int err);
-
 unsigned long parse_int (const char * str);
 
 extern void * initial_stack;
@@ -797,13 +794,7 @@ void set_rlimit_cur(int resource, uint64_t rlim);
 
 int object_wait_with_retry(PAL_HANDLE handle);
 
-/* this struct is passed as the second argument to release_clear_child_id() */
-struct clear_child_tid_struct {
-    int* clear_child_tid;         /* passed to LibOS's clone() from user app */
-    int  clear_child_tid_val_pal; /* ptr to it is passed to PAL's DkThreadExit() */
-};
-
-void release_clear_child_id(IDTYPE caller, void* clear_child_tids);
+void release_clear_child_tid(int* clear_child_tid);
 
 #ifdef __x86_64__
 #define __SWITCH_STACK(stack_top, func, arg)                    \

LibOS/shim/include/shim_thread.h

@@ -50,7 +50,9 @@ struct shim_thread {
     struct shim_handle_map * handle_map;
 
     /* child tid */
-    int * set_child_tid, * clear_child_tid;
+    int* set_child_tid;
+    int* clear_child_tid;      /* LibOS zeroes it to notify parent that thread exited */
+    int  clear_child_tid_pal;  /* PAL zeroes it to notify LibOS that thread exited */
 
     /* signal handling */
     __sigset_t signal_mask;
@@ -272,7 +274,8 @@ void del_thread (struct shim_thread * thread);
 void add_simple_thread (struct shim_simple_thread * thread);
 void del_simple_thread (struct shim_simple_thread * thread);
 
-int check_last_thread (struct shim_thread * self);
+void cleanup_thread(IDTYPE caller, void* thread);
+int check_last_thread(struct shim_thread* self);
 
 #ifndef ALIAS_VFORK_AS_FORK
 void switch_dummy_thread (struct shim_thread * thread);
@@ -313,10 +316,7 @@ void set_handle_map (struct shim_thread * thread,
     thread->handle_map = map;
 }
 
-/* shim exit callback */
-int thread_exit(struct shim_thread* self, bool send_ipc, int** clear_child_tid_pal_ptr);
-/* If the process was killed by a signal, pass it in the second
- *  argument, else pass zero */
+int thread_exit(struct shim_thread* self, bool send_ipc);
 noreturn void thread_or_process_exit(int error_code, int term_signal);
 
 /* thread cloning helpers */

LibOS/shim/src/bookkeep/shim_thread.c

@@ -463,26 +463,62 @@ void del_simple_thread (struct shim_simple_thread * thread)
     put_simple_thread(thread);
 }
 
-int check_last_thread (struct shim_thread * self)
-{
-    struct shim_thread * tmp;
+static int _check_last_thread(struct shim_thread* self) {
+    IDTYPE self_tid = self ? self->tid : 0;
 
-    lock(&thread_list_lock);
-    /* find out if there is any thread that is
-       1) no current thread 2) in current vm
-       3) still alive */
-    LISTP_FOR_EACH_ENTRY(tmp, &thread_list, list) {
-        if (tmp->tid &&
-            (!self || tmp->tid != self->tid) && tmp->in_vm && tmp->is_alive) {
-            debug("check_last_thread: thread %d is alive\n", tmp->tid);
-            unlock(&thread_list_lock);
-            return tmp->tid;
+    struct shim_thread* thread;
+    LISTP_FOR_EACH_ENTRY(thread, &thread_list, list) {
+        if (thread->tid && thread->tid != self_tid && thread->in_vm && thread->is_alive) {
+            return thread->tid;
         }
     }
+    return 0;
+}
+
+/* Checks for any alive threads apart from thread self. Returns tid of the first found alive thread
+ * or 0 if there are no alive threads. self can be NULL, then all threads are checked. */
+int check_last_thread(struct shim_thread* self) {
+    lock(&thread_list_lock);
+    int alive_thread_tid = _check_last_thread(self);
+    unlock(&thread_list_lock);
+    return alive_thread_tid;
+}
+
+/* This function is called by Async Helper thread to wait on thread->clear_child_tid_pal to be
+ * zeroed (PAL does it when thread finally exits). Since it is a callback to Async Helper thread,
+ * this function must follow the `void (*callback) (IDTYPE caller, void* arg)` signature. */
+void cleanup_thread(IDTYPE caller, void* arg) {
+    __UNUSED(caller);
+
+    struct shim_thread* thread = (struct shim_thread*)arg;
+    assert(thread);
+
+    int exit_code = thread->term_signal ? : thread->exit_code;
+
+    /* wait on clear_child_tid_pal; this signals that PAL layer exited child thread */
+    while (__atomic_load_n(&thread->clear_child_tid_pal, __ATOMIC_RELAXED) != 0) {
+        __asm__ volatile ("pause");
+    }
+
+    /* notify parent if any */
+    release_clear_child_tid(thread->clear_child_tid);
+
+    /* clean up the thread itself */
+    lock(&thread_list_lock);
+    thread->is_alive = false;
+    LISTP_DEL_INIT(thread, &thread_list, list);
+
+    put_thread(thread);
+
+    if (!_check_last_thread(NULL)) {
+        /* corner case when all application threads exited via exit(), only Async helper
+         * and IPC helper threads are left at this point so simply exit process (recall
+         * that typically processes exit via exit_group()) */
+        unlock(&thread_list_lock);
+        shim_clean_and_exit(exit_code);
+    }
 
-    debug("this is the only thread %d\n", self->tid);
     unlock(&thread_list_lock);
-    return 0;
 }
 
 int walk_thread_list (int (*callback) (struct shim_thread *, void *, bool *),

LibOS/shim/src/ipc/shim_ipc_child.c

@@ -80,7 +80,7 @@ static int child_thread_exit(struct shim_thread* thread, void* arg, bool* unlock
             /* remote thread is "virtually" exited: SIGCHLD is generated for
              * the parent thread and exit events are arranged for subsequent
              * wait4(). */
-            thread_exit(thread, /*send_ipc=*/false, /*clear_child_tid_pal_ptr*/NULL);
+            thread_exit(thread, /*send_ipc=*/false);
             goto out;
         }
     }
@@ -207,7 +207,7 @@ int ipc_cld_exit_callback(struct shim_ipc_msg* msg, struct shim_ipc_port* port)
 
         /* Remote thread is "virtually" exited: SIGCHLD is generated for the
          * parent thread and exit events are arranged for subsequent wait4(). */
-        ret = thread_exit(thread, /*send_ipc=*/false, /*clear_child_tid_pal_ptr*/NULL);
+        ret = thread_exit(thread, /*send_ipc=*/false);
         put_thread(thread);
     } else {
         /* Uncommon case: remote child thread was already exited and deleted

LibOS/shim/src/shim_async.c

@@ -85,7 +85,7 @@ int64_t install_async_event(PAL_HANDLE object, uint64_t time,
 
     lock(&async_helper_lock);
 
-    if (callback != &release_clear_child_id && !object) {
+    if (callback != &cleanup_thread && !object) {
         /* This is alarm() or setitimer() emulation, treat both according to
          * alarm() syscall semantics: cancel any pending alarm/timer. */
         struct async_event * tmp, * n;
@@ -210,8 +210,8 @@ static void shim_async_helper(void * arg) {
              *   1. Exited child:  trigger callback and remove from the list;
              *   2. IO events:     trigger callback and keep in the list;
              *   3. alarms/timers: trigger callback and remove from the list. */
-            if (tmp->callback == &release_clear_child_id) {
-                debug("Child exited, notifying parent if any\n");
+            if (tmp->callback == &cleanup_thread) {
+                debug("Thread exited, cleaning up\n");
                 LISTP_DEL(tmp, &async_list, list);
                 LISTP_ADD_TAIL(tmp, &triggered, list);
                 continue;

LibOS/shim/src/shim_init.c

@@ -66,7 +66,7 @@ static void handle_failure (PAL_PTR event, PAL_NUM arg, PAL_CONTEXT * context)
 
 noreturn void __abort(void) {
     PAUSE();
-    shim_terminate(-ENOTRECOVERABLE);
+    shim_clean_and_exit(-ENOTRECOVERABLE);
 }
 
 void warn (const char *format, ...)
@@ -662,7 +662,7 @@ DEFINE_PROFILE_INTERVAL(init_signal,                init);
         int _err = CALL_INIT(func, ##__VA_ARGS__);                      \
         if (_err < 0) {                                                 \
             SYS_PRINTF("shim_init() in " #func " (%d)\n", _err);        \
-            shim_terminate(_err);                                       \
+            shim_clean_and_exit(_err);                                  \
         }                                                               \
         SAVE_PROFILE_INTERVAL(func);                                    \
     } while (0)
@@ -691,7 +691,7 @@ noreturn void* shim_init (int argc, void * args)
     g_pal_alloc_align = PAL_CB(alloc_align);
     if (!IS_POWER_OF_2(g_pal_alloc_align)) {
         SYS_PRINTF("shim_init(): error: PAL allocation alignment not a power of 2\n");
-        shim_terminate(-EINVAL);
+        shim_clean_and_exit(-EINVAL);
     }
 
     create_lock(&__master_lock);
@@ -1129,26 +1129,16 @@ static void print_profile_result (PAL_HANDLE hdl, struct shim_profile * root,
 }
 #endif /* PROFILE */
 
-static struct atomic_int in_terminate = { .counter = 0, };
-
-noreturn void shim_terminate (int err)
-{
-    debug("teminating the whole process (%d)\n", err);
-
-    /* do last clean-up of the process */
-    shim_clean_and_exit(err);
-}
-
-/* cleanup and terminate process, preserve exit code if err == 0 */
-noreturn void shim_clean_and_exit(int err) {
-    if (atomic_inc_return(&in_terminate) > 1) {
+noreturn void shim_clean_and_exit(int exit_code) {
+    static int in_terminate = 0;
+    if (__atomic_add_fetch(&in_terminate, 1, __ATOMIC_RELAXED) > 1) {
         while (true) {
             /* nothing */
         }
     }
 
-    if (err != 0)
-        cur_process.exit_code = err;
+    cur_process.exit_code = exit_code;
+
     store_all_msg_persist();
 
 #ifdef PROFILE

LibOS/shim/src/sys/shim_clone.c

@@ -287,7 +287,6 @@ int shim_do_clone (int flags, void * user_stack_addr, int * parent_tidptr,
     }
 
     if (flags & CLONE_CHILD_CLEARTID)
-        /* Implemented in shim_futex.c: release_clear_child_id */
         thread->clear_child_tid = child_tidptr;
 
     unsigned long fs_base = 0;

LibOS/shim/src/sys/shim_exec.c

@@ -202,7 +202,7 @@ retry_dump_vmas:
 
 error:
     debug("execve: failed %d\n", ret);
-    shim_terminate(ret);
+    shim_clean_and_exit(ret);
 }
 
 static int shim_do_execve_rtld (struct shim_handle * hdl, const char ** argv,

LibOS/shim/src/sys/shim_exit.c

@@ -40,7 +40,7 @@
 
 void release_robust_list (struct robust_list_head * head);
 
-int thread_exit(struct shim_thread* self, bool send_ipc, int** clear_child_tid_pal_ptr) {
+int thread_exit(struct shim_thread* self, bool send_ipc) {
     bool sent_exit_msg = false;
 
     /* Chia-Che: Broadcast exit message as early as possible,
@@ -64,7 +64,11 @@ int thread_exit(struct shim_thread* self, bool send_ipc, int** clear_child_tid_p
     #endif
 
     int exit_code = self->exit_code;
-    self->is_alive = false;
+
+    if (!self->in_vm) {
+        /* thread is in another process, we cannot rely on Async Helper thread to clean it */
+        self->is_alive = false;
+    }
 
     if (is_internal(self))
         goto out;
@@ -119,24 +123,6 @@ int thread_exit(struct shim_thread* self, bool send_ipc, int** clear_child_tid_p
     if (robust_list)
         release_robust_list(robust_list);
 
-    if (parent && self->in_vm && self->clear_child_tid) {
-        /* ask Async Helper thread to wake up parent when this child thread finally exits;
-         * we must alloc clear_child_tids on heap instead of this thread's stack; it is
-         * freed in release_clear_child_id() */
-        struct clear_child_tid_struct* clear_child_tids = malloc(sizeof(*clear_child_tids));
-        if (clear_child_tids) {
-            clear_child_tids->clear_child_tid         = self->clear_child_tid;
-            clear_child_tids->clear_child_tid_val_pal = 1; /* any non-zero value suffices */
-            install_async_event(NULL, 0, &release_clear_child_id, clear_child_tids);
-
-            if (clear_child_tid_pal_ptr) {
-                /* caller wants to performs DkThreadExit() and needs to know which address
-                 * PAL must set to inform the Async Helper thread */
-                *clear_child_tid_pal_ptr = &clear_child_tids->clear_child_tid_val_pal;
-            }
-        }
-    }
-
     DkEventSet(self->exit_event);
     return 0;
 }
@@ -146,15 +132,17 @@ noreturn void thread_or_process_exit(int error_code, int term_signal) {
     struct shim_thread * cur_thread = get_cur_thread();
 
     cur_thread->exit_code = -error_code;
-    cur_process.exit_code = term_signal ? term_signal : error_code;
     cur_thread->term_signal = term_signal;
 
-    int* clear_child_tid_pal = NULL;
     if (cur_thread->in_vm)
-        thread_exit(cur_thread, true, &clear_child_tid_pal);
+        thread_exit(cur_thread, true);
 
     if (check_last_thread(cur_thread)) {
-        DkThreadExit(clear_child_tid_pal);
+        /* ask Async Helper thread to cleanup this thread */
+        cur_thread->clear_child_tid_pal = 1; /* any non-zero value suffices */
+        install_async_event(NULL, 0, &cleanup_thread, cur_thread);
+
+        DkThreadExit(&cur_thread->clear_child_tid_pal);
     }
 
     struct shim_thread * async_thread = terminate_async_helper();
@@ -173,7 +161,7 @@ noreturn void thread_or_process_exit(int error_code, int term_signal) {
          */
         put_thread(ipc_thread); /* free resources of the thread */
 
-    shim_clean_and_exit(0);
+    shim_clean_and_exit(term_signal ? term_signal : error_code);
 }
 
 noreturn int shim_do_exit_group (int error_code)
@@ -197,19 +185,13 @@ noreturn int shim_do_exit_group (int error_code)
 #ifndef ALIAS_VFORK_AS_FORK
     if (cur_thread->dummy) {
         cur_thread->term_signal = 0;
-        thread_exit(cur_thread, true, NULL);
+        thread_exit(cur_thread, true);
         switch_dummy_thread(cur_thread);
     }
 #endif
 
     debug("now kill other threads in the process\n");
     do_kill_proc(cur_thread->tgid, cur_thread->tgid, SIGKILL, false);
-    /* This loop ensures that the current thread, which issues exit_group(), wins in setting the
-     * process's exit code. thread_or_process_exit() first sets the exit_code before updating the
-     * thread's state to "dead". Once check_last_thread() indicates that the current thread is the
-     * last thread, all the children will already have set thread->exit_code. Hence, this thread's
-     * execution of thread_or_process_exit() gets to determine the final exit_code, which is the
-     * desired outcome. */
     while (check_last_thread(cur_thread)) {
         DkThreadYieldExecution();
     }
@@ -237,7 +219,7 @@ noreturn int shim_do_exit (int error_code)
 #ifndef ALIAS_VFORK_AS_FORK
     if (cur_thread->dummy) {
         cur_thread->term_signal = 0;
-        thread_exit(cur_thread, true, NULL);
+        thread_exit(cur_thread, true);
         switch_dummy_thread(cur_thread);
     }
 #endif

LibOS/shim/src/sys/shim_futex.c

@@ -427,28 +427,10 @@ void release_robust_list(struct robust_list_head* head) {
     }
 }
 
-/* Function is called by Async Helper thread to wait on clear_child_tid_val_pal to be set to 0
- * (PAL does it when child thread finally exits). Next, *clear_child_tid is set to 0 and parent
- * threads are woken up. Since it is a callback to Async Helper thread, it must follow the
- * `void (*callback) (IDTYPE caller, void * arg)` signature even though we don't use caller. */
-void release_clear_child_id(IDTYPE caller, void* clear_child_tids) {
-    __UNUSED(caller);
-
-    struct clear_child_tid_struct* child = (struct clear_child_tid_struct*)clear_child_tids;
-    if (!child || !child->clear_child_tid)
-        goto out;
-
-    /* wait on clear_child_tid_val_pal; this signals that PAL layer exited child thread */
-    while (__atomic_load_n(&child->clear_child_tid_val_pal, __ATOMIC_RELAXED) != 0) {
-        __asm__ volatile ("pause");
-    }
-
-    /* child thread exited, now parent can wake up; note that PAL layer can't set clear_child_tid
-     * itself, because parent thread could spuriously wake up, notice 0 on clear_child_tid, and
-     * continue its execution without waiting for this function to succeed first */
-    __atomic_store_n(child->clear_child_tid, 0, __ATOMIC_RELAXED);
+void release_clear_child_tid(int* clear_child_tid) {
+    /* child thread exited, now parent can wake up */
+    __atomic_store_n(clear_child_tid, 0, __ATOMIC_RELAXED);
 
-    /* at this point, child thread finally exited, can wake up parents if any */
     create_lock_runtime(&futex_list_lock);
 
     struct shim_futex_handle* tmp;
@@ -456,19 +438,17 @@ void release_clear_child_id(IDTYPE caller, void* clear_child_tids) {
 
     lock(&futex_list_lock);
     LISTP_FOR_EACH_ENTRY(tmp, &futex_list, list) {
-        if (tmp->uaddr == (void*)child->clear_child_tid) {
+        if (tmp->uaddr == (void*)clear_child_tid) {
             futex = tmp;
             break;
         }
     }
     unlock(&futex_list_lock);
 
-    if (!futex) {
-        /* no parent threads waiting on this child to exit */
-        goto out;
-    }
+    if (!futex)
+        return;
 
-    debug("release futex at %p\n", child->clear_child_tid);
+    debug("release futex at %p\n", clear_child_tid);
     struct futex_waiter* waiter;
     struct futex_waiter* wtmp;
     struct shim_handle* hdl = container_of(futex, struct shim_handle, info.futex);
@@ -481,7 +461,4 @@ void release_clear_child_id(IDTYPE caller, void* clear_child_tids) {
     }
     unlock(&hdl->lock);
     put_handle(hdl);
-
-out:
-    free(child);
 }