Skip to content

Commit 85efe08

Browse files
authored
wasi-nn: protect the backend lookup table with a lock (#4319)
this would avoid potential issues when multiple instances happen to make an attempt to load a backend at the same time. Fixes: #4314
1 parent 93ef19b commit 85efe08

File tree

1 file changed

+28
-3
lines changed

1 file changed

+28
-3
lines changed

core/iwasm/libraries/wasi-nn/src/wasi_nn.c

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,14 @@
3131
#define LLAMACPP_BACKEND_LIB "libwasi_nn_llamacpp" LIB_EXTENTION
3232

3333
/* Global variables */
34+
static korp_mutex wasi_nn_lock;
35+
/*
36+
* the "lookup" table is protected by wasi_nn_lock.
37+
*
38+
* an exception: during wasm_runtime_destroy, wasi_nn_destroy tears down
39+
* the table without acquiring the lock. it's ok because there should be
40+
* no other threads using the runtime at this point.
41+
*/
3442
struct backends_api_functions {
3543
void *backend_handle;
3644
api_function functions;
@@ -109,12 +117,18 @@ wasi_nn_initialize()
109117
{
110118
NN_DBG_PRINTF("[WASI NN General] Initializing wasi-nn");
111119

120+
if (os_mutex_init(&wasi_nn_lock)) {
121+
NN_ERR_PRINTF("Error while initializing global lock");
122+
return false;
123+
}
124+
112125
// hashmap { instance: wasi_nn_ctx }
113126
hashmap = bh_hash_map_create(HASHMAP_INITIAL_SIZE, true, hash_func,
114127
key_equal_func, key_destroy_func,
115128
value_destroy_func);
116129
if (hashmap == NULL) {
117130
NN_ERR_PRINTF("Error while initializing hashmap");
131+
os_mutex_destroy(&wasi_nn_lock);
118132
return false;
119133
}
120134

@@ -175,6 +189,8 @@ wasi_nn_destroy()
175189

176190
memset(&lookup[i].functions, 0, sizeof(api_function));
177191
}
192+
193+
os_mutex_destroy(&wasi_nn_lock);
178194
}
179195

180196
/* Utils */
@@ -349,6 +365,8 @@ static bool
349365
detect_and_load_backend(graph_encoding backend_hint,
350366
graph_encoding *loaded_backend)
351367
{
368+
bool ret;
369+
352370
if (backend_hint > autodetect)
353371
return false;
354372

@@ -360,16 +378,23 @@ detect_and_load_backend(graph_encoding backend_hint,
360378

361379
*loaded_backend = backend_hint;
362380

381+
os_mutex_lock(&wasi_nn_lock);
363382
/* if already loaded */
364-
if (lookup[backend_hint].backend_handle)
383+
if (lookup[backend_hint].backend_handle) {
384+
os_mutex_unlock(&wasi_nn_lock);
365385
return true;
386+
}
366387

367388
const char *backend_lib_name =
368389
graph_encoding_to_backend_lib_name(backend_hint);
369-
if (!backend_lib_name)
390+
if (!backend_lib_name) {
391+
os_mutex_unlock(&wasi_nn_lock);
370392
return false;
393+
}
371394

372-
return prepare_backend(backend_lib_name, lookup + backend_hint);
395+
ret = prepare_backend(backend_lib_name, lookup + backend_hint);
396+
os_mutex_unlock(&wasi_nn_lock);
397+
return ret;
373398
}
374399

375400
/* WASI-NN implementation */

0 commit comments

Comments
 (0)