fix: modified implementation

Author: Hritik Gupta

core/iwasm/common/wasm_shared_memory.c

@@ -10,6 +10,10 @@ static bh_list shared_memory_list_head;
 static bh_list *const shared_memory_list = &shared_memory_list_head;
 static korp_mutex shared_memory_list_lock;
 
+static bh_list atomic_wait_address_list_head;
+static bh_list *const atomic_wait_address_list = &atomic_wait_address_list_head;
+static korp_mutex atomic_wait_address_list_lock;
+
 /* clang-format off */
 enum {
     S_WAITING,
@@ -30,6 +34,11 @@ typedef struct AtomicWaitNode {
     korp_cond wait_cond;
 } AtomicWaitNode;
 
+typedef struct AtomicWaitAddressNode {
+    bh_list_link l;
+    void *addr;
+} AtomicWaitAddressNode;
+
 /* Atomic wait map */
 static HashMap *wait_map;
 
@@ -47,6 +56,8 @@ wasm_shared_memory_init()
 {
     if (os_mutex_init(&shared_memory_list_lock) != 0)
         return false;
+    if (os_mutex_init(&atomic_wait_address_list_lock) != 0)
+        return false;
     /* wait map not exists, create new map */
     if (!(wait_map = bh_hash_map_create(32, true, (HashFunc)wait_address_hash,
                                         (KeyEqualFunc)wait_address_equal, NULL,
@@ -62,6 +73,7 @@ void
 wasm_shared_memory_destroy()
 {
     os_mutex_destroy(&shared_memory_list_lock);
+    os_mutex_destroy(&atomic_wait_address_list_lock);
     if (wait_map) {
         bh_hash_map_destroy(wait_map);
     }
@@ -87,41 +99,18 @@ search_module(WASMModuleCommon *module)
     return NULL;
 }
 
-static void
-notify_stale_threads(void *key, void *value, void *user_data)
-{
-    WASMModuleInstanceCommon *module_inst =
-        (WASMModuleInstanceCommon *)user_data;
-    wasm_runtime_atomic_notify(module_inst, key, UINT32_MAX);
-    bh_hash_map_remove(wait_map, key, NULL, NULL);
-}
-
-static void
-create_copy_of_wait_map(void *key, void *value, void *map_copy)
-{
-    bh_hash_map_insert((HashMap *)map_copy, key, value);
-}
-
 void
 notify_stale_threads_on_exception(WASMModuleInstanceCommon *module_inst)
 {
-    static HashMap *wait_map_without_lock;
-    wait_map_without_lock = bh_hash_map_create(
-        32, false, (HashFunc)wait_address_hash,
-        (KeyEqualFunc)wait_address_equal, NULL, destroy_wait_info);
-    /**
-     * create a copy of wait_map with no lock
-     * this allows invoking methods in the callback
-     * which require acquiring lock on the original map
-     */
-    os_mutex_lock(&shared_memory_list_lock);
-    bh_hash_map_traverse(wait_map, create_copy_of_wait_map,
-                         wait_map_without_lock);
-    os_mutex_unlock(&shared_memory_list_lock);
+    AtomicWaitAddressNode *node;
 
-    // terminate stale threads
-    bh_hash_map_traverse(wait_map_without_lock, notify_stale_threads,
-                         module_inst);
+    os_mutex_lock(&atomic_wait_address_list_lock);
+    node = bh_list_first_elem(atomic_wait_address_list);
+    while (node) {
+        wasm_runtime_atomic_notify(module_inst, node->addr, UINT32_MAX);
+        node = bh_list_elem_next(node);
+    }
+    os_mutex_unlock(&atomic_wait_address_list_lock);
 }
 
 WASMSharedMemNode *
@@ -341,6 +330,11 @@ release_wait_info(HashMap *wait_map_, AtomicWaitInfo *wait_info, void *address)
 
     if (wait_info->wait_list->len == 0) {
         bh_hash_map_remove(wait_map_, address, NULL, NULL);
+
+        os_mutex_lock(&atomic_wait_address_list_lock);
+        bh_list_remove(atomic_wait_address_list, address);
+        os_mutex_unlock(&atomic_wait_address_list_lock);
+
         destroy_wait_info(wait_info);
     }
 
@@ -382,6 +376,22 @@ wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,
 
     os_mutex_lock(&wait_info->wait_list_lock);
 
+    /* add waiter address to atomic wait list */
+    AtomicWaitAddressNode *wait_address_node;
+
+    os_mutex_lock(&atomic_wait_address_list_lock);
+    if (!(wait_address_node =
+              wasm_runtime_malloc(sizeof(AtomicWaitAddressNode)))) {
+        wasm_runtime_set_exception(module,
+                                   "failed to create wait address node");
+        os_mutex_unlock(&atomic_wait_address_list_lock);
+        return -1;
+    }
+    wait_address_node->addr = address;
+    bh_assert(bh_list_insert(atomic_wait_address_list, wait_address_node)
+              == BH_LIST_SUCCESS);
+    os_mutex_unlock(&atomic_wait_address_list_lock);
+
     if ((!wait64 && *(uint32 *)address != (uint32)expect)
         || (wait64 && *(uint64 *)address != expect)) {
         os_mutex_unlock(&wait_info->wait_list_lock);