Plug some memory leaks (#45)

---------

Co-authored-by: Frank Emrich <[email protected]>

cranelift/filetests/src/test_wasm/env.rs

@@ -798,6 +798,16 @@ impl<'a> FuncEnvironment for FuncEnv<'a> {
             .typed_continuations_cont_ref_get_cont_obj(builder, contref)
     }
 
+    /// TODO
+    fn typed_continuations_drop_cont_obj(
+        &mut self,
+        builder: &mut cranelift_frontend::FunctionBuilder,
+        contobj: ir::Value,
+    ) {
+        self.inner
+            .typed_continuations_drop_cont_obj(builder, contobj)
+    }
+
     fn typed_continuations_load_tag_return_values(
         &mut self,
         builder: &mut cranelift_frontend::FunctionBuilder,

cranelift/wasm/src/code_translator.rs

@@ -2577,6 +2577,12 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
                     &returns,
                     original_contobj,
                 );
+
+                // The continuation has returned and all `ContinuationReferences`
+                // to it should have been be invalidated. We may safely deallocate
+                // it.
+                environ.typed_continuations_drop_cont_obj(builder, original_contobj);
+
                 state.pushn(&values);
             }
         }

cranelift/wasm/src/environ/dummy.rs

@@ -798,6 +798,15 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
     ) -> ir::Value {
         unimplemented!()
     }
+
+    /// TODO
+    fn typed_continuations_drop_cont_obj(
+        &mut self,
+        _builder: &mut FunctionBuilder,
+        _contobj: ir::Value,
+    ) {
+        unimplemented!()
+    }
 }
 
 impl TypeConvert for DummyEnvironment {

cranelift/wasm/src/environ/spec.rs

@@ -650,6 +650,13 @@ pub trait FuncEnvironment: TargetEnvironment {
         contobj: ir::Value,
     );
 
+    /// TODO
+    fn typed_continuations_drop_cont_obj(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        contobj: ir::Value,
+    );
+
     /// TODO
     fn tag_params(&self, tag_index: u32) -> &[wasmtime_types::WasmType];
 

crates/cranelift/src/func_environ.rs

@@ -2382,12 +2382,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
                 offset += self.offsets.ptr.maximum_value_size() as i32;
             }
 
-            generate_builtin_call_no_return_val!(
-                self,
-                builder,
-                dealllocate_payload_buffer,
-                [nargs]
-            );
+            generate_builtin_call_no_return_val!(self, builder, deallocate_payload_buffer, [nargs]);
         }
         values
     }
@@ -2523,7 +2518,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
             let nargs = builder.ins().iconst(I32, values.len() as i64);
 
             let (_vmctx, payload_addr) =
-                generate_builtin_call!(self, builder, alllocate_payload_buffer, [nargs]);
+                generate_builtin_call!(self, builder, allocate_payload_buffer, [nargs]);
 
             let mut offset = 0;
             for value in values {
@@ -2554,6 +2549,15 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         return contref;
     }
 
+    /// TODO
+    fn typed_continuations_drop_cont_obj(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        contobj: ir::Value,
+    ) {
+        generate_builtin_call_no_return_val!(self, builder, drop_cont_obj, [contobj]);
+    }
+
     fn typed_continuations_load_return_values(
         &mut self,
         builder: &mut FunctionBuilder,

crates/environ/src/builtin.rs

@@ -84,11 +84,11 @@ macro_rules! foreach_builtin_function {
             /// Returns a pointer to that buffer.
             /// Such a payload buffer is only used to store payloads provided
             /// at a suspend site and read in a corresponding handler.
-            alllocate_payload_buffer(vmctx: vmctx, element_count: i32) -> pointer;
+            allocate_payload_buffer(vmctx: vmctx, element_count: i32) -> pointer;
             /// Counterpart to `alllocate_payload_buffer`, deallocating the
             /// buffer. For debugging purposes, `expected_element_capacity`
             /// should be the same value passed when allocating.
-            dealllocate_payload_buffer(vmctx: vmctx, expected_element_capacity: i32);
+            deallocate_payload_buffer(vmctx: vmctx, expected_element_capacity: i32);
             /// Returns pointer to the payload buffer, whose function was described earlier.
             /// `expected_element_capacity` should be the same value passed when
             /// allocating.
@@ -106,6 +106,8 @@ macro_rules! foreach_builtin_function {
             cont_obj_get_tag_return_values_buffer(vmctx: vmctx, contobj: pointer, expected_count : i32) -> pointer;
             /// Deallocated the tag return value buffer within the continuation object.
             cont_obj_deallocate_tag_return_values_buffer(vmctx: vmctx, contobj: pointer);
+            /// TODO
+            drop_cont_obj(vmctx: vmctx, contobj: pointer);
         }
     };
 }

crates/runtime/src/continuation.rs

@@ -160,8 +160,9 @@ pub fn cont_obj_get_tag_return_values_buffer(
 pub fn cont_obj_deallocate_tag_return_values_buffer(obj: *mut ContinuationObject) {
     let obj = unsafe { obj.as_mut().unwrap() };
     assert!(obj.state == State::Invoked);
-    let existing = obj.tag_return_values.take().unwrap();
-    mem::drop((*existing).data);
+    let payloads: Box<Payloads> = obj.tag_return_values.take().unwrap();
+    let _: Vec<u128> =
+        unsafe { Vec::from_raw_parts((*payloads).data, (*payloads).length, (*payloads).capacity) };
     obj.tag_return_values = None;
 }
 
@@ -187,19 +188,24 @@ pub fn new_cont_ref(contobj: *mut ContinuationObject) -> *mut ContinuationRefere
 /// TODO
 #[inline(always)]
 pub fn drop_cont_obj(contobj: *mut ContinuationObject) {
-    mem::drop(unsafe { (*contobj).fiber });
+    let _: Box<ContinuationFiber> = unsafe { Box::from_raw((*contobj).fiber) };
+    //mem::drop(unsafe { (*contobj).fiber });
     unsafe {
-        mem::drop((*contobj).args.data);
+        let _: Vec<u128> = Vec::from_raw_parts(
+            (*contobj).args.data,
+            (*contobj).args.length,
+            (*contobj).args.capacity,
+        );
     };
     let tag_return_vals = &mut unsafe { contobj.as_mut().unwrap() }.tag_return_values;
     match tag_return_vals {
         None => (),
-        Some(b) => mem::drop((*b).data),
+        Some(_) => cont_obj_deallocate_tag_return_values_buffer(contobj),
     }
 }
 
 /// TODO
-pub fn alllocate_payload_buffer(instance: &mut Instance, element_count: usize) -> *mut u128 {
+pub fn allocate_payload_buffer(instance: &mut Instance, element_count: usize) -> *mut u128 {
     // In the current design, we allocate a `Vec<u128>` and store a pointer to
     // it in the `VMContext` payloads pointer slot. We then return the pointer
     // to the `Vec`'s data, not to the `Vec` itself.
@@ -226,7 +232,7 @@ pub fn alllocate_payload_buffer(instance: &mut Instance, element_count: usize) -
 }
 
 /// TODO
-pub fn dealllocate_payload_buffer(instance: &mut Instance, element_count: usize) {
+pub fn deallocate_payload_buffer(instance: &mut Instance, element_count: usize) {
     let payload_ptr =
         unsafe { instance.get_typed_continuations_payloads_ptr_mut() as *mut *mut Vec<u128> };
 

crates/runtime/src/libcalls.rs

@@ -718,14 +718,18 @@ fn cont_obj_deallocate_tag_return_values_buffer(_instance: &mut Instance, contob
     );
 }
 
-fn alllocate_payload_buffer(instance: &mut Instance, element_count: u32) -> *mut u8 {
-    crate::continuation::alllocate_payload_buffer(instance, element_count as usize) as *mut u8
+fn allocate_payload_buffer(instance: &mut Instance, element_count: u32) -> *mut u8 {
+    crate::continuation::allocate_payload_buffer(instance, element_count as usize) as *mut u8
 }
 
-fn dealllocate_payload_buffer(instance: &mut Instance, expected_element_capacity: u32) {
-    crate::continuation::dealllocate_payload_buffer(instance, expected_element_capacity as usize);
+fn deallocate_payload_buffer(instance: &mut Instance, expected_element_capacity: u32) {
+    crate::continuation::deallocate_payload_buffer(instance, expected_element_capacity as usize);
 }
 
 fn get_payload_buffer(instance: &mut Instance, expected_element_capacity: u32) -> *mut u8 {
     crate::continuation::get_payload_buffer(instance, expected_element_capacity as usize) as *mut u8
 }
+
+fn drop_cont_obj(_instance: &mut Instance, contobj: *mut u8) {
+    crate::continuation::drop_cont_obj(contobj as *mut crate::continuation::ContinuationObject)
+}