Fix memory sharing (#2415)

- Inherit shared memory from the parent instance, instead of
  trying to look it up by the underlying module. The old method
  works correctly only when every cluster uses different module.
- Use reference count in WASMMemoryInstance/AOTMemoryInstance
  to mark whether the memory is shared or not
- Retire WASMSharedMemNode
- For atomic opcode implementations in the interpreters, use
  a global lock for now
- Update the internal API users
  (wasi-threads, lib-pthread, wasm_runtime_spawn_thread)

Fixes https://github.com/bytecodealliance/wasm-micro-runtime/issues/1962

core/iwasm/aot/aot_runtime.c

@@ -339,11 +339,8 @@ memories_deinstantiate(AOTModuleInstance *module_inst)
         memory_inst = module_inst->memories[i];
         if (memory_inst) {
 #if WASM_ENABLE_SHARED_MEMORY != 0
-            if (memory_inst->is_shared) {
-                int32 ref_count = shared_memory_dec_reference(
-                    (WASMModuleCommon *)module_inst->module);
-                bh_assert(ref_count >= 0);
-
+            if (shared_memory_is_shared(memory_inst)) {
+                uint32 ref_count = shared_memory_dec_reference(memory_inst);
                 /* if the reference count is not zero,
                     don't free the memory */
                 if (ref_count > 0)
@@ -373,9 +370,10 @@ memories_deinstantiate(AOTModuleInstance *module_inst)
 }
 
 static AOTMemoryInstance *
-memory_instantiate(AOTModuleInstance *module_inst, AOTModule *module,
-                   AOTMemoryInstance *memory_inst, AOTMemory *memory,
-                   uint32 heap_size, char *error_buf, uint32 error_buf_size)
+memory_instantiate(AOTModuleInstance *module_inst, AOTModuleInstance *parent,
+                   AOTModule *module, AOTMemoryInstance *memory_inst,
+                   AOTMemory *memory, uint32 memory_idx, uint32 heap_size,
+                   char *error_buf, uint32 error_buf_size)
 {
     void *heap_handle;
     uint32 num_bytes_per_page = memory->num_bytes_per_page;
@@ -396,23 +394,13 @@ memory_instantiate(AOTModuleInstance *module_inst, AOTModule *module,
     bool is_shared_memory = memory->memory_flags & 0x02 ? true : false;
 
     /* Shared memory */
-    if (is_shared_memory) {
+    if (is_shared_memory && parent != NULL) {
         AOTMemoryInstance *shared_memory_instance;
-        WASMSharedMemNode *node =
-            wasm_module_get_shared_memory((WASMModuleCommon *)module);
-        /* If the memory of this module has been instantiated,
-            return the memory instance directly */
-        if (node) {
-            uint32 ref_count;
-            ref_count = shared_memory_inc_reference((WASMModuleCommon *)module);
-            bh_assert(ref_count > 0);
-            shared_memory_instance =
-                (AOTMemoryInstance *)shared_memory_get_memory_inst(node);
-            bh_assert(shared_memory_instance);
-
-            (void)ref_count;
-            return shared_memory_instance;
-        }
+        bh_assert(memory_idx == 0);
+        bh_assert(parent->memory_count > memory_idx);
+        shared_memory_instance = parent->memories[memory_idx];
+        shared_memory_inc_reference(shared_memory_instance);
+        return shared_memory_instance;
     }
 #endif
 
@@ -609,23 +597,12 @@ memory_instantiate(AOTModuleInstance *module_inst, AOTModule *module,
 
 #if WASM_ENABLE_SHARED_MEMORY != 0
     if (is_shared_memory) {
-        memory_inst->is_shared = true;
-        if (!shared_memory_set_memory_inst(
-                (WASMModuleCommon *)module,
-                (WASMMemoryInstanceCommon *)memory_inst)) {
-            set_error_buf(error_buf, error_buf_size, "allocate memory failed");
-            goto fail3;
-        }
+        memory_inst->ref_count = 1;
     }
 #endif
 
     return memory_inst;
 
-#if WASM_ENABLE_SHARED_MEMORY != 0
-fail3:
-    if (heap_size > 0)
-        mem_allocator_destroy(memory_inst->heap_handle);
-#endif
 fail2:
     if (heap_size > 0)
         wasm_runtime_free(memory_inst->heap_handle);
@@ -654,8 +631,9 @@ aot_get_default_memory(AOTModuleInstance *module_inst)
 }
 
 static bool
-memories_instantiate(AOTModuleInstance *module_inst, AOTModule *module,
-                     uint32 heap_size, char *error_buf, uint32 error_buf_size)
+memories_instantiate(AOTModuleInstance *module_inst, AOTModuleInstance *parent,
+                     AOTModule *module, uint32 heap_size, char *error_buf,
+                     uint32 error_buf_size)
 {
     uint32 global_index, global_data_offset, base_offset, length;
     uint32 i, memory_count = module->memory_count;
@@ -672,8 +650,8 @@ memories_instantiate(AOTModuleInstance *module_inst, AOTModule *module,
 
     memories = module_inst->global_table_data.memory_instances;
     for (i = 0; i < memory_count; i++, memories++) {
-        memory_inst = memory_instantiate(module_inst, module, memories,
-                                         &module->memories[i], heap_size,
+        memory_inst = memory_instantiate(module_inst, parent, module, memories,
+                                         &module->memories[i], i, heap_size,
                                          error_buf, error_buf_size);
         if (!memory_inst) {
             return false;
@@ -1100,9 +1078,9 @@ check_linked_symbol(AOTModule *module, char *error_buf, uint32 error_buf_size)
 }
 
 AOTModuleInstance *
-aot_instantiate(AOTModule *module, bool is_sub_inst, WASMExecEnv *exec_env_main,
-                uint32 stack_size, uint32 heap_size, char *error_buf,
-                uint32 error_buf_size)
+aot_instantiate(AOTModule *module, AOTModuleInstance *parent,
+                WASMExecEnv *exec_env_main, uint32 stack_size, uint32 heap_size,
+                char *error_buf, uint32 error_buf_size)
 {
     AOTModuleInstance *module_inst;
     const uint32 module_inst_struct_size =
@@ -1112,6 +1090,7 @@ aot_instantiate(AOTModule *module, bool is_sub_inst, WASMExecEnv *exec_env_main,
     uint64 total_size, table_size = 0;
     uint8 *p;
     uint32 i, extra_info_offset;
+    const bool is_sub_inst = parent != NULL;
 
     /* Check heap size */
     heap_size = align_uint(heap_size, 8);
@@ -1171,7 +1150,7 @@ aot_instantiate(AOTModule *module, bool is_sub_inst, WASMExecEnv *exec_env_main,
         goto fail;
 
     /* Initialize memory space */
-    if (!memories_instantiate(module_inst, module, heap_size, error_buf,
+    if (!memories_instantiate(module_inst, parent, module, heap_size, error_buf,
                               error_buf_size))
         goto fail;
 

core/iwasm/aot/aot_runtime.h

@@ -406,7 +406,7 @@ aot_unload(AOTModule *module);
  * Instantiate a AOT module.
  *
  * @param module the AOT module to instantiate
- * @param is_sub_inst the flag of sub instance
+ * @param parent the parent module instance
  * @param heap_size the default heap size of the module instance, a heap will
  *        be created besides the app memory space. Both wasm app and native
  *        function can allocate memory from the heap. If heap_size is 0, the
@@ -417,9 +417,9 @@ aot_unload(AOTModule *module);
  * @return return the instantiated AOT module instance, NULL if failed
  */
 AOTModuleInstance *
-aot_instantiate(AOTModule *module, bool is_sub_inst, WASMExecEnv *exec_env_main,
-                uint32 stack_size, uint32 heap_size, char *error_buf,
-                uint32 error_buf_size);
+aot_instantiate(AOTModule *module, AOTModuleInstance *parent,
+                WASMExecEnv *exec_env_main, uint32 stack_size, uint32 heap_size,
+                char *error_buf, uint32 error_buf_size);
 
 /**
  * Deinstantiate a AOT module instance, destroy the resources.

core/iwasm/common/wasm_memory.c

@@ -608,7 +608,7 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count)
     }
 
 #if WASM_ENABLE_SHARED_MEMORY != 0
-    if (memory->is_shared) {
+    if (shared_memory_is_shared(memory)) {
         memory->num_bytes_per_page = num_bytes_per_page;
         memory->cur_page_count = total_page_count;
         memory->max_page_count = max_page_count;
@@ -769,15 +769,13 @@ wasm_enlarge_memory(WASMModuleInstance *module, uint32 inc_page_count)
     bool ret = false;
 
 #if WASM_ENABLE_SHARED_MEMORY != 0
-    WASMSharedMemNode *node =
-        wasm_module_get_shared_memory((WASMModuleCommon *)module->module);
-    if (node)
-        os_mutex_lock(&node->shared_mem_lock);
+    if (module->memory_count > 0)
+        shared_memory_lock(module->memories[0]);
 #endif
     ret = wasm_enlarge_memory_internal(module, inc_page_count);
 #if WASM_ENABLE_SHARED_MEMORY != 0
-    if (node)
-        os_mutex_unlock(&node->shared_mem_lock);
+    if (module->memory_count > 0)
+        shared_memory_unlock(module->memories[0]);
 #endif
 
     return ret;

core/iwasm/common/wasm_runtime_common.c

@@ -1196,7 +1196,8 @@ wasm_runtime_unload(WASMModuleCommon *module)
 }
 
 WASMModuleInstanceCommon *
-wasm_runtime_instantiate_internal(WASMModuleCommon *module, bool is_sub_inst,
+wasm_runtime_instantiate_internal(WASMModuleCommon *module,
+                                  WASMModuleInstanceCommon *parent,
                                   WASMExecEnv *exec_env_main, uint32 stack_size,
                                   uint32 heap_size, char *error_buf,
                                   uint32 error_buf_size)
@@ -1204,14 +1205,14 @@ wasm_runtime_instantiate_internal(WASMModuleCommon *module, bool is_sub_inst,
 #if WASM_ENABLE_INTERP != 0
     if (module->module_type == Wasm_Module_Bytecode)
         return (WASMModuleInstanceCommon *)wasm_instantiate(
-            (WASMModule *)module, is_sub_inst, exec_env_main, stack_size,
-            heap_size, error_buf, error_buf_size);
+            (WASMModule *)module, (WASMModuleInstance *)parent, exec_env_main,
+            stack_size, heap_size, error_buf, error_buf_size);
 #endif
 #if WASM_ENABLE_AOT != 0
     if (module->module_type == Wasm_Module_AoT)
         return (WASMModuleInstanceCommon *)aot_instantiate(
-            (AOTModule *)module, is_sub_inst, exec_env_main, stack_size,
-            heap_size, error_buf, error_buf_size);
+            (AOTModule *)module, (AOTModuleInstance *)parent, exec_env_main,
+            stack_size, heap_size, error_buf, error_buf_size);
 #endif
     set_error_buf(error_buf, error_buf_size,
                   "Instantiate module failed, invalid module type");
@@ -1224,7 +1225,7 @@ wasm_runtime_instantiate(WASMModuleCommon *module, uint32 stack_size,
                          uint32 error_buf_size)
 {
     return wasm_runtime_instantiate_internal(
-        module, false, NULL, stack_size, heap_size, error_buf, error_buf_size);
+        module, NULL, NULL, stack_size, heap_size, error_buf, error_buf_size);
 }
 
 void
@@ -2310,10 +2311,8 @@ wasm_set_exception(WASMModuleInstance *module_inst, const char *exception)
     WASMExecEnv *exec_env = NULL;
 
 #if WASM_ENABLE_SHARED_MEMORY != 0
-    WASMSharedMemNode *node =
-        wasm_module_get_shared_memory((WASMModuleCommon *)module_inst->module);
-    if (node)
-        os_mutex_lock(&node->shared_mem_lock);
+    if (module_inst->memory_count > 0)
+        shared_memory_lock(module_inst->memories[0]);
 #endif
     if (exception) {
         snprintf(module_inst->cur_exception, sizeof(module_inst->cur_exception),
@@ -2323,8 +2322,8 @@ wasm_set_exception(WASMModuleInstance *module_inst, const char *exception)
         module_inst->cur_exception[0] = '\0';
     }
 #if WASM_ENABLE_SHARED_MEMORY != 0
-    if (node)
-        os_mutex_unlock(&node->shared_mem_lock);
+    if (module_inst->memory_count > 0)
+        shared_memory_unlock(module_inst->memories[0]);
 #endif
 
 #if WASM_ENABLE_THREAD_MGR != 0
@@ -2386,10 +2385,8 @@ wasm_copy_exception(WASMModuleInstance *module_inst, char *exception_buf)
     bool has_exception = false;
 
 #if WASM_ENABLE_SHARED_MEMORY != 0
-    WASMSharedMemNode *node =
-        wasm_module_get_shared_memory((WASMModuleCommon *)module_inst->module);
-    if (node)
-        os_mutex_lock(&node->shared_mem_lock);
+    if (module_inst->memory_count > 0)
+        shared_memory_lock(module_inst->memories[0]);
 #endif
     if (module_inst->cur_exception[0] != '\0') {
         /* NULL is passed if the caller is not interested in getting the
@@ -2403,8 +2400,8 @@ wasm_copy_exception(WASMModuleInstance *module_inst, char *exception_buf)
         has_exception = true;
     }
 #if WASM_ENABLE_SHARED_MEMORY != 0
-    if (node)
-        os_mutex_unlock(&node->shared_mem_lock);
+    if (module_inst->memory_count > 0)
+        shared_memory_unlock(module_inst->memories[0]);
 #endif
 
     return has_exception;

core/iwasm/common/wasm_runtime_common.h

@@ -498,7 +498,8 @@ wasm_runtime_unload(WASMModuleCommon *module);
 
 /* Internal API */
 WASMModuleInstanceCommon *
-wasm_runtime_instantiate_internal(WASMModuleCommon *module, bool is_sub_inst,
+wasm_runtime_instantiate_internal(WASMModuleCommon *module,
+                                  WASMModuleInstanceCommon *parent,
                                   WASMExecEnv *exec_env_main, uint32 stack_size,
                                   uint32 heap_size, char *error_buf,
                                   uint32 error_buf_size);

core/iwasm/common/wasm_shared_memory.c

@@ -9,9 +9,16 @@
 #include "../libraries/thread-mgr/thread_manager.h"
 #endif
 
-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;
+/*
+ * Note: this lock can be per memory.
+ *
+ * For now, just use a global because:
+ * - it's a bit cumbersome to extend WASMMemoryInstance w/o breaking
+ *   the AOT ABI.
+ * - If you care performance, it's better to make the interpreters
+ *   use atomic ops.
+ */
+static korp_mutex _shared_memory_lock;
 
 /* clang-format off */
 enum {
@@ -48,17 +55,15 @@ destroy_wait_info(void *wait_info);
 bool
 wasm_shared_memory_init()
 {
-    if (os_mutex_init(&shared_memory_list_lock) != 0)
+    if (os_mutex_init(&_shared_memory_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,
                                         destroy_wait_info))) {
-        os_mutex_destroy(&shared_memory_list_lock);
+        os_mutex_destroy(&_shared_memory_lock);
         return false;
     }
-
     return true;
 }
 
@@ -66,110 +71,79 @@ void
 wasm_shared_memory_destroy()
 {
     bh_hash_map_destroy(wait_map);
-    os_mutex_destroy(&shared_memory_list_lock);
+    os_mutex_destroy(&_shared_memory_lock);
 }
 
-static WASMSharedMemNode *
-search_module(WASMModuleCommon *module)
+uint32
+shared_memory_inc_reference(WASMMemoryInstance *memory)
 {
-    WASMSharedMemNode *node;
-
-    os_mutex_lock(&shared_memory_list_lock);
-    node = bh_list_first_elem(shared_memory_list);
-
-    while (node) {
-        if (module == node->module) {
-            os_mutex_unlock(&shared_memory_list_lock);
-            return node;
-        }
-        node = bh_list_elem_next(node);
-    }
-
-    os_mutex_unlock(&shared_memory_list_lock);
-    return NULL;
+    bh_assert(shared_memory_is_shared(memory));
+    uint32 old;
+#if BH_ATOMIC_32_IS_ATOMIC == 0
+    os_mutex_lock(&_shared_memory_lock);
+#endif
+    old = BH_ATOMIC_32_FETCH_ADD(memory->ref_count, 1);
+#if BH_ATOMIC_32_IS_ATOMIC == 0
+    os_mutex_unlock(&_shared_memory_lock);
+#endif
+    bh_assert(old >= 1);
+    bh_assert(old < UINT32_MAX);
+    return old + 1;
 }
 
-WASMSharedMemNode *
-wasm_module_get_shared_memory(WASMModuleCommon *module)
+uint32
+shared_memory_dec_reference(WASMMemoryInstance *memory)
 {
-    return search_module(module);
+    bh_assert(shared_memory_is_shared(memory));
+    uint32 old;
+#if BH_ATOMIC_32_IS_ATOMIC == 0
+    os_mutex_lock(&_shared_memory_lock);
+#endif
+    old = BH_ATOMIC_32_FETCH_SUB(memory->ref_count, 1);
+#if BH_ATOMIC_32_IS_ATOMIC == 0
+    os_mutex_unlock(&_shared_memory_lock);
+#endif
+    bh_assert(old > 0);
+    return old - 1;
 }
 
-int32
-shared_memory_inc_reference(WASMModuleCommon *module)
+bool
+shared_memory_is_shared(WASMMemoryInstance *memory)
 {
-    WASMSharedMemNode *node = search_module(module);
-    uint32 ref_count = -1;
-    if (node) {
-        os_mutex_lock(&node->lock);
-        ref_count = ++node->ref_count;
-        os_mutex_unlock(&node->lock);
-    }
-    return ref_count;
+    uint32 old;
+#if BH_ATOMIC_32_IS_ATOMIC == 0
+    os_mutex_lock(&_shared_memory_lock);
+#endif
+    old = BH_ATOMIC_32_LOAD(memory->ref_count);
+#if BH_ATOMIC_32_IS_ATOMIC == 0
+    os_mutex_unlock(&_shared_memory_lock);
+#endif
+    return old > 0;
 }
 
-int32
-shared_memory_dec_reference(WASMModuleCommon *module)
+static korp_mutex *
+shared_memory_get_lock_pointer(WASMMemoryInstance *memory)
 {
-    WASMSharedMemNode *node = search_module(module);
-    uint32 ref_count = 0;
-    if (node) {
-        os_mutex_lock(&node->lock);
-        ref_count = --node->ref_count;
-        os_mutex_unlock(&node->lock);
-        if (ref_count == 0) {
-            os_mutex_lock(&shared_memory_list_lock);
-            bh_list_remove(shared_memory_list, node);
-            os_mutex_unlock(&shared_memory_list_lock);
-
-            os_mutex_destroy(&node->shared_mem_lock);
-            os_mutex_destroy(&node->lock);
-            wasm_runtime_free(node);
-        }
-        return ref_count;
-    }
-
-    return -1;
+    bh_assert(memory != NULL);
+    return &_shared_memory_lock;
 }
 
-WASMMemoryInstanceCommon *
-shared_memory_get_memory_inst(WASMSharedMemNode *node)
+void
+shared_memory_lock(WASMMemoryInstance *memory)
 {
-    return node->memory_inst;
+    /*
+     * Note: exception logic is currently abusing this lock.
+     * cf. https://github.com/bytecodealliance/wasm-micro-runtime/issues/2407
+     */
+    bh_assert(memory != NULL);
+    os_mutex_lock(&_shared_memory_lock);
 }
 
-WASMSharedMemNode *
-shared_memory_set_memory_inst(WASMModuleCommon *module,
-                              WASMMemoryInstanceCommon *memory)
+void
+shared_memory_unlock(WASMMemoryInstance *memory)
 {
-    WASMSharedMemNode *node;
-    bh_list_status ret;
-
-    if (!(node = wasm_runtime_malloc(sizeof(WASMSharedMemNode))))
-        return NULL;
-
-    node->module = module;
-    node->memory_inst = memory;
-    node->ref_count = 1;
-
-    if (os_mutex_init(&node->shared_mem_lock) != 0) {
-        wasm_runtime_free(node);
-        return NULL;
-    }
-
-    if (os_mutex_init(&node->lock) != 0) {
-        os_mutex_destroy(&node->shared_mem_lock);
-        wasm_runtime_free(node);
-        return NULL;
-    }
-
-    os_mutex_lock(&shared_memory_list_lock);
-    ret = bh_list_insert(shared_memory_list, node);
-    bh_assert(ret == BH_LIST_SUCCESS);
-    os_mutex_unlock(&shared_memory_list_lock);
-
-    (void)ret;
-    return node;
+    bh_assert(memory != NULL);
+    os_mutex_unlock(&_shared_memory_lock);
 }
 
 /* Atomics wait && notify APIs */
@@ -307,7 +281,7 @@ wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,
     WASMModuleInstance *module_inst = (WASMModuleInstance *)module;
     AtomicWaitInfo *wait_info;
     AtomicWaitNode *wait_node;
-    WASMSharedMemNode *node;
+    korp_mutex *lock;
 #if WASM_ENABLE_THREAD_MGR != 0
     WASMExecEnv *exec_env;
 #endif
@@ -322,7 +296,7 @@ wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,
     }
 
     /* Currently we have only one memory instance */
-    if (!module_inst->memories[0]->is_shared) {
+    if (!shared_memory_is_shared(module_inst->memories[0])) {
         wasm_runtime_set_exception(module, "expected shared memory");
         return -1;
     }
@@ -340,30 +314,29 @@ wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,
     bh_assert(exec_env);
 #endif
 
-    node = search_module((WASMModuleCommon *)module_inst->module);
-    bh_assert(node);
+    lock = shared_memory_get_lock_pointer(module_inst->memories[0]);
 
     /* Lock the shared_mem_lock for the whole atomic wait process,
        and use it to os_cond_reltimedwait */
-    os_mutex_lock(&node->shared_mem_lock);
+    os_mutex_lock(lock);
 
     no_wait = (!wait64 && *(uint32 *)address != (uint32)expect)
               || (wait64 && *(uint64 *)address != expect);
 
     if (no_wait) {
-        os_mutex_unlock(&node->shared_mem_lock);
+        os_mutex_unlock(lock);
         return 1;
     }
 
     if (!(wait_node = wasm_runtime_malloc(sizeof(AtomicWaitNode)))) {
-        os_mutex_unlock(&node->shared_mem_lock);
+        os_mutex_unlock(lock);
         wasm_runtime_set_exception(module, "failed to create wait node");
         return -1;
     }
     memset(wait_node, 0, sizeof(AtomicWaitNode));
 
     if (0 != os_cond_init(&wait_node->wait_cond)) {
-        os_mutex_unlock(&node->shared_mem_lock);
+        os_mutex_unlock(lock);
         wasm_runtime_free(wait_node);
         wasm_runtime_set_exception(module, "failed to init wait cond");
         return -1;
@@ -375,7 +348,7 @@ wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,
     wait_info = acquire_wait_info(address, wait_node);
 
     if (!wait_info) {
-        os_mutex_unlock(&node->shared_mem_lock);
+        os_mutex_unlock(lock);
         os_cond_destroy(&wait_node->wait_cond);
         wasm_runtime_free(wait_node);
         wasm_runtime_set_exception(module, "failed to acquire wait_info");
@@ -390,7 +363,7 @@ wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,
         if (timeout < 0) {
             /* wait forever until it is notified or terminatied
                here we keep waiting and checking every second */
-            os_cond_reltimedwait(&wait_node->wait_cond, &node->shared_mem_lock,
+            os_cond_reltimedwait(&wait_node->wait_cond, lock,
                                  (uint64)timeout_1sec);
             if (wait_node->status == S_NOTIFIED /* notified by atomic.notify */
 #if WASM_ENABLE_THREAD_MGR != 0
@@ -404,8 +377,7 @@ wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,
         else {
             timeout_wait =
                 timeout_left < timeout_1sec ? timeout_left : timeout_1sec;
-            os_cond_reltimedwait(&wait_node->wait_cond, &node->shared_mem_lock,
-                                 timeout_wait);
+            os_cond_reltimedwait(&wait_node->wait_cond, lock, timeout_wait);
             if (wait_node->status == S_NOTIFIED /* notified by atomic.notify */
                 || timeout_left <= timeout_wait /* time out */
 #if WASM_ENABLE_THREAD_MGR != 0
@@ -433,7 +405,7 @@ wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,
     /* Release wait info if no wait nodes are attached */
     map_try_release_wait_info(wait_map, wait_info, address);
 
-    os_mutex_unlock(&node->shared_mem_lock);
+    os_mutex_unlock(lock);
 
     return is_timeout ? 2 : 0;
 }
@@ -445,7 +417,7 @@ wasm_runtime_atomic_notify(WASMModuleInstanceCommon *module, void *address,
     WASMModuleInstance *module_inst = (WASMModuleInstance *)module;
     uint32 notify_result;
     AtomicWaitInfo *wait_info;
-    WASMSharedMemNode *node;
+    korp_mutex *lock;
     bool out_of_bounds;
 
     bh_assert(module->module_type == Wasm_Module_Bytecode
@@ -461,31 +433,30 @@ wasm_runtime_atomic_notify(WASMModuleInstanceCommon *module, void *address,
     }
 
     /* Currently we have only one memory instance */
-    if (!module_inst->memories[0]->is_shared) {
+    if (!shared_memory_is_shared(module_inst->memories[0])) {
         /* Always return 0 for ushared linear memory since there is
            no way to create a waiter on it */
         return 0;
     }
 
-    node = search_module((WASMModuleCommon *)module_inst->module);
-    bh_assert(node);
+    lock = shared_memory_get_lock_pointer(module_inst->memories[0]);
 
     /* Lock the shared_mem_lock for the whole atomic notify process,
        and use it to os_cond_signal */
-    os_mutex_lock(&node->shared_mem_lock);
+    os_mutex_lock(lock);
 
     wait_info = acquire_wait_info(address, NULL);
 
     /* Nobody wait on this address */
     if (!wait_info) {
-        os_mutex_unlock(&node->shared_mem_lock);
+        os_mutex_unlock(lock);
         return 0;
     }
 
     /* Notify each wait node in the wait list */
     notify_result = notify_wait_list(wait_info->wait_list, count);
 
-    os_mutex_unlock(&node->shared_mem_lock);
+    os_mutex_unlock(lock);
 
     return notify_result;
 }

core/iwasm/common/wasm_shared_memory.h

@@ -7,53 +7,33 @@
 #define _WASM_SHARED_MEMORY_H
 
 #include "bh_common.h"
-#if WASM_ENABLE_INTERP != 0
-#include "wasm_runtime.h"
-#endif
-#if WASM_ENABLE_AOT != 0
-#include "aot_runtime.h"
-#endif
+#include "../interpreter/wasm_runtime.h"
+#include "wasm_runtime_common.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-typedef struct WASMSharedMemNode {
-    bh_list_link l;
-    /* Lock */
-    korp_mutex lock;
-    /* The module reference */
-    WASMModuleCommon *module;
-    /* The memory information */
-    WASMMemoryInstanceCommon *memory_inst;
-    /* Lock used for atomic operations */
-    korp_mutex shared_mem_lock;
-
-    /* reference count */
-    uint32 ref_count;
-} WASMSharedMemNode;
-
 bool
 wasm_shared_memory_init();
 
 void
 wasm_shared_memory_destroy();
 
-WASMSharedMemNode *
-wasm_module_get_shared_memory(WASMModuleCommon *module);
+uint32
+shared_memory_inc_reference(WASMMemoryInstance *memory);
 
-int32
-shared_memory_inc_reference(WASMModuleCommon *module);
+uint32
+shared_memory_dec_reference(WASMMemoryInstance *memory);
 
-int32
-shared_memory_dec_reference(WASMModuleCommon *module);
+bool
+shared_memory_is_shared(WASMMemoryInstance *memory);
 
-WASMMemoryInstanceCommon *
-shared_memory_get_memory_inst(WASMSharedMemNode *node);
+void
+shared_memory_lock(WASMMemoryInstance *memory);
 
-WASMSharedMemNode *
-shared_memory_set_memory_inst(WASMModuleCommon *module,
-                              WASMMemoryInstanceCommon *memory);
+void
+shared_memory_unlock(WASMMemoryInstance *memory);
 
 uint32
 wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,

core/iwasm/interpreter/wasm_interp_classic.c

@@ -710,28 +710,28 @@ trunc_f64_to_int(WASMModuleInstance *module, uint32 *frame_sp, float64 src_min,
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS();                            \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint32)(*(uint8 *)maddr);                       \
             *(uint8 *)maddr = (uint8)(readv op sval);                \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else if (opcode == WASM_OP_ATOMIC_RMW_I32_##OP_NAME##16_U) { \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS();                            \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint32)LOAD_U16(maddr);                         \
             STORE_U16(maddr, (uint16)(readv op sval));               \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else {                                                       \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS();                            \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = LOAD_I32(maddr);                                 \
             STORE_U32(maddr, readv op sval);                         \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         PUSH_I32(readv);                                             \
         break;                                                       \
@@ -750,39 +750,39 @@ trunc_f64_to_int(WASMModuleInstance *module, uint32 *frame_sp, float64 src_min,
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS();                            \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint64)(*(uint8 *)maddr);                       \
             *(uint8 *)maddr = (uint8)(readv op sval);                \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else if (opcode == WASM_OP_ATOMIC_RMW_I64_##OP_NAME##16_U) { \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS();                            \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint64)LOAD_U16(maddr);                         \
             STORE_U16(maddr, (uint16)(readv op sval));               \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else if (opcode == WASM_OP_ATOMIC_RMW_I64_##OP_NAME##32_U) { \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS();                            \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint64)LOAD_U32(maddr);                         \
             STORE_U32(maddr, (uint32)(readv op sval));               \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else {                                                       \
             uint64 op_result;                                        \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 8, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS();                            \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint64)LOAD_I64(maddr);                         \
             op_result = readv op sval;                               \
             STORE_I64(maddr, op_result);                             \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         PUSH_I64(readv);                                             \
         break;                                                       \
@@ -1156,10 +1156,6 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                                WASMFunctionInstance *cur_func,
                                WASMInterpFrame *prev_frame)
 {
-#if WASM_ENABLE_SHARED_MEMORY != 0
-    WASMSharedMemNode *node =
-        wasm_module_get_shared_memory((WASMModuleCommon *)module->module);
-#endif
     WASMMemoryInstance *memory = wasm_get_default_memory(module);
 #if !defined(OS_ENABLE_HW_BOUND_CHECK)              \
     || WASM_CPU_SUPPORTS_UNALIGNED_ADDR_ACCESS == 0 \
@@ -3526,23 +3522,23 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                         if (opcode == WASM_OP_ATOMIC_I32_LOAD8_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint32)(*(uint8 *)maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I32_LOAD16_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint32)LOAD_U16(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = LOAD_I32(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
 
                         PUSH_I32(readv);
@@ -3561,30 +3557,30 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                         if (opcode == WASM_OP_ATOMIC_I64_LOAD8_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)(*(uint8 *)maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I64_LOAD16_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_U16(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I64_LOAD32_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_U32(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 8, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = LOAD_I64(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
 
                         PUSH_I64(readv);
@@ -3603,23 +3599,23 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                         if (opcode == WASM_OP_ATOMIC_I32_STORE8) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             *(uint8 *)maddr = (uint8)sval;
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I32_STORE16) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             STORE_U16(maddr, (uint16)sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             STORE_U32(maddr, sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         break;
                     }
@@ -3637,30 +3633,30 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                         if (opcode == WASM_OP_ATOMIC_I64_STORE8) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             *(uint8 *)maddr = (uint8)sval;
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I64_STORE16) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             STORE_U16(maddr, (uint16)sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I64_STORE32) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             STORE_U32(maddr, (uint32)sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 8, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             PUT_I64_TO_ADDR((uint32 *)maddr, sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         break;
                     }
@@ -3680,32 +3676,32 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                             CHECK_ATOMIC_MEMORY_ACCESS();
 
                             expect = (uint8)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint32)(*(uint8 *)maddr);
                             if (readv == expect)
                                 *(uint8 *)maddr = (uint8)(sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_RMW_I32_CMPXCHG16_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
 
                             expect = (uint16)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint32)LOAD_U16(maddr);
                             if (readv == expect)
                                 STORE_U16(maddr, (uint16)(sval));
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
 
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = LOAD_I32(maddr);
                             if (readv == expect)
                                 STORE_U32(maddr, sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         PUSH_I32(readv);
                         break;
@@ -3726,43 +3722,43 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                             CHECK_ATOMIC_MEMORY_ACCESS();
 
                             expect = (uint8)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)(*(uint8 *)maddr);
                             if (readv == expect)
                                 *(uint8 *)maddr = (uint8)(sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_RMW_I64_CMPXCHG16_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
 
                             expect = (uint16)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_U16(maddr);
                             if (readv == expect)
                                 STORE_U16(maddr, (uint16)(sval));
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_RMW_I64_CMPXCHG32_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
 
                             expect = (uint32)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_U32(maddr);
                             if (readv == expect)
                                 STORE_U32(maddr, (uint32)(sval));
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 8, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS();
 
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_I64(maddr);
                             if (readv == expect)
                                 STORE_I64(maddr, sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         PUSH_I64(readv);
                         break;

core/iwasm/interpreter/wasm_interp_fast.c

@@ -482,28 +482,28 @@ LOAD_PTR(void *addr)
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS(1);                           \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint32)(*(uint8 *)maddr);                       \
             *(uint8 *)maddr = (uint8)(readv op sval);                \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else if (opcode == WASM_OP_ATOMIC_RMW_I32_##OP_NAME##16_U) { \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS(2);                           \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint32)LOAD_U16(maddr);                         \
             STORE_U16(maddr, (uint16)(readv op sval));               \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else {                                                       \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS(4);                           \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = LOAD_I32(maddr);                                 \
             STORE_U32(maddr, readv op sval);                         \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         PUSH_I32(readv);                                             \
         break;                                                       \
@@ -522,39 +522,39 @@ LOAD_PTR(void *addr)
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS(1);                           \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint64)(*(uint8 *)maddr);                       \
             *(uint8 *)maddr = (uint8)(readv op sval);                \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else if (opcode == WASM_OP_ATOMIC_RMW_I64_##OP_NAME##16_U) { \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS(2);                           \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint64)LOAD_U16(maddr);                         \
             STORE_U16(maddr, (uint16)(readv op sval));               \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else if (opcode == WASM_OP_ATOMIC_RMW_I64_##OP_NAME##32_U) { \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS(4);                           \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint64)LOAD_U32(maddr);                         \
             STORE_U32(maddr, (uint32)(readv op sval));               \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         else {                                                       \
             uint64 op_result;                                        \
             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 8, maddr);     \
             CHECK_ATOMIC_MEMORY_ACCESS(8);                           \
                                                                      \
-            os_mutex_lock(&node->shared_mem_lock);                   \
+            shared_memory_lock(memory);                              \
             readv = (uint64)LOAD_I64(maddr);                         \
             op_result = readv op sval;                               \
             STORE_I64(maddr, op_result);                             \
-            os_mutex_unlock(&node->shared_mem_lock);                 \
+            shared_memory_unlock(memory);                            \
         }                                                            \
         PUSH_I64(readv);                                             \
         break;                                                       \
@@ -1166,10 +1166,6 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                                WASMFunctionInstance *cur_func,
                                WASMInterpFrame *prev_frame)
 {
-#if WASM_ENABLE_SHARED_MEMORY != 0
-    WASMSharedMemNode *node =
-        wasm_module_get_shared_memory((WASMModuleCommon *)module->module);
-#endif
     WASMMemoryInstance *memory = wasm_get_default_memory(module);
 #if !defined(OS_ENABLE_HW_BOUND_CHECK)              \
     || WASM_CPU_SUPPORTS_UNALIGNED_ADDR_ACCESS == 0 \
@@ -3353,23 +3349,23 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                         if (opcode == WASM_OP_ATOMIC_I32_LOAD8_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(1);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint32)(*(uint8 *)maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I32_LOAD16_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(2);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint32)LOAD_U16(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(4);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = LOAD_I32(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
 
                         PUSH_I32(readv);
@@ -3388,30 +3384,30 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                         if (opcode == WASM_OP_ATOMIC_I64_LOAD8_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(1);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)(*(uint8 *)maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I64_LOAD16_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(2);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_U16(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I64_LOAD32_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(4);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_U32(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 8, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(8);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = LOAD_I64(maddr);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
 
                         PUSH_I64(readv);
@@ -3429,23 +3425,23 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                         if (opcode == WASM_OP_ATOMIC_I32_STORE8) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(1);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             *(uint8 *)maddr = (uint8)sval;
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I32_STORE16) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(2);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             STORE_U16(maddr, (uint16)sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(4);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             STORE_U32(maddr, sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         break;
                     }
@@ -3463,30 +3459,30 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                         if (opcode == WASM_OP_ATOMIC_I64_STORE8) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 1, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(1);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             *(uint8 *)maddr = (uint8)sval;
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I64_STORE16) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(2);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             STORE_U16(maddr, (uint16)sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_I64_STORE32) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(4);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             STORE_U32(maddr, (uint32)sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 8, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(8);
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             STORE_I64(maddr, sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         break;
                     }
@@ -3506,32 +3502,32 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                             CHECK_ATOMIC_MEMORY_ACCESS(1);
 
                             expect = (uint8)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint32)(*(uint8 *)maddr);
                             if (readv == expect)
                                 *(uint8 *)maddr = (uint8)(sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_RMW_I32_CMPXCHG16_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(2);
 
                             expect = (uint16)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint32)LOAD_U16(maddr);
                             if (readv == expect)
                                 STORE_U16(maddr, (uint16)(sval));
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(4);
 
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = LOAD_I32(maddr);
                             if (readv == expect)
                                 STORE_U32(maddr, sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         PUSH_I32(readv);
                         break;
@@ -3552,43 +3548,43 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                             CHECK_ATOMIC_MEMORY_ACCESS(1);
 
                             expect = (uint8)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)(*(uint8 *)maddr);
                             if (readv == expect)
                                 *(uint8 *)maddr = (uint8)(sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_RMW_I64_CMPXCHG16_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 2, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(2);
 
                             expect = (uint16)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_U16(maddr);
                             if (readv == expect)
                                 STORE_U16(maddr, (uint16)(sval));
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else if (opcode == WASM_OP_ATOMIC_RMW_I64_CMPXCHG32_U) {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 4, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(4);
 
                             expect = (uint32)expect;
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_U32(maddr);
                             if (readv == expect)
                                 STORE_U32(maddr, (uint32)(sval));
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         else {
                             CHECK_BULK_MEMORY_OVERFLOW(addr + offset, 8, maddr);
                             CHECK_ATOMIC_MEMORY_ACCESS(8);
 
-                            os_mutex_lock(&node->shared_mem_lock);
+                            shared_memory_lock(memory);
                             readv = (uint64)LOAD_I64(maddr);
                             if (readv == expect)
                                 STORE_I64(maddr, sval);
-                            os_mutex_unlock(&node->shared_mem_lock);
+                            shared_memory_unlock(memory);
                         }
                         PUSH_I64(readv);
                         break;

core/iwasm/interpreter/wasm_runtime.c

@@ -122,11 +122,8 @@ memories_deinstantiate(WASMModuleInstance *module_inst,
                 }
 #endif
 #if WASM_ENABLE_SHARED_MEMORY != 0
-                if (memories[i]->is_shared) {
-                    int32 ref_count = shared_memory_dec_reference(
-                        (WASMModuleCommon *)module_inst->module);
-                    bh_assert(ref_count >= 0);
-
+                if (shared_memory_is_shared(memories[i])) {
+                    uint32 ref_count = shared_memory_dec_reference(memories[i]);
                     /* if the reference count is not zero,
                         don't free the memory */
                     if (ref_count > 0)
@@ -159,7 +156,8 @@ memories_deinstantiate(WASMModuleInstance *module_inst,
 }
 
 static WASMMemoryInstance *
-memory_instantiate(WASMModuleInstance *module_inst, WASMMemoryInstance *memory,
+memory_instantiate(WASMModuleInstance *module_inst, WASMModuleInstance *parent,
+                   WASMMemoryInstance *memory, uint32 memory_idx,
                    uint32 num_bytes_per_page, uint32 init_page_count,
                    uint32 max_page_count, uint32 heap_size, uint32 flags,
                    char *error_buf, uint32 error_buf_size)
@@ -180,22 +178,11 @@ memory_instantiate(WASMModuleInstance *module_inst, WASMMemoryInstance *memory,
     bool is_shared_memory = flags & 0x02 ? true : false;
 
     /* shared memory */
-    if (is_shared_memory) {
-        WASMSharedMemNode *node = wasm_module_get_shared_memory(
-            (WASMModuleCommon *)module_inst->module);
-        /* If the memory of this module has been instantiated,
-            return the memory instance directly */
-        if (node) {
-            uint32 ref_count;
-            ref_count = shared_memory_inc_reference(
-                (WASMModuleCommon *)module_inst->module);
-            bh_assert(ref_count > 0);
-            memory = (WASMMemoryInstance *)shared_memory_get_memory_inst(node);
-            bh_assert(memory);
-
-            (void)ref_count;
-            return memory;
-        }
+    if (is_shared_memory && parent != NULL) {
+        bh_assert(parent->memory_count > memory_idx);
+        memory = parent->memories[memory_idx];
+        shared_memory_inc_reference(memory);
+        return memory;
     }
 #endif /* end of WASM_ENABLE_SHARED_MEMORY */
 
@@ -388,24 +375,13 @@ memory_instantiate(WASMModuleInstance *module_inst, WASMMemoryInstance *memory,
 
 #if WASM_ENABLE_SHARED_MEMORY != 0
     if (is_shared_memory) {
-        memory->is_shared = true;
-        if (!shared_memory_set_memory_inst(
-                (WASMModuleCommon *)module_inst->module,
-                (WASMMemoryInstanceCommon *)memory)) {
-            set_error_buf(error_buf, error_buf_size, "allocate memory failed");
-            goto fail4;
-        }
+        memory->ref_count = 1;
     }
 #endif
 
     LOG_VERBOSE("Memory instantiate success.");
     return memory;
 
-#if WASM_ENABLE_SHARED_MEMORY != 0
-fail4:
-    if (heap_size > 0)
-        mem_allocator_destroy(memory->heap_handle);
-#endif
 fail3:
     if (heap_size > 0)
         wasm_runtime_free(memory->heap_handle);
@@ -428,7 +404,8 @@ fail1:
  */
 static WASMMemoryInstance **
 memories_instantiate(const WASMModule *module, WASMModuleInstance *module_inst,
-                     uint32 heap_size, char *error_buf, uint32 error_buf_size)
+                     WASMModuleInstance *parent, uint32 heap_size,
+                     char *error_buf, uint32 error_buf_size)
 {
     WASMImport *import;
     uint32 mem_index = 0, i,
@@ -474,26 +451,29 @@ memories_instantiate(const WASMModule *module, WASMModuleInstance *module_inst,
         else
 #endif
         {
-            if (!(memories[mem_index++] = memory_instantiate(
-                      module_inst, memory, num_bytes_per_page, init_page_count,
-                      max_page_count, actual_heap_size, flags, error_buf,
-                      error_buf_size))) {
+            if (!(memories[mem_index] = memory_instantiate(
+                      module_inst, parent, memory, mem_index,
+                      num_bytes_per_page, init_page_count, max_page_count,
+                      actual_heap_size, flags, error_buf, error_buf_size))) {
                 memories_deinstantiate(module_inst, memories, memory_count);
                 return NULL;
             }
+            mem_index++;
         }
     }
 
     /* instantiate memories from memory section */
     for (i = 0; i < module->memory_count; i++, memory++) {
-        if (!(memories[mem_index++] = memory_instantiate(
-                  module_inst, memory, module->memories[i].num_bytes_per_page,
+        if (!(memories[mem_index] = memory_instantiate(
+                  module_inst, parent, memory, mem_index,
+                  module->memories[i].num_bytes_per_page,
                   module->memories[i].init_page_count,
                   module->memories[i].max_page_count, heap_size,
                   module->memories[i].flags, error_buf, error_buf_size))) {
             memories_deinstantiate(module_inst, memories, memory_count);
             return NULL;
         }
+        mem_index++;
     }
 
     bh_assert(mem_index == memory_count);
@@ -1301,7 +1281,7 @@ sub_module_instantiate(WASMModule *module, WASMModuleInstance *module_inst,
         WASMModuleInstance *sub_module_inst = NULL;
 
         sub_module_inst =
-            wasm_instantiate(sub_module, false, NULL, stack_size, heap_size,
+            wasm_instantiate(sub_module, NULL, NULL, stack_size, heap_size,
                              error_buf, error_buf_size);
         if (!sub_module_inst) {
             LOG_DEBUG("instantiate %s failed",
@@ -1646,7 +1626,7 @@ wasm_set_running_mode(WASMModuleInstance *module_inst, RunningMode running_mode)
  * Instantiate module
  */
 WASMModuleInstance *
-wasm_instantiate(WASMModule *module, bool is_sub_inst,
+wasm_instantiate(WASMModule *module, WASMModuleInstance *parent,
                  WASMExecEnv *exec_env_main, uint32 stack_size,
                  uint32 heap_size, char *error_buf, uint32 error_buf_size)
 {
@@ -1663,6 +1643,7 @@ wasm_instantiate(WASMModule *module, bool is_sub_inst,
 #if WASM_ENABLE_MULTI_MODULE != 0
     bool ret = false;
 #endif
+    const bool is_sub_inst = parent != NULL;
 
     if (!module)
         return NULL;
@@ -1781,8 +1762,9 @@ wasm_instantiate(WASMModule *module, bool is_sub_inst,
 
     /* Instantiate memories/tables/functions */
     if ((module_inst->memory_count > 0
-         && !(module_inst->memories = memories_instantiate(
-                  module, module_inst, heap_size, error_buf, error_buf_size)))
+         && !(module_inst->memories =
+                  memories_instantiate(module, module_inst, parent, heap_size,
+                                       error_buf, error_buf_size)))
         || (module_inst->table_count > 0
             && !(module_inst->tables =
                      tables_instantiate(module, module_inst, first_table,

core/iwasm/interpreter/wasm_runtime.h

@@ -7,6 +7,7 @@
 #define _WASM_RUNTIME_H
 
 #include "wasm.h"
+#include "bh_atomic.h"
 #include "bh_hashmap.h"
 #include "../common/wasm_runtime_common.h"
 #include "../common/wasm_exec_env.h"
@@ -79,7 +80,7 @@ struct WASMMemoryInstance {
     /* Module type */
     uint32 module_type;
     /* Shared memory flag */
-    bool is_shared;
+    bh_atomic_32_t ref_count; /* 0: non-shared, > 0: reference count */
 
     /* Number bytes per page */
     uint32 num_bytes_per_page;
@@ -400,7 +401,7 @@ void
 wasm_unload(WASMModule *module);
 
 WASMModuleInstance *
-wasm_instantiate(WASMModule *module, bool is_sub_inst,
+wasm_instantiate(WASMModule *module, WASMModuleInstance *parent,
                  WASMExecEnv *exec_env_main, uint32 stack_size,
                  uint32 heap_size, char *error_buf, uint32 error_buf_size);
 

core/iwasm/libraries/lib-pthread/lib_pthread_wrapper.c

@@ -581,7 +581,7 @@ pthread_create_wrapper(wasm_exec_env_t exec_env,
 #endif
 
     if (!(new_module_inst = wasm_runtime_instantiate_internal(
-              module, true, exec_env, stack_size, 0, NULL, 0)))
+              module, module_inst, exec_env, stack_size, 0, NULL, 0)))
         return -1;
 
     /* Set custom_data to new module instance */

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

@@ -90,7 +90,7 @@ thread_spawn_wrapper(wasm_exec_env_t exec_env, uint32 start_arg)
     stack_size = ((WASMModuleInstance *)module_inst)->default_wasm_stack_size;
 
     if (!(new_module_inst = wasm_runtime_instantiate_internal(
-              module, true, exec_env, stack_size, 0, NULL, 0)))
+              module, module_inst, exec_env, stack_size, 0, NULL, 0)))
         return -1;
 
     wasm_runtime_set_custom_data_internal(

core/iwasm/libraries/thread-mgr/thread_manager.c

@@ -509,7 +509,7 @@ wasm_cluster_spawn_exec_env(WASMExecEnv *exec_env)
 #endif
 
     if (!(new_module_inst = wasm_runtime_instantiate_internal(
-              module, true, exec_env, stack_size, 0, NULL, 0))) {
+              module, module_inst, exec_env, stack_size, 0, NULL, 0))) {
         goto fail1;
     }
 
@@ -1254,10 +1254,8 @@ set_exception_visitor(void *node, void *user_data)
 
         /* Only spread non "wasi proc exit" exception */
 #if WASM_ENABLE_SHARED_MEMORY != 0
-        WASMSharedMemNode *shared_mem_node = wasm_module_get_shared_memory(
-            (WASMModuleCommon *)curr_wasm_inst->module);
-        if (shared_mem_node)
-            os_mutex_lock(&shared_mem_node->shared_mem_lock);
+        if (curr_wasm_inst->memory_count > 0)
+            shared_memory_lock(curr_wasm_inst->memories[0]);
 #endif
         if (!strstr(wasm_inst->cur_exception, "wasi proc exit")) {
             bh_memcpy_s(curr_wasm_inst->cur_exception,
@@ -1266,8 +1264,8 @@ set_exception_visitor(void *node, void *user_data)
                         sizeof(wasm_inst->cur_exception));
         }
 #if WASM_ENABLE_SHARED_MEMORY != 0
-        if (shared_mem_node)
-            os_mutex_unlock(&shared_mem_node->shared_mem_lock);
+        if (curr_wasm_inst->memory_count > 0)
+            shared_memory_unlock(curr_wasm_inst->memories[0]);
 #endif
 
         /* Terminate the thread so it can exit from dead loops */
@@ -1286,15 +1284,13 @@ clear_exception_visitor(void *node, void *user_data)
             (WASMModuleInstance *)get_module_inst(curr_exec_env);
 
 #if WASM_ENABLE_SHARED_MEMORY != 0
-        WASMSharedMemNode *shared_mem_node = wasm_module_get_shared_memory(
-            (WASMModuleCommon *)curr_wasm_inst->module);
-        if (shared_mem_node)
-            os_mutex_lock(&shared_mem_node->shared_mem_lock);
+        if (curr_wasm_inst->memory_count > 0)
+            shared_memory_lock(curr_wasm_inst->memories[0]);
 #endif
         curr_wasm_inst->cur_exception[0] = '\0';
 #if WASM_ENABLE_SHARED_MEMORY != 0
-        if (shared_mem_node)
-            os_mutex_unlock(&shared_mem_node->shared_mem_lock);
+        if (curr_wasm_inst->memory_count > 0)
+            shared_memory_unlock(curr_wasm_inst->memories[0]);
 #endif
     }
 }