allow to deregister modules

src/vm/errors.rs

@@ -45,6 +45,9 @@ pub enum FCEError {
 
     /// Returns where there is no module with such name.
     NoSuchModule,
+
+    /// Indicates that modules currently in use and couldn't be deleted.
+    ModuleInUse,
 }
 
 impl Error for FCEError {}
@@ -62,6 +65,9 @@ impl std::fmt::Display for FCEError {
             }
             FCEError::NonUniqueModuleName => write!(f, "FCE already has module with such name"),
             FCEError::NoSuchModule => write!(f, "FCE doesn't have a module with such name"),
+            FCEError::ModuleInUse => {
+                write!(f, "Module is used by other modules and couldn't be deleted")
+            }
         }
     }
 }

src/vm/fce.rs

@@ -41,18 +41,18 @@ impl FCE {
     }
 
     /// Extracts ABI of a module into Namespace.
-    fn create_import_object(module: &FCEModule, config: &Config) -> Namespace {
+    fn create_namespace_from_module(module: &FCEModule, config: &Config) -> Namespace {
         let mut namespace = Namespace::new();
-        let module_abi = module.get_abi();
+        let module_abi = module.acquire_abi();
 
         // TODO: introduce a macro for such things
-        let allocate = module_abi.allocate.clone().unwrap();
+        let allocate = module_abi.allocate.unwrap();
         namespace.insert(
             config.allocate_fn_name.clone(),
             func!(move |size: i32| -> i32 { allocate.call(size).expect("allocate failed") }),
         );
 
-        let invoke = module_abi.invoke.clone().unwrap();
+        let invoke = module_abi.invoke.unwrap();
         namespace.insert(
             config.invoke_fn_name.clone(),
             func!(move |offset: i32, size: i32| -> i32 {
@@ -60,7 +60,7 @@ impl FCE {
             }),
         );
 
-        let deallocate = module_abi.deallocate.clone().unwrap();
+        let deallocate = module_abi.deallocate.unwrap();
         namespace.insert(
             config.deallocate_fn_name.clone(),
             func!(move |ptr: i32, size: i32| {
@@ -68,7 +68,7 @@ impl FCE {
             }),
         );
 
-        let store = module_abi.store.clone().unwrap();
+        let store = module_abi.store.unwrap();
         namespace.insert(
             config.store_fn_name.clone(),
             func!(move |offset: i32, value: i32| {
@@ -76,7 +76,7 @@ impl FCE {
             }),
         );
 
-        let load = module_abi.load.clone().unwrap();
+        let load = module_abi.load.unwrap();
         namespace.insert(
             config.load_fn_name.clone(),
             func!(move |offset: i32| -> i32 { load.call(offset).expect("load failed") }),
@@ -119,7 +119,7 @@ impl FCEService for FCE {
         )?;
 
         // registers ABI of newly registered module in abi_import_object
-        let namespace = FCE::create_import_object(&module, &config);
+        let namespace = FCE::create_namespace_from_module(&module, &config);
         let module_name: String = module_name.into();
         self.abi_import_object
             .register(module_name.clone(), namespace);
@@ -134,12 +134,17 @@ impl FCEService for FCE {
     }
 
     fn unregister_module(&mut self, module_name: &str) -> Result<(), FCEError> {
-        self.modules
-            .remove(module_name)
-            .ok_or_else(|| FCEError::NoSuchModule)?;
-        // unregister abi from a dispatcher
+        match self.modules.entry(module_name.to_string()) {
+            Entry::Vacant(_) => Err(FCEError::NoSuchModule),
 
-        Ok(())
+            Entry::Occupied(module) => {
+                if module.get().get_abi_ref_counter() != 0 {
+                    return Err(FCEError::ModuleInUse)
+                }
+                module.remove_entry();
+                Ok(())
+            }
+        }
     }
 
     fn compute_state_hash(

src/vm/module/fce_module.rs

@@ -21,6 +21,7 @@ use crate::vm::module::{ModuleABI, ModuleAPI};
 
 use sha2::digest::generic_array::GenericArray;
 use sha2::digest::FixedOutput;
+use std::cell::RefCell;
 use wasmer_runtime::{compile, func, imports, Ctx, Func, Instance};
 use wasmer_runtime_core::import::ImportObject;
 use wasmer_runtime_core::memory::ptr::{Array, WasmPtr};
@@ -45,6 +46,8 @@ pub(crate) struct ABI {
 pub(crate) struct FCEModule {
     instance: Instance,
     abi: ABI,
+    // due to the Wasmer architecture :(
+    abi_ref_counter: RefCell<i32>,
 }
 
 impl FCEModule {
@@ -71,7 +74,11 @@ impl FCEModule {
         let instance = compile(&wasm_bytes)?.instantiate(&import_object)?;
         let abi = FCEModule::create_abi(&instance, &config_copy)?;
 
-        Ok(Self { instance, abi })
+        Ok(Self {
+            instance,
+            abi,
+            abi_ref_counter: RefCell::new(0),
+        })
     }
 
     #[rustfmt::skip]
@@ -111,8 +118,23 @@ impl FCEModule {
     /// Returns ABI of a module.
     /// (!) Be carefull and delete all instances of ABI before dropping corresponding module.
     /// There is no any internal ref counter due to the internals of Wasmer.
-    pub fn get_abi(&self) -> &ABI {
-        &self.abi
+    pub fn acquire_abi(&self) -> ABI {
+        use std::ops::AddAssign;
+        self.abi_ref_counter.borrow_mut().add_assign(1);
+        self.abi.clone()
+    }
+
+    pub fn release_abi(&self) {
+        if *self.abi_ref_counter.borrow() < 0 {
+            return;
+        }
+
+        use std::ops::SubAssign;
+        self.abi_ref_counter.borrow_mut().sub_assign(1);
+    }
+
+    pub fn get_abi_ref_counter(&self) -> i32 {
+        *self.abi_ref_counter.borrow()
     }
 
     /// Prints utf8 string of the given size from the given offset. Called from the wasm.
@@ -219,7 +241,3 @@ impl ModuleABI for FCEModule {
         Ok(self.abi.load.as_ref().unwrap().call(address)?)
     }
 }
-
-impl Drop for FCEModule {
-    fn drop(&mut self) {}
-}