@@ -280,7 +280,11 @@ WasmBase::WasmBase(std::unique_ptr<WasmVm> wasm_vm, std::string_view vm_id,
280280 }
281281}
282282
283- WasmBase::~WasmBase () {}
283+ WasmBase::~WasmBase () {
284+ root_contexts_.clear ();
285+ pending_done_.clear ();
286+ pending_delete_.clear ();
287+ }
284288
285289bool WasmBase::initialize (const std::string &code, bool allow_precompiled) {
286290 if (!wasm_vm_) {
@@ -319,22 +323,19 @@ bool WasmBase::initialize(const std::string &code, bool allow_precompiled) {
319323 return !isFailed ();
320324}
321325
322- ContextBase *WasmBase::getRootContext (std::string_view root_id) {
323- auto it = root_contexts_.find (std::string (root_id));
324- if (it == root_contexts_.end ()) {
325- return nullptr ;
326+ ContextBase *WasmBase::getRootContext (const std::shared_ptr<PluginBase> &plugin,
327+ bool allow_closed) {
328+ auto it = root_contexts_.find (plugin->key ());
329+ if (it != root_contexts_.end ()) {
330+ return it->second .get ();
326331 }
327- return it->second .get ();
328- }
329-
330- ContextBase *WasmBase::getOrCreateRootContext (const std::shared_ptr<PluginBase> &plugin) {
331- auto root_context = getRootContext (plugin->root_id_ );
332- if (!root_context) {
333- auto context = std::unique_ptr<ContextBase>(createRootContext (plugin));
334- root_context = context.get ();
335- root_contexts_[plugin->root_id_ ] = std::move (context);
332+ if (allow_closed) {
333+ it = pending_done_.find (plugin->key ());
334+ if (it != pending_done_.end ()) {
335+ return it->second .get ();
336+ }
336337 }
337- return root_context ;
338+ return nullptr ;
338339}
339340
340341void WasmBase::startVm (ContextBase *root_context) {
@@ -352,15 +353,14 @@ bool WasmBase::configure(ContextBase *root_context, std::shared_ptr<PluginBase>
352353}
353354
354355ContextBase *WasmBase::start (std::shared_ptr<PluginBase> plugin) {
355- auto root_id = plugin->root_id_ ;
356- auto it = root_contexts_.find (root_id);
356+ auto it = root_contexts_.find (plugin->key ());
357357 if (it != root_contexts_.end ()) {
358358 it->second ->onStart (plugin);
359359 return it->second .get ();
360360 }
361361 auto context = std::unique_ptr<ContextBase>(createRootContext (plugin));
362362 auto context_ptr = context.get ();
363- root_contexts_[root_id ] = std::move (context);
363+ root_contexts_[plugin-> key () ] = std::move (context);
364364 if (!context_ptr->onStart (plugin)) {
365365 return nullptr ;
366366 }
@@ -377,38 +377,49 @@ uint32_t WasmBase::allocContextId() {
377377 }
378378}
379379
380- void WasmBase::startShutdown () {
381- bool all_done = true ;
382- for (auto &p : root_contexts_) {
383- if (!p.second ->onDone ()) {
384- all_done = false ;
385- pending_done_.insert (p.second .get ());
380+ void WasmBase::startShutdown (const std::shared_ptr<PluginBase> &plugin) {
381+ auto it = root_contexts_.find (plugin->key ());
382+ if (it != root_contexts_.end ()) {
383+ if (it->second ->onDone ()) {
384+ it->second ->onDelete ();
385+ } else {
386+ pending_done_[it->first ] = std::move (it->second );
386387 }
388+ root_contexts_.erase (it);
387389 }
388- if (!all_done) {
389- shutdown_handle_ = std::make_unique<ShutdownHandle>(shared_from_this ());
390- } else {
391- finishShutdown ();
390+ }
391+
392+ void WasmBase::startShutdown () {
393+ auto it = root_contexts_.begin ();
394+ while (it != root_contexts_.end ()) {
395+ if (it->second ->onDone ()) {
396+ it->second ->onDelete ();
397+ } else {
398+ pending_done_[it->first ] = std::move (it->second );
399+ }
400+ it = root_contexts_.erase (it);
392401 }
393402}
394403
395404WasmResult WasmBase::done (ContextBase *root_context) {
396- auto it = pending_done_.find (root_context);
405+ auto it = pending_done_.find (root_context-> plugin_ -> key () );
397406 if (it == pending_done_.end ()) {
398407 return WasmResult::NotFound;
399408 }
409+ pending_delete_.insert (std::move (it->second ));
400410 pending_done_.erase (it);
401- if (pending_done_.empty () && shutdown_handle_) {
402- // Defer the delete so that onDelete is not called from within the done() handler.
403- addAfterVmCallAction (
404- [shutdown_handle = shutdown_handle_.release ()]() { delete shutdown_handle; });
405- }
411+ // Defer the delete so that onDelete is not called from within the done() handler.
412+ shutdown_handle_ = std::make_unique<ShutdownHandle>(shared_from_this ());
413+ addAfterVmCallAction (
414+ [shutdown_handle = shutdown_handle_.release ()]() { delete shutdown_handle; });
406415 return WasmResult::Ok;
407416}
408417
409418void WasmBase::finishShutdown () {
410- for (auto &p : root_contexts_) {
411- p.second ->onDelete ();
419+ auto it = pending_delete_.begin ();
420+ while (it != pending_delete_.end ()) {
421+ (*it)->onDelete ();
422+ it = pending_delete_.erase (it);
412423 }
413424}
414425
@@ -520,11 +531,18 @@ getOrCreateThreadLocalWasm(std::shared_ptr<WasmHandleBase> base_wasm,
520531 WasmHandleCloneFactory clone_factory) {
521532 auto wasm_handle = getThreadLocalWasm (base_wasm->wasm ()->vm_key ());
522533 if (wasm_handle) {
523- auto root_context = wasm_handle->wasm ()->getOrCreateRootContext (plugin);
524- if (!wasm_handle->wasm ()->configure (root_context, plugin)) {
525- base_wasm->wasm ()->fail (FailState::ConfigureFailed,
526- " Failed to configure thread-local Wasm code" );
527- return nullptr ;
534+ auto root_context = wasm_handle->wasm ()->getRootContext (plugin, false );
535+ if (!root_context) {
536+ root_context = wasm_handle->wasm ()->start (plugin);
537+ if (!root_context) {
538+ base_wasm->wasm ()->fail (FailState::StartFailed, " Failed to start thread-local Wasm" );
539+ return nullptr ;
540+ }
541+ if (!wasm_handle->wasm ()->configure (root_context, plugin)) {
542+ base_wasm->wasm ()->fail (FailState::ConfigureFailed,
543+ " Failed to configure thread-local Wasm plugin" );
544+ return nullptr ;
545+ }
528546 }
529547 return wasm_handle;
530548 }
0 commit comments