[WASI Threads] Instance for threads

Signed-off-by: Sylveon <sylveon@secondstate.io>
WasmEdge · Jul 18, 2023 · 7294d18 · 7294d18
1 parent 961879b
commit 7294d18
Show file tree

Hide file tree

Showing 12 changed files with 135 additions and 77 deletions.
diff --git a/include/executor/executor.h b/include/executor/executor.h
@@ -117,6 +117,19 @@ class Executor {
   Expect<std::unique_ptr<Runtime::Instance::ModuleInstance>>
   instantiateModule(Runtime::StoreManager &StoreMgr, const AST::Module &Mod);
 
+  ///
+  void setHostCallback(
+      std::function<std::vector<
+          std::unique_ptr<WasmEdge::Runtime::Instance::ModuleInstance>>()>
+          Func);
+
+  /// Reinstantiate a WASM Module into an anonymous module instance.
+  Expect<std::unique_ptr<Runtime::Instance::ModuleInstance>>
+  reinstantiateModule(
+      Runtime::StoreManager &StoreMgr,
+      std::vector<std::unique_ptr<Runtime::Instance::ModuleInstance>>
+          &HostModsHolder);
+
   /// Instantiate and register a WASM module into a named module instance.
   Expect<std::unique_ptr<Runtime::Instance::ModuleInstance>>
   registerModule(Runtime::StoreManager &StoreMgr, const AST::Module &Mod,
@@ -690,6 +703,17 @@ class Executor {
   Statistics::Statistics *Stat;
   /// Stop Execution
   std::atomic_uint32_t StopToken = 0;
+
+private:
+  /// For reinstantiateModule on threads
+  struct ExecutionModuleSourceStruct {
+    const AST::Module *SourceMod = nullptr;
+    std::function<std::vector<
+        std::unique_ptr<WasmEdge::Runtime::Instance::ModuleInstance>>()>
+        HostCallback;
+  };
+
+  ExecutionModuleSourceStruct ModuleSourceStructContext;
 };
 
 } // namespace Executor

diff --git a/include/host/mock/wasi_threads_func.h b/include/host/mock/wasi_threads_func.h
@@ -16,7 +16,7 @@ using namespace std::literals;
 
 class ThreadSpawn : public Runtime::HostFunction<ThreadSpawn> {
 public:
-  Expect<void> body(const Runtime::CallingFrame &, uint32_t) {
+  Expect<void> body(const Runtime::CallingFrame &, int32_t) {
     printPluginMock("Wasi-Threads"sv);
     return Unexpect(ErrCode::Value::HostFuncError);
   }

diff --git a/include/runtime/instance/memory.h b/include/runtime/instance/memory.h
@@ -338,41 +338,15 @@ class MemoryInstance {
   /// @}
 };
 
-class SharedMemoryInstance : private MemoryInstance {
-  SharedMemoryInstance() = delete;
-  SharedMemoryInstance(MemoryInstance &&Inst) noexcept
-      : MemoryInstance(std::move(Inst)) {}
-  SharedMemoryInstance(const AST::MemoryType &MType,
-                       uint32_t PageLim = UINT32_C(65536)) noexcept
-      : MemoryInstance(MType, PageLim) {}
-
-  bool isShared() const noexcept {
-    std::shared_lock Lock(Mutex);
-    return MemoryInstance::isShared();
-  }
-
-  uint32_t getPageSize() const noexcept {
-    std::shared_lock Lock(Mutex);
-    return MemoryInstance::getPageSize();
-  }
-
-  const AST::MemoryType &getMemoryType() const noexcept {
-    std::shared_lock Lock(Mutex);
-    return MemoryInstance::getMemoryType();
-  }
-
-  bool checkAccessBound(uint32_t Offset, uint32_t Length) const noexcept {
-    std::shared_lock Lock(Mutex);
-    return MemoryInstance::checkAccessBound(Offset, Length);
-  }
-
-  uint32_t getBoundIdx() const noexcept {
-    std::shared_lock Lock(Mutex);
-    return MemoryInstance::getBoundIdx();
+class SharedMemory: public MemoryInstance {
+public:
+  SharedMemory() = delete;
+  SharedMemory(const AST::MemoryType &MType,
+                 uint32_t PageLim = UINT32_C(65536)) noexcept
+      : MemoryInstance(MType, PageLim) {
   }
-
 private:
-  mutable std::shared_mutex Mutex;
+  mutable std::shared_mutex Lock;
 };
 
 } // namespace Instance

diff --git a/include/runtime/instance/module.h b/include/runtime/instance/module.h
@@ -87,30 +87,15 @@ class ModuleInstance {
 
   void *getHostData() const noexcept { return HostData; }
 
-  ModuleInstance *Clone() const {
+  std::unique_ptr<ModuleInstance> CloneMemoryInEnv() const {
     std::unique_lock Lock(Mutex);
-    return UnsafeClone();
+    return UnsafeCloneMemoryInEnv();
   }
 
-  ModuleInstance *UnsafeClone() const {
-    ModuleInstance *Res = new ModuleInstance(ModName);
-
-    Res->FuncInsts = FuncInsts;
-    Res->TabInsts = TabInsts;
+  std::unique_ptr<ModuleInstance> UnsafeCloneMemoryInEnv() const {
+    std::unique_ptr<ModuleInstance> Res(new ModuleInstance("env"));
     Res->MemInsts = MemInsts;
-    Res->GlobInsts = GlobInsts;
-    Res->ElemInsts = ElemInsts;
-    Res->DataInsts = DataInsts;
-
-    Res->ImpGlobalNum = ImpGlobalNum;
-
-    Res->ExpFuncs = ExpFuncs;
-    Res->ExpTables = ExpTables;
     Res->ExpMems = ExpMems;
-    Res->ExpGlobals = ExpGlobals;
-
-    Res->StartFunc = StartFunc;
-
     return Res;
   }
 

diff --git a/include/vm/vm.h b/include/vm/vm.h
@@ -266,6 +266,9 @@ class VM {
   void unsafeLoadPlugInHosts();
   void unsafeRegisterBuiltInHosts();
   void unsafeRegisterPlugInHosts();
+  std::function<std::vector<
+      std::unique_ptr<WasmEdge::Runtime::Instance::ModuleInstance>>()>
+  unsafeGetRegisterHostsCallback();
 
   /// Helper function for execution.
   Expect<std::vector<std::pair<ValVariant, ValType>>>

diff --git a/lib/executor/executor.cpp b/lib/executor/executor.cpp
@@ -13,6 +13,7 @@ namespace Executor {
 Expect<std::unique_ptr<Runtime::Instance::ModuleInstance>>
 Executor::instantiateModule(Runtime::StoreManager &StoreMgr,
                             const AST::Module &Mod) {
+  ModuleSourceStructContext.SourceMod = &Mod;
   if (auto Res = instantiate(StoreMgr, Mod)) {
     return Res;
   } else {
@@ -26,6 +27,24 @@ Executor::instantiateModule(Runtime::StoreManager &StoreMgr,
   }
 }
 
+void Executor::setHostCallback(
+    std::function<std::vector<
+        std::unique_ptr<WasmEdge::Runtime::Instance::ModuleInstance>>()>
+        Func) {
+  ModuleSourceStructContext.HostCallback = Func;
+}
+
+Expect<std::unique_ptr<Runtime::Instance::ModuleInstance>>
+Executor::reinstantiateModule(
+    Runtime::StoreManager &StoreMgr,
+    std::vector<std::unique_ptr<Runtime::Instance::ModuleInstance>>
+        &HostModsHolder) {
+  HostModsHolder = ModuleSourceStructContext.HostCallback();
+  for (auto &It : HostModsHolder)
+    registerModule(StoreMgr, *It.get());
+  return instantiateModule(StoreMgr, *ModuleSourceStructContext.SourceMod);
+}
+
 /// Register a named WASM module. See "include/executor/executor.h".
 Expect<std::unique_ptr<Runtime::Instance::ModuleInstance>>
 Executor::registerModule(Runtime::StoreManager &StoreMgr,

diff --git a/lib/executor/instantiate/import.cpp b/lib/executor/instantiate/import.cpp
@@ -117,7 +117,7 @@ Expect<void> Executor::instantiate(Runtime::StoreManager &StoreMgr,
             auto EnvMod = new Runtime::Instance::ModuleInstance("env");
             EnvMod->addHostMemory(
                 ExtName,
-                std::make_unique<WasmEdge::Runtime::Instance::MemoryInstance>(
+                std::make_unique<WasmEdge::Runtime::Instance::SharedMemory>(
                     MemType));
             StoreMgr.registerModule(EnvMod);
           }

diff --git a/lib/vm/vm.cpp b/lib/vm/vm.cpp
@@ -57,6 +57,8 @@ void VM::unsafeInitVM() {
   // Register all module instances.
   unsafeRegisterBuiltInHosts();
   unsafeRegisterPlugInHosts();
+
+  ExecutorEngine.setHostCallback(unsafeGetRegisterHostsCallback());
 }
 
 void VM::unsafeLoadBuiltInHosts() {
@@ -141,6 +143,42 @@ void VM::unsafeRegisterPlugInHosts() {
   }
 }
 
+std::function<
+    std::vector<std::unique_ptr<WasmEdge::Runtime::Instance::ModuleInstance>>()>
+VM::unsafeGetRegisterHostsCallback() {
+  auto Callback = [&]() {
+    // Note: if copy unsafeLoadPlugInHosts() and unsafeLoadBuiltInHosts() to
+    // here can make this to lock-free
+    static std::shared_mutex Mutex;
+    std::unique_lock Lock(Mutex);
+
+    std::vector<std::unique_ptr<Runtime::Instance::ModuleInstance>> Res;
+    /// Built-in module instances mapped to the configurations. For WASI.
+    std::unordered_map<HostRegistration,
+                       std::unique_ptr<Runtime::Instance::ModuleInstance>>
+        BuiltInModInstsBackUp;
+    /// Loaded module instances from plug-ins.
+    std::vector<std::unique_ptr<Runtime::Instance::ModuleInstance>>
+        PlugInModInstsBackUp;
+
+    BuiltInModInstsBackUp.swap(BuiltInModInsts);
+    PlugInModInstsBackUp.swap(PlugInModInsts);
+    unsafeLoadBuiltInHosts();
+    unsafeLoadPlugInHosts();
+
+    for (auto &It : BuiltInModInsts)
+      Res.emplace_back(std::move(It.second));
+    for (auto &It : PlugInModInsts)
+      Res.emplace_back(std::move(It));
+    spdlog::info("add {} mods", Res.size());
+    // restore
+    BuiltInModInstsBackUp.swap(BuiltInModInsts);
+    PlugInModInstsBackUp.swap(PlugInModInsts);
+    return Res;
+  };
+  return Callback;
+}
+
 Expect<void> VM::unsafeRegisterModule(std::string_view Name,
                                       const std::filesystem::path &Path) {
   if (Stage == VMStage::Instantiated) {

diff --git a/plugins/wasi_threads/threadenv.cpp b/plugins/wasi_threads/threadenv.cpp
@@ -33,36 +33,46 @@ Plugin::Plugin::PluginDescriptor Descriptor{
 
 Plugin::PluginRegister WasiThreadsEnvironment::Register(&Descriptor);
 
-Expect<wasmedge_tid_t>
-WasiThreadsEnvironment::wasiThreadSpawn(Executor::Executor *Exec,
-                                        Runtime::Instance::ModuleInstance *Mods,
-                                        uint32_t ThreadStartArg) {
+Expect<wasmedge_tid_t> WasiThreadsEnvironment::wasiThreadSpawn(
+    Executor::Executor *Exec,
+    std::unique_ptr<Runtime::Instance::ModuleInstance> EnvMod,
+    int32_t ThreadStartArg) {
 
   auto Tid = Manager.allocate();
-  auto Warper = [=]() -> void {
-    std::vector<ValVariant> Params{Tid, ThreadStartArg};
+  auto Warper =
+      [=](std::unique_ptr<Runtime::Instance::ModuleInstance> EnvMod) -> void {
+    Runtime::StoreManager SM;
+    Exec->registerModule(SM, *EnvMod);
 
-    auto WasiThreadStartFunc = Mods->findFuncExports(WASI_ENTRY_POINT);
+    std::vector<std::unique_ptr<Runtime::Instance::ModuleInstance>>
+        HostModsHolder;
+    auto Inst = Exec->reinstantiateModule(SM, HostModsHolder);
+    if (!Inst) {
+      spdlog::error("reinstantiateModule fail");
+      std::exit(1);
+    }
 
+    auto NewMods = Inst.value().get();
+    auto WasiThreadStartFunc = NewMods->findFuncExports(WASI_ENTRY_POINT);
     if (WasiThreadStartFunc == nullptr) {
       spdlog::error("can not find wasi_thread_start");
       spdlog::error(ErrCode::Value::HostFuncError);
-      // exit() will terminate whole program
-      std::exit(int(ErrCode::Value::HostFuncError));
+      std::exit(1);
     }
 
-    auto Res = Exec->invoke(*WasiThreadStartFunc, Params,
+    std::vector<ValVariant> Params{Tid, ThreadStartArg};
+    auto Res = Exec->invoke(WasiThreadStartFunc, Params,
                             {ValType::I32, ValType::I32});
 
-    if (Res) {
-      spdlog::info("thread end good");
-    } else {
+    if (!Res) {
+
       spdlog::error("thread {} exit with error code: {}", Tid, Res.error());
       std::exit(1);
+    } else {
+      spdlog::info("thread end good");
     }
   };
-  spdlog::warn("std::thread Thread(Warper)");
-  std::thread Thread(Warper);
+  std::thread Thread(Warper, std::move(EnvMod));
   Thread.detach();
 
   return {Tid};

diff --git a/plugins/wasi_threads/threadenv.h b/plugins/wasi_threads/threadenv.h
@@ -77,8 +77,8 @@ class WasiThreadsEnvironment {
 
   Expect<wasmedge_tid_t>
   wasiThreadSpawn(Executor::Executor *Exec,
-                  Runtime::Instance::ModuleInstance *Mods,
-                  uint32_t ThreadStartArg);
+                  std::unique_ptr<Runtime::Instance::ModuleInstance> EnvMod,
+                  int32_t ThreadStartArg);
 
   static Plugin::PluginRegister Register;
 

diff --git a/plugins/wasi_threads/threadfunc.cpp b/plugins/wasi_threads/threadfunc.cpp
@@ -24,18 +24,23 @@ static bool HasWasiEntryPoint(const Runtime::Instance::ModuleInstance *Mods) {
 }
 
 Expect<int32_t> WasiThreadSpawn::body(const Runtime::CallingFrame &Frame,
-                                      uint32_t ThreadStartArg) {
-  // TODO: do this check on initialize
-  spdlog::warn("WasiThreadSpawn body");
+                                      int32_t ThreadStartArg) {
   if (!HasWasiEntryPoint(Frame.getModule())) {
     return Unexpect(ErrCode::Value::HostFuncError);
   }
 
   // Executor is thread_local variable, do not copy it
   auto NewThreadExec = Frame.getExecutor();
-  auto NewThreadModule = Frame.getModule()->Clone();
+  auto NewThreadModule = Frame.getModule()->UnsafeCloneMemoryInEnv();
 
-  return Env.wasiThreadSpawn(NewThreadExec, NewThreadModule, ThreadStartArg);
+  if (!NewThreadModule ||
+      NewThreadModule->findMemoryExports("memory") == nullptr) {
+    spdlog::info("Can not find shared memory");
+    return Unexpect(ErrCode::Value::HostFuncError);
+  }
+
+  return Env.wasiThreadSpawn(NewThreadExec, std::move(NewThreadModule),
+                             ThreadStartArg);
 }
 
 } // namespace Host

diff --git a/plugins/wasi_threads/threadfunc.h b/plugins/wasi_threads/threadfunc.h
@@ -16,7 +16,7 @@ class WasiThreadSpawn : public WasiThreads<WasiThreadSpawn> {
   WasiThreadSpawn(WasiThreadsEnvironment &HostEnv) : WasiThreads(HostEnv) {}
 
   Expect<int32_t> body(const Runtime::CallingFrame &Frame,
-                       uint32_t ThreadStartArg);
+                       int32_t ThreadStartArg);
 };
 
 } // namespace Host