Reserve TID values for WASI threads (#1862)

According to the [WASI thread specification](WebAssembly/wasi-threads#16),
some thread identifiers are reserved and should not be used. In fact, only IDs between `1` and
`0x1FFFFFFF` are valid.

The thread ID allocator has been moved to a separate class to avoid polluting the
`lib_wasi_threads_wrapper` logic.

Author: eloparco

core/iwasm/libraries/lib-wasi-threads/lib_wasi_threads.cmake

@@ -8,4 +8,5 @@ add_definitions (-DWASM_ENABLE_LIB_WASI_THREADS=1 -DWASM_ENABLE_HEAP_AUX_STACK_A
 include_directories(${LIB_WASI_THREADS_DIR})
 
 set (LIB_WASI_THREADS_SOURCE
-    ${LIB_WASI_THREADS_DIR}/lib_wasi_threads_wrapper.c)
+    ${LIB_WASI_THREADS_DIR}/lib_wasi_threads_wrapper.c
+    ${LIB_WASI_THREADS_DIR}/tid_allocator.c)

core/iwasm/libraries/lib-wasi-threads/lib_wasi_threads_wrapper.c

@@ -5,6 +5,7 @@
 
 #include "bh_log.h"
 #include "thread_manager.h"
+#include "tid_allocator.h"
 
 #if WASM_ENABLE_INTERP != 0
 #include "wasm_runtime.h"
@@ -15,16 +16,8 @@
 #endif
 
 static const char *THREAD_START_FUNCTION = "wasi_thread_start";
-
 static korp_mutex thread_id_lock;
-
-// Stack data structure to track available thread identifiers
-#define AVAIL_TIDS_INIT_SIZE CLUSTER_MAX_THREAD_NUM
-typedef struct {
-    int32 *ids;
-    uint32 pos, size;
-} AvailableThreadIds;
-static AvailableThreadIds avail_tids;
+static TidAllocator tid_allocator;
 
 typedef struct {
     /* app's entry function */
@@ -38,53 +31,18 @@ typedef struct {
 static int32
 allocate_thread_id()
 {
-    int32 id = -1;
-
     os_mutex_lock(&thread_id_lock);
-    if (avail_tids.pos == 0) { // Resize stack and push new thread ids
-        uint32 old_size = avail_tids.size;
-        uint32 new_size = avail_tids.size * 2;
-        if (new_size / 2 != avail_tids.size) {
-            LOG_ERROR("Overflow detected during new size calculation");
-            goto return_id;
-        }
-
-        size_t realloc_size = new_size * sizeof(int32);
-        if (realloc_size / sizeof(int32) != new_size) {
-            LOG_ERROR("Overflow detected during realloc");
-            goto return_id;
-        }
-        int32 *tmp =
-            (int32 *)wasm_runtime_realloc(avail_tids.ids, realloc_size);
-        if (tmp == NULL) {
-            LOG_ERROR("Thread ID allocator realloc failed");
-            goto return_id;
-        }
-
-        avail_tids.size = new_size;
-        avail_tids.pos = old_size;
-        avail_tids.ids = tmp;
-        for (uint32 i = 0; i < old_size; i++)
-            avail_tids.ids[i] = new_size - i;
-    }
-
-    // Pop available thread identifier from `avail_tids` stack
-    id = avail_tids.ids[--avail_tids.pos];
-
-return_id:
+    int32 id = tid_allocator_get_tid(&tid_allocator);
     os_mutex_unlock(&thread_id_lock);
+
     return id;
 }
 
 void
 deallocate_thread_id(int32 thread_id)
 {
     os_mutex_lock(&thread_id_lock);
-
-    // Release thread identifier by pushing it into `avail_tids` stack
-    bh_assert(avail_tids.pos < avail_tids.size);
-    avail_tids.ids[avail_tids.pos++] = thread_id;
-
+    tid_allocator_release_tid(&tid_allocator, thread_id);
     os_mutex_unlock(&thread_id_lock);
 }
 
@@ -212,25 +170,17 @@ lib_wasi_threads_init(void)
     if (0 != os_mutex_init(&thread_id_lock))
         return false;
 
-    // Initialize stack to store thread identifiers
-    avail_tids.size = AVAIL_TIDS_INIT_SIZE;
-    avail_tids.pos = avail_tids.size;
-    avail_tids.ids =
-        (int32 *)wasm_runtime_malloc(avail_tids.size * sizeof(int32));
-    if (avail_tids.ids == NULL) {
+    if (!tid_allocator_init(&tid_allocator)) {
         os_mutex_destroy(&thread_id_lock);
         return false;
     }
-    for (uint32 i = 0; i < avail_tids.size; i++)
-        avail_tids.ids[i] = avail_tids.size - i;
 
     return true;
 }
 
 void
 lib_wasi_threads_destroy(void)
 {
-    wasm_runtime_free(avail_tids.ids);
-    avail_tids.ids = NULL;
+    tid_allocator_deinit(&tid_allocator);
     os_mutex_destroy(&thread_id_lock);
 }

core/iwasm/libraries/lib-wasi-threads/tid_allocator.c

@@ -0,0 +1,80 @@
+/*
+ * Copyright (C) 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ */
+
+#include "tid_allocator.h"
+#include "wasm_export.h"
+#include "bh_log.h"
+
+bh_static_assert(TID_MIN <= TID_MAX);
+#define MIN(a, b) (((a) < (b)) ? (a) : (b))
+
+bool
+tid_allocator_init(TidAllocator *tid_allocator)
+{
+    tid_allocator->size = MIN(TID_ALLOCATOR_INIT_SIZE, TID_MAX - TID_MIN + 1);
+    tid_allocator->pos = tid_allocator->size;
+    tid_allocator->ids =
+        wasm_runtime_malloc(tid_allocator->size * sizeof(int32));
+    if (tid_allocator->ids == NULL)
+        return false;
+
+    for (int64 i = tid_allocator->pos - 1; i >= 0; i--)
+        tid_allocator->ids[i] = TID_MIN + (tid_allocator->pos - 1 - i);
+
+    return true;
+}
+
+void
+tid_allocator_deinit(TidAllocator *tid_allocator)
+{
+    wasm_runtime_free(tid_allocator->ids);
+}
+
+int32
+tid_allocator_get_tid(TidAllocator *tid_allocator)
+{
+    if (tid_allocator->pos == 0) { // Resize stack and push new thread ids
+        if (tid_allocator->size == TID_MAX - TID_MIN + 1) {
+            LOG_ERROR("Maximum thread identifier reached");
+            return -1;
+        }
+
+        uint32 old_size = tid_allocator->size;
+        uint32 new_size = MIN(tid_allocator->size * 2, TID_MAX - TID_MIN + 1);
+        if (new_size != TID_MAX - TID_MIN + 1
+            && new_size / 2 != tid_allocator->size) {
+            LOG_ERROR("Overflow detected during new size calculation");
+            return -1;
+        }
+
+        size_t realloc_size = new_size * sizeof(int32);
+        if (realloc_size / sizeof(int32) != new_size) {
+            LOG_ERROR("Overflow detected during realloc");
+            return -1;
+        }
+        int32 *tmp = wasm_runtime_realloc(tid_allocator->ids, realloc_size);
+        if (tmp == NULL) {
+            LOG_ERROR("Thread ID allocator realloc failed");
+            return -1;
+        }
+
+        tid_allocator->size = new_size;
+        tid_allocator->pos = new_size - old_size;
+        tid_allocator->ids = tmp;
+        for (int64 i = tid_allocator->pos - 1; i >= 0; i--)
+            tid_allocator->ids[i] = TID_MIN + (tid_allocator->size - 1 - i);
+    }
+
+    // Pop available thread identifier from the stack
+    return tid_allocator->ids[--tid_allocator->pos];
+}
+
+void
+tid_allocator_release_tid(TidAllocator *tid_allocator, int32 thread_id)
+{
+    // Release thread identifier by pushing it into the stack
+    bh_assert(tid_allocator->pos < tid_allocator->size);
+    tid_allocator->ids[tid_allocator->pos++] = thread_id;
+}

core/iwasm/libraries/lib-wasi-threads/tid_allocator.h

@@ -0,0 +1,36 @@
+/*
+ * Copyright (C) 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ */
+
+#ifndef _TID_ALLOCATOR_H
+#define _TID_ALLOCATOR_H
+
+#include "platform_common.h"
+
+#define TID_ALLOCATOR_INIT_SIZE CLUSTER_MAX_THREAD_NUM
+enum {
+    TID_MIN = 1,
+    TID_MAX = 0x1FFFFFFF
+}; // Reserved TIDs (WASI specification)
+
+/* Stack data structure to track available thread identifiers */
+typedef struct {
+    int32 *ids;  // Array used to store the stack
+    uint32 size; // Stack capacity
+    uint32 pos;  // Index of the element after the stack top
+} TidAllocator;
+
+bool
+tid_allocator_init(TidAllocator *tid_allocator);
+
+void
+tid_allocator_deinit(TidAllocator *tid_allocator);
+
+int32
+tid_allocator_get_tid(TidAllocator *tid_allocator);
+
+void
+tid_allocator_release_tid(TidAllocator *tid_allocator, int32 thread_id);
+
+#endif /* _TID_ALLOCATOR_H */