[LibOS] Make get_new_id immediately move ID ownership if needed

LibOS/shim/include/shim_thread.h

@@ -143,14 +143,15 @@ static inline bool is_internal(struct shim_thread* thread) {
 /*!
  * \brief Allocates new ID
  *
- * \param remove_from_owned if `true`, ID is removed from owned and locally tracked IDs
+ * \param move_ownership_to VMID of the process to pass ownership of ID to; if set, ID is removed
+ *                          from owned and locally tracked IDs
  *
  * \returns new ID on success, `0` on failure (`0` is an invalid ID)
  *
- * If \p remove_from_owned is `true`, the returned ID cannot be freed with #release_id since it's
- * no longer locally tracked. You probably want to call #ipc_change_id_owner afterwards.
+ * If \p move_ownership_to is set, the returned ID should not be freed with #release_id since it's
+ * no longer locally tracked and the current process no longer owns it.
  */
-IDTYPE get_new_id(bool remove_from_owned);
+IDTYPE get_new_id(IDTYPE move_ownership_to);
 /*!
  * \brief Releases (frees) previously allocated ID
  *

LibOS/shim/src/bookkeep/shim_pid.c

@@ -66,7 +66,7 @@ int init_id_ranges(IDTYPE preload_tid) {
     return 0;
 }
 
-IDTYPE get_new_id(bool remove_from_owned) {
+IDTYPE get_new_id(IDTYPE move_ownership_to) {
     IDTYPE ret_id = 0;
     lock(&g_ranges_lock);
     if (!g_last_range) {
@@ -99,7 +99,7 @@ IDTYPE get_new_id(bool remove_from_owned) {
     ret_id = ++g_last_used_id;
     g_last_range->taken_count++;
 
-    if (remove_from_owned) {
+    if (move_ownership_to) {
         g_last_range->taken_count--;
         if (g_last_range->start == g_last_range->end) {
             assert(g_last_range->taken_count == 0);
@@ -128,6 +128,11 @@ IDTYPE get_new_id(bool remove_from_owned) {
             avl_tree_insert(&g_used_ranges_tree, &g_last_range->node);
             g_last_range = range;
         }
+        if (ipc_change_id_owner(ret_id, move_ownership_to) < 0) {
+            /* Good luck unwinding all of above operations. Better just kill everything. */
+            log_error("Unrecoverable error in %s:%d", __FILE__, __LINE__);
+            DkProcessExit(1);
+        }
     } else {
         if (g_last_used_id == g_last_range->end) {
             avl_tree_insert(&g_used_ranges_tree, &g_last_range->node);

LibOS/shim/src/bookkeep/shim_thread.c

@@ -141,7 +141,7 @@ static int init_main_thread(void) {
         return -ENOMEM;
     }
 
-    cur_thread->tid = get_new_id(/*remove_from_owned=*/false);
+    cur_thread->tid = get_new_id(/*move_ownership_to=*/0);
     if (!cur_thread->tid) {
         log_error("Cannot allocate pid for the initial thread!");
         put_thread(cur_thread);

LibOS/shim/src/ipc/shim_ipc_pid.c

@@ -202,7 +202,9 @@ static void release_id_range(IDTYPE start, IDTYPE end) {
         BUG();
     }
     struct id_range* range = container_of(node, struct id_range, node);
-    assert(range->start == start && range->end == end);
+    if (range->start != start || range->end != end) {
+        BUG();
+    }
     avl_tree_delete(&g_id_owners_tree, &range->node);
 
     unlock(&g_id_owners_tree_lock);
@@ -345,7 +347,7 @@ int ipc_change_id_owner(IDTYPE id, IDTYPE new_owner) {
     init_ipc_msg(msg, IPC_MSG_CHANGE_ID_OWNER, msg_size);
     memcpy(&msg->data, &owner_msg, sizeof(owner_msg));
 
-    log_debug("%s: sending a request (%u..%u)", __func__, id, new_owner);
+    log_debug("%s: sending a request (%u, %u)", __func__, id, new_owner);
 
     int ret = ipc_send_msg_and_get_response(g_process_ipc_ids.leader_vmid, msg, /*resp=*/NULL);
     log_debug("%s: ipc_send_msg_and_get_response: %d", __func__, ret);
@@ -356,7 +358,7 @@ int ipc_change_id_owner(IDTYPE id, IDTYPE new_owner) {
 int ipc_change_id_owner_callback(IDTYPE src, void* data, uint64_t seq) {
     struct ipc_id_owner_msg* owner_msg = data;
     int ret = change_id_owner(owner_msg->id, owner_msg->owner);
-    log_debug("%s: change_id_owner(%u..%u): %d", __func__, owner_msg->id, owner_msg->owner, ret);
+    log_debug("%s: change_id_owner(%u, %u): %d", __func__, owner_msg->id, owner_msg->owner, ret);
     if (ret < 0) {
         return ret;
     }

LibOS/shim/src/shim_init.c

@@ -456,13 +456,15 @@ noreturn void* shim_init(int argc, void* args) {
             DkProcessExit(1);
         }
 
-        /* This has also a (very much desired) side effect of the IPC leader making a connection to
-         * this process, so that it's included in all broadcast messages. */
-        ret = ipc_change_id_owner(g_process.pid, g_process_ipc_ids.self_vmid);
+        /* Send a dummy request causing the IPC leader to connect to this process, so that it is
+         * included in all broadcast messages. */
+        IDTYPE dummy = 0;
+        ret = ipc_get_id_owner(/*id=*/0, /*out_owner=*/&dummy);
         if (ret < 0) {
-            log_debug("shim_init: failed to change child process PID ownership: %d", ret);
+            log_debug("shim_init: failed to get a connection from IPC leader to us: %d", ret);
             DkProcessExit(1);
         }
+        assert(dummy == 0); // Nobody should own ID `0`.
 
         /* Notify the parent process we are done. */
         char dummy_c = 0;

LibOS/shim/src/sys/shim_clone.c

@@ -128,8 +128,8 @@ static int migrate_fork(struct shim_cp_store* store, struct shim_process* proces
     return START_MIGRATE(store, fork, process_description, thread_description, process_ipc_ids);
 }
 
-static long do_clone_new_vm(unsigned long flags, struct shim_thread* thread, unsigned long tls,
-                            unsigned long user_stack_addr, int* set_parent_tid) {
+static long do_clone_new_vm(IDTYPE child_vmid, unsigned long flags, struct shim_thread* thread,
+                            unsigned long tls, unsigned long user_stack_addr, int* set_parent_tid) {
     assert(!(flags & CLONE_VM));
 
     struct shim_child_process* child_process = create_child_process();
@@ -193,11 +193,10 @@ static long do_clone_new_vm(unsigned long flags, struct shim_thread* thread, uns
     child_process->pid = process_description.pid;
     child_process->child_termination_signal = flags & CSIGNAL;
     child_process->uid = thread->uid;
-    long ret = ipc_get_new_vmid(&child_process->vmid);
-    if (!ret) {
-        ret = create_process_and_send_checkpoint(&migrate_fork, child_process, &process_description,
-                                                 thread);
-    }
+    child_process->vmid = child_vmid;
+
+    long ret = create_process_and_send_checkpoint(&migrate_fork, child_process,
+                                                  &process_description, thread);
 
     if (parent_stack) {
         pal_context_set_sp(self->shim_tcb->context.regs, parent_stack);
@@ -335,7 +334,17 @@ long shim_do_clone(unsigned long flags, unsigned long user_stack_addr, int* pare
         tls = get_tls();
     }
 
-    IDTYPE tid = get_new_id(/*remove_from_owned=*/clone_new_process);
+    IDTYPE new_vmid = 0;
+    if (clone_new_process) {
+        ret = ipc_get_new_vmid(&new_vmid);
+        if (ret < 0) {
+            log_error("Cound not allocate new vmid!");
+            ret = -EAGAIN;
+            goto failed;
+        }
+    }
+
+    IDTYPE tid = get_new_id(/*move_ownership_to=*/new_vmid);
     if (!tid) {
         log_error("Could not allocate a tid!");
         ret = -EAGAIN;
@@ -344,7 +353,7 @@ long shim_do_clone(unsigned long flags, unsigned long user_stack_addr, int* pare
     thread->tid = tid;
 
     if (clone_new_process) {
-        ret = do_clone_new_vm(flags, thread, tls, user_stack_addr, set_parent_tid);
+        ret = do_clone_new_vm(new_vmid, flags, thread, tls, user_stack_addr, set_parent_tid);
 
         /* We should not have saved any references to this thread anywhere and `put_thread` below
          * should free it. */
@@ -356,9 +365,7 @@ long shim_do_clone(unsigned long flags, unsigned long user_stack_addr, int* pare
         put_thread(thread);
 
         if (ret < 0) {
-            /* The child process might have already taken the ownership of `tid`, let's change it
-             * back (if we still own it, this call will split `tid` from any other range, if it is
-             * a part of one)... */
+            /* The child process have already taken the ownership of `tid`, let's change it back. */
             int tmp_ret = ipc_change_id_owner(tid, g_process_ipc_ids.self_vmid);
             if (tmp_ret < 0) {
                 log_debug("Failed to change back ID %u owner: %d", tid, tmp_ret);