Compile-time option for bypassing ContinuationReference creation (#42)

Cargo.toml

@@ -268,6 +268,15 @@ component-model = [
 ]
 winch = ["wasmtime/winch"]
 
+# Until we implement proper reference counting for `ContinuationReference` objects,
+# we may use this flag to bypass the creation of ContinuationReference objects,
+# directly using the corresponding ContinuationObject instead.
+# This is to allow running benchmarks that may create a lot of such objects and
+# may otherwise run out of memory.
+# Note that enabling this is highly unsafe, as it makes it impossible to detect
+# at runtime when an already taken continuation is used again.
+use_contobj_as_contref = []
+
 [[test]]
 name = "host_segfault"
 harness = false

cranelift/wasm/src/code_translator.rs

@@ -2417,13 +2417,10 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             let arg_types = environ.continuation_arguments(*type_index).to_vec();
             let result_types = environ.continuation_returns(*type_index).to_vec();
             let r = state.pop1();
-            state.push1(environ.translate_cont_new(
-                builder,
-                state,
-                r,
-                &arg_types,
-                &result_types,
-            )?);
+            let contobj =
+                environ.translate_cont_new(builder, state, r, &arg_types, &result_types)?;
+            let contref = environ.typed_continuations_new_cont_ref(builder, contobj);
+            state.push1(contref);
         }
         Operator::Resume {
             type_index,

crates/cranelift/src/func_environ.rs

@@ -2265,10 +2265,10 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         let nargs = builder.ins().iconst(I64, arg_types.len() as i64);
         let nreturns = builder.ins().iconst(I64, return_types.len() as i64);
 
-        let (_vmctx, contref) =
+        let (_vmctx, contobj) =
             generate_builtin_call!(self, builder, cont_new, [func, nargs, nreturns]);
 
-        Ok(contref)
+        Ok(contobj)
     }
 
     // TODO(dhil): Currently, this function invokes
@@ -2434,9 +2434,14 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         builder: &mut FunctionBuilder,
         contref: ir::Value,
     ) -> ir::Value {
-        let (_vmctx, contobj) =
-            generate_builtin_call!(self, builder, cont_ref_get_cont_obj, [contref]);
-        return contobj;
+        if cfg!(feature = "use_contobj_as_contref") {
+            // The "contref" is a contobj already
+            return contref;
+        } else {
+            let (_vmctx, contobj) =
+                generate_builtin_call!(self, builder, cont_ref_get_cont_obj, [contref]);
+            return contobj;
+        }
     }
 
     /// TODO
@@ -2545,8 +2550,13 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         builder: &mut FunctionBuilder,
         contobj_addr: ir::Value,
     ) -> ir::Value {
-        let (_vmctx, contref) = generate_builtin_call!(self, builder, new_cont_ref, [contobj_addr]);
-        return contref;
+        if cfg!(feature = "use_contobj_as_contref") {
+            return contobj_addr;
+        } else {
+            let (_vmctx, contref) =
+                generate_builtin_call!(self, builder, new_cont_ref, [contobj_addr]);
+            return contref;
+        }
     }
 
     /// TODO

crates/runtime/src/continuation.rs

@@ -92,6 +92,10 @@ pub fn cont_ref_get_cont_obj(
     contref: *mut ContinuationReference,
 ) -> Result<*mut ContinuationObject, TrapReason> {
     //FIXME rename to indicate that this invalidates the cont ref
+
+    // If this is enabled, we should never call this function.
+    assert!(!cfg!(feature = "use_contobj_as_contref"));
+
     let contopt = unsafe { contref.as_mut().unwrap().0 };
     match contopt {
         None => Err(TrapReason::user_with_backtrace(anyhow::Error::msg(
@@ -181,6 +185,9 @@ pub fn cont_obj_has_state_invoked(obj: *mut ContinuationObject) -> bool {
 /// TODO
 #[inline(always)]
 pub fn new_cont_ref(contobj: *mut ContinuationObject) -> *mut ContinuationReference {
+    // If this is enabled, we should never call this function.
+    assert!(!cfg!(feature = "use_contobj_as_contref"));
+
     let contref = Box::new(ContinuationReference(Some(contobj)));
     Box::into_raw(contref)
 }
@@ -267,7 +274,7 @@ pub fn cont_new(
     func: *mut u8,
     param_count: usize,
     result_count: usize,
-) -> *mut ContinuationReference {
+) -> *mut ContinuationObject {
     let func = func as *mut VMFuncRef;
     let callee_ctx = unsafe { (*func).vmctx };
     let caller_ctx = VMOpaqueContext::from_vmcontext(instance.vmctx());
@@ -298,9 +305,9 @@ pub fn cont_new(
         tag_return_values: None,
         state: State::Allocated,
     });
-    let contref = new_cont_ref(Box::into_raw(contobj));
-    contref // TODO(dhil): we need memory clean up of
-            // continuation reference objects.
+    // TODO(dhil): we need memory clean up of
+    // continuation reference objects.
+    return Box::into_raw(contobj);
 }
 
 /// TODO