gnh1201/wasm-micro-runtime
1
0
Fork 0
mirror of https://github.com/bytecodealliance/wasm-micro-runtime.git synced 2025-06-18 02:59:21 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

wasi-nn: protect the backend lookup table with a lock (#4319)

Browse Source 
this would avoid potential issues when multiple instances happen to
make an attempt to load a backend at the same time.

Fixes: https://github.com/bytecodealliance/wasm-micro-runtime/issues/4314
This commit is contained in:
YAMAMOTO Takashi 2025-06-05 10:19:46 +09:00 committed by GitHub
parent 93ef19b0ca
commit 85efe08431
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 28 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
core/iwasm/libraries/wasi-nn/src/wasi_nn.c
View File

@ -31,6 +31,14 @@
#define LLAMACPP_BACKEND_LIB "libwasi_nn_llamacpp" LIB_EXTENTION #define LLAMACPP_BACKEND_LIB "libwasi_nn_llamacpp" LIB_EXTENTION
/* Global variables */ /* Global variables */
static korp_mutex wasi_nn_lock;
/*
* the "lookup" table is protected by wasi_nn_lock.
*
* an exception: during wasm_runtime_destroy, wasi_nn_destroy tears down
* the table without acquiring the lock. it's ok because there should be
* no other threads using the runtime at this point.
*/
struct backends_api_functions { struct backends_api_functions {
void *backend_handle; void *backend_handle;
api_function functions; api_function functions;
@ -109,12 +117,18 @@ wasi_nn_initialize()
{ {
NN_DBG_PRINTF("[WASI NN General] Initializing wasi-nn"); NN_DBG_PRINTF("[WASI NN General] Initializing wasi-nn");
if (os_mutex_init(&wasi_nn_lock)) {
NN_ERR_PRINTF("Error while initializing global lock");
return false;
}
// hashmap { instance: wasi_nn_ctx } // hashmap { instance: wasi_nn_ctx }
hashmap = bh_hash_map_create(HASHMAP_INITIAL_SIZE, true, hash_func, hashmap = bh_hash_map_create(HASHMAP_INITIAL_SIZE, true, hash_func,
key_equal_func, key_destroy_func, key_equal_func, key_destroy_func,
value_destroy_func); value_destroy_func);
if (hashmap == NULL) { if (hashmap == NULL) {
NN_ERR_PRINTF("Error while initializing hashmap"); NN_ERR_PRINTF("Error while initializing hashmap");
os_mutex_destroy(&wasi_nn_lock);
return false; return false;
} }
@ -175,6 +189,8 @@ wasi_nn_destroy()
memset(&lookup[i].functions, 0, sizeof(api_function)); memset(&lookup[i].functions, 0, sizeof(api_function));
} }
os_mutex_destroy(&wasi_nn_lock);
} }
/* Utils */ /* Utils */
@ -349,6 +365,8 @@ static bool
detect_and_load_backend(graph_encoding backend_hint, detect_and_load_backend(graph_encoding backend_hint,
graph_encoding *loaded_backend) graph_encoding *loaded_backend)
{ {
bool ret;
if (backend_hint > autodetect) if (backend_hint > autodetect)
return false; return false;
@ -360,16 +378,23 @@ detect_and_load_backend(graph_encoding backend_hint,
*loaded_backend = backend_hint; *loaded_backend = backend_hint;
os_mutex_lock(&wasi_nn_lock);
/* if already loaded */ /* if already loaded */
if (lookup[backend_hint].backend_handle) if (lookup[backend_hint].backend_handle) {
os_mutex_unlock(&wasi_nn_lock);
return true; return true;
}
const char *backend_lib_name = const char *backend_lib_name =
graph_encoding_to_backend_lib_name(backend_hint); graph_encoding_to_backend_lib_name(backend_hint);
if (!backend_lib_name) if (!backend_lib_name) {
os_mutex_unlock(&wasi_nn_lock);
return false; return false;
}
return prepare_backend(backend_lib_name, lookup + backend_hint); ret = prepare_backend(backend_lib_name, lookup + backend_hint);
os_mutex_unlock(&wasi_nn_lock);
return ret;
} }
/* WASI-NN implementation */ /* WASI-NN implementation */