Просмотр исходного кода

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: https://github.com/bytecodealliance/wasm-micro-runtime/issues/4314
YAMAMOTO Takashi 7 месяцев назад
Родитель
Сommit
85efe08431
1 измененных файлов с 28 добавлено и 3 удалено
  1. 28 3
      core/iwasm/libraries/wasi-nn/src/wasi_nn.c

+ 28 - 3
core/iwasm/libraries/wasi-nn/src/wasi_nn.c

@@ -31,6 +31,14 @@
 #define LLAMACPP_BACKEND_LIB "libwasi_nn_llamacpp" LIB_EXTENTION
 
 /* 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 {
     void *backend_handle;
     api_function functions;
@@ -109,12 +117,18 @@ wasi_nn_initialize()
 {
     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 = bh_hash_map_create(HASHMAP_INITIAL_SIZE, true, hash_func,
                                  key_equal_func, key_destroy_func,
                                  value_destroy_func);
     if (hashmap == NULL) {
         NN_ERR_PRINTF("Error while initializing hashmap");
+        os_mutex_destroy(&wasi_nn_lock);
         return false;
     }
 
@@ -175,6 +189,8 @@ wasi_nn_destroy()
 
         memset(&lookup[i].functions, 0, sizeof(api_function));
     }
+
+    os_mutex_destroy(&wasi_nn_lock);
 }
 
 /* Utils */
@@ -349,6 +365,8 @@ static bool
 detect_and_load_backend(graph_encoding backend_hint,
                         graph_encoding *loaded_backend)
 {
+    bool ret;
+
     if (backend_hint > autodetect)
         return false;
 
@@ -360,16 +378,23 @@ detect_and_load_backend(graph_encoding backend_hint,
 
     *loaded_backend = backend_hint;
 
+    os_mutex_lock(&wasi_nn_lock);
     /* if already loaded */
-    if (lookup[backend_hint].backend_handle)
+    if (lookup[backend_hint].backend_handle) {
+        os_mutex_unlock(&wasi_nn_lock);
         return true;
+    }
 
     const char *backend_lib_name =
         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 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 */