call InstanceFlags::set_may_enter where appropriate

There's still more work to do to fully implement (and test) the reentrance rules
for concurrent tasks, but this is a start.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -3,7 +3,8 @@ use {
         component::func::{self, Func, Lower as _, LowerContext, Options},
         vm::{
             component::{
-                CallContext, ComponentInstance, ResourceTables, VMComponentContext, WaitableState,
+                CallContext, ComponentInstance, InstanceFlags, ResourceTables, VMComponentContext,
+                WaitableState,
             },
             mpk::{self, ProtectionMask},
             AsyncWasmCallState, PreviousAsyncWasmCallState, SendSyncPtr, VMFuncRef,
@@ -771,14 +772,25 @@ fn maybe_send_event<'a, T>(
 
         maybe_push_call_context(&mut store, guest_task)?;
 
+        let mut flags = unsafe {
+            (*store
+                .concurrent_state()
+                .component_instance
+                .unwrap()
+                .as_ptr())
+            .instance_flags(callback.instance)
+        };
+
         let params = &mut [
             ValRaw::u32(callback.context),
             ValRaw::u32(event as u32),
             ValRaw::u32(handle),
             ValRaw::u32(result),
         ];
         unsafe {
+            flags.set_may_enter(false);
             crate::Func::call_unchecked_raw(&mut store, callback.function.as_non_null(), params)?;
+            flags.set_may_enter(true);
         }
 
         maybe_pop_call_context(&mut store, guest_task)?;
@@ -1733,6 +1745,7 @@ fn make_call<T>(
     callee: SendSyncPtr<VMFuncRef>,
     param_count: usize,
     result_count: usize,
+    flags: Option<InstanceFlags>,
 ) -> impl FnOnce(
     StoreContextMut<T>,
 ) -> Result<([MaybeUninit<ValRaw>; MAX_FLAT_PARAMS], StoreContextMut<T>)>
@@ -1753,11 +1766,17 @@ fn make_call<T>(
         let mut cx = unsafe { StoreContextMut::<T>(&mut *cx.cast()) };
 
         unsafe {
+            if let Some(mut flags) = flags {
+                flags.set_may_enter(false);
+            }
             crate::Func::call_unchecked_raw(
                 &mut cx,
                 callee.as_non_null(),
                 &mut storage[..param_count.max(result_count)] as *mut [MaybeUninit<ValRaw>] as _,
             )?;
+            if let Some(mut flags) = flags {
+                flags.set_may_enter(true);
+            }
         }
 
         Ok((storage, cx))
@@ -1863,9 +1882,7 @@ fn do_start_call<'a, T>(
                 })?;
                 let mut cx = unsafe { StoreContextMut::<T>(&mut *cx.cast()) };
 
-                unsafe {
-                    flags.set_needs_post_return(false);
-                }
+                unsafe { flags.set_needs_post_return(false) }
 
                 if let Some(func) = post_return {
                     let arg = match result_count {
@@ -1882,9 +1899,7 @@ fn do_start_call<'a, T>(
                     }
                 }
 
-                unsafe {
-                    flags.set_may_enter(true);
-                }
+                unsafe { flags.set_may_enter(true) }
 
                 let (calls, host_table, _) = cx.0.component_resource_state();
                 ResourceTables {
@@ -2003,7 +2018,17 @@ pub(crate) extern "C" fn async_exit<T>(
             let result_count = usize::try_from(result_count).unwrap();
             assert!(result_count <= MAX_FLAT_RESULTS);
 
-            let call = make_call(guest_task, callee, param_count, result_count);
+            let call = make_call(
+                guest_task,
+                callee,
+                param_count,
+                result_count,
+                if callback.is_null() {
+                    None
+                } else {
+                    Some((*instance).instance_flags(callee_instance))
+                },
+            );
 
             let (guest_context, new_cx) = do_start_call(
                 cx,
@@ -2093,17 +2118,22 @@ pub(crate) fn start_call<'a, T: Send, LowerParams: Copy, R: 'static>(
 
     log::trace!("starting call {}", guest_task.rep());
 
+    let instance = store.0[instance.0].as_ref().unwrap().instance_ptr();
+
     let call = make_call(
         guest_task,
         SendSyncPtr::new(callee),
         mem::size_of::<LowerParams>() / mem::size_of::<ValRaw>(),
         1,
+        if callback.is_none() {
+            None
+        } else {
+            Some(unsafe { (*instance).instance_flags(component_instance) })
+        },
     );
 
     store.concurrent_state().guest_task = Some(guest_task);
 
-    let instance = store.0[instance.0].as_ref().unwrap().instance_ptr();
-
     store = do_start_call(
         store,
         instance,