Fix the current guest task when calling realloc (#12718)

* Fix the current guest task when calling `realloc`

This commit fixes some fallout from #12550 where after that PR the
current task registered when a guest's `realloc` was called was
incorrect. This additionally fixes issues where guest-to-guest
trampolines also had the wrong task registered when `realloc` was being
called. Finally this adjusts a few locations that "no borrows should be
active" traps happen to ensure that it happens around the time of when a
task returns rather than after the lowering of its values happens.

* Fix non-async build

Author: alexcrichton

crates/environ/src/fact/trampoline.rs

@@ -840,6 +840,25 @@ impl<'a, 'b> Compiler<'a, 'b> {
         }
         result_locals.reverse();
 
+        // Handle a few things related to the concurrent task infrastructure
+        // after the callee has finished, such as:
+        //
+        // * Validate that the callee dropped all its borrows
+        // * Transition the current running task back to the caller.
+        //
+        // This is not necessary if there are no resources in this call, nor if
+        // concurrency support is disabled (no tasks). Note that this must
+        // happen before lowering below because semantically that's where the
+        // "you forgot to drop borrows" trap shows up and additionally the
+        // lowering below may call realloc which is in the context of the
+        // caller's task, not the callee.
+        //
+        // FIXME: Apply the optimizations described in #12311.
+        if self.emit_resource_call || self.module.tunables.concurrency_support {
+            let exit_sync_call = self.module.import_exit_sync_call();
+            self.instruction(Call(exit_sync_call.as_u32()));
+        }
+
         // Like above during the translation of results the caller cannot be
         // left (as we might invoke things like `realloc`). Again the precise
         // order of everything doesn't matter since intermediate states cannot
@@ -865,14 +884,6 @@ impl<'a, 'b> Compiler<'a, 'b> {
             self.free_temp_local(tmp);
         }
 
-        if self.emit_resource_call || self.module.tunables.concurrency_support {
-            // Pop the task we pushed earlier off of the current task stack.
-            //
-            // FIXME: Apply the optimizations described in #12311.
-            let exit_sync_call = self.module.import_exit_sync_call();
-            self.instruction(Call(exit_sync_call.as_u32()));
-        }
-
         if self.module.tunables.concurrency_support {
             // Restore old `may_block_field`
             if let Some(old_task_may_block) = old_task_may_block {

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

@@ -1547,34 +1547,50 @@ impl StoreOpaque {
     /// relatively expensive table manipulations. This would ideally be
     /// optimized to avoid the full allocation of a `HostTask` in at least some
     /// situations.
-    pub fn enter_host_call(&mut self) -> Result<()> {
+    pub(crate) fn host_task_create(&mut self) -> Result<Option<TableId<HostTask>>> {
         if !self.concurrency_support() {
             self.enter_call_not_concurrent();
-            return Ok(());
+            return Ok(None);
         }
         let state = self.concurrent_state_mut();
         let caller = state.current_guest_thread()?;
         let task = state.push(HostTask::new(caller, HostTaskState::CalleeStarted))?;
         log::trace!("new host task {task:?}");
         self.set_thread(task)?;
+        Ok(Some(task))
+    }
+
+    /// Invoked before lowering the results of a host task to the guest.
+    ///
+    /// This is used to update the current thread annotations within the store
+    /// to ensure that it reflects the guest task, not the host task, since
+    /// lowering may execute guest code.
+    pub fn host_task_reenter_caller(&mut self) -> Result<()> {
+        if !self.concurrency_support() {
+            return Ok(());
+        }
+        let task = self.concurrent_state_mut().current_host_thread()?;
+        let caller = self.concurrent_state_mut().get_mut(task)?.caller;
+        self.set_thread(caller)?;
         Ok(())
     }
 
-    /// Dual of `enter_host_call` and signifies that the host has finished and
+    /// Dual of `host_task_create` and signifies that the host has finished and
     /// will be cleaned up.
     ///
     /// Note that this isn't invoked when the host is invoked asynchronously and
     /// the host isn't complete yet. In that situation the host task persists
-    /// and will be cleaned up separately.
-    pub fn exit_host_call(&mut self) -> Result<()> {
-        if !self.concurrency_support() {
-            self.exit_call_not_concurrent();
-            return Ok(());
+    /// and will be cleaned up separately in `subtask_drop`
+    pub(crate) fn host_task_delete(&mut self, task: Option<TableId<HostTask>>) -> Result<()> {
+        match task {
+            Some(task) => {
+                log::trace!("delete host task {task:?}");
+                self.concurrent_state_mut().delete(task)?;
+            }
+            None => {
+                self.exit_call_not_concurrent();
+            }
         }
-        let task = self.concurrent_state_mut().current_host_thread()?;
-        log::trace!("delete host task {task:?}");
-        let task = self.concurrent_state_mut().delete(task)?;
-        self.set_thread(task.caller)?;
         Ok(())
     }
 
@@ -4183,7 +4199,7 @@ type HostTaskFuture = Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>
 /// Represents the state of a pending host task.
 ///
 /// This is used to represent tasks when the guest calls into the host.
-struct HostTask {
+pub(crate) struct HostTask {
     common: WaitableCommon,
 
     /// Guest thread which called the host.

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

@@ -175,11 +175,15 @@ impl StoreOpaque {
         Ok(self.exit_call_not_concurrent())
     }
 
-    pub(crate) fn enter_host_call(&mut self) -> Result<()> {
+    pub(crate) fn host_task_create(&mut self) -> Result<()> {
         Ok(self.enter_call_not_concurrent())
     }
 
-    pub(crate) fn exit_host_call(&mut self) -> Result<()> {
+    pub(crate) fn host_task_reenter_caller(&mut self) -> Result<()> {
+        Ok(())
+    }
+
+    pub(crate) fn host_task_delete(&mut self, (): ()) -> Result<()> {
         Ok(self.exit_call_not_concurrent())
     }
 

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

@@ -649,6 +649,13 @@ impl Func {
             )?;
         }
 
+        // Validate that the task, after returning, has no more active borrows
+        // as they're required to have been dropped by this point.
+        store
+            .0
+            .component_resource_tables(Some(self.instance))
+            .validate_scope_exit()?;
+
         // SAFETY: We're relying on the correctness of the structure of
         // `LowerReturn` and the type-checking performed to acquire the
         // `TypedFunc` to make this safe. It should be the case that
@@ -706,11 +713,6 @@ impl Func {
 
         unsafe {
             call_post_return(&mut store, post_return, arg, flags)?;
-
-            store
-                .0
-                .component_resource_tables(Some(self.instance))
-                .validate_scope_exit()?;
             store.0.exit_guest_sync_call()?;
         }
         Ok(())

crates/wasmtime/src/runtime/component/func/host.rs

@@ -378,9 +378,9 @@ where
         }
 
         // Enter the host by pushing a `HostTask` into the concurrent state.
-        store.0.enter_host_call()?;
+        let host_task = store.0.host_task_create()?;
 
-        let task_exited = if async_ {
+        let host_task_complete = if async_ {
             #[cfg(feature = "component-model-async")]
             {
                 self.call_async_lower(store.as_context_mut(), instance, ty, options, storage)?
@@ -395,13 +395,13 @@ where
             true
         };
 
-        // If the host task exited, then it's popped and deallocated.
+        // If the host task completed, then it's deallocated.
         //
         // Note that if the host task did not exit then the `call_async_lower`
         // function transitively would have updated the current guest thread to
         // the caller of this host function.
-        if task_exited {
-            store.0.exit_host_call()?;
+        if host_task_complete {
+            store.0.host_task_delete(host_task)?;
         }
 
         Ok(())
@@ -571,6 +571,16 @@ where
         ret: Option<R>,
         dst: Destination<'_>,
     ) -> Result<()> {
+        // Before lowering below semantically ensure that the caller has dropped
+        // all of its borrows and such.
+        lower.validate_scope_exit()?;
+
+        // At this point we're transitioning back to the caller task which means
+        // that the current task needs to be updated. This will restore the
+        // currently running thread as the caller's thread, for example if
+        // lowering below calls `realloc` it'll use the right context.
+        lower.store.0.host_task_reenter_caller()?;
+
         if let Some(ret) = ret {
             let caller_instance = lower.options().instance;
             let mut flags = lower.instance_mut().instance_flags(caller_instance);
@@ -582,7 +592,6 @@ where
                 flags.set_may_leave(true);
             }
         }
-        lower.validate_scope_exit()?;
         Ok(())
     }
 }

crates/wast/src/spectest.rs

@@ -218,5 +218,6 @@ pub fn link_component_spectest<T>(linker: &mut component::Linker<T>) -> Result<(
             Box::pin(async move { std::future::pending::<Result<()>>().await })
         },
     )?;
+    i.func_wrap("return-hi", |_cx, (): ()| Ok(("hi".to_string(),)))?;
     Ok(())
 }

tests/all/component_model/resources.rs

@@ -814,7 +814,7 @@ fn thread_through_borrow() -> Result<()> {
 }
 
 #[test]
-fn cannot_use_borrow_for_own() -> Result<()> {
+fn borrows_must_be_dropped_before_lifting() -> Result<()> {
     let engine = super::engine();
     let c = Component::new(
         &engine,
@@ -824,6 +824,9 @@ fn cannot_use_borrow_for_own() -> Result<()> {
 
                 (core module $m
                     (func (export "f") (param i32) (result i32)
+                        ;; this is a `borrow $t`, and we're returning it as an
+                        ;; `own $t`. This should have an error about remaining
+                        ;; borrows rather than a mismatch of own/borrow.
                         local.get 0
                     )
                 )
@@ -848,7 +851,63 @@ fn cannot_use_borrow_for_own() -> Result<()> {
 
     let resource = Resource::new_own(100);
     let err = f.call(&mut store, (&resource,)).unwrap_err();
-    assert_eq!(err.to_string(), "cannot lift own resource from a borrow");
+    assert_eq!(
+        err.to_string(),
+        "borrow handles still remain at the end of the call"
+    );
+    Ok(())
+}
+
+#[test]
+fn cannot_use_borrow_for_own() -> Result<()> {
+    let engine = super::engine();
+    let c = Component::new(
+        &engine,
+        r#"
+            (component
+                (import "t" (type $t (sub resource)))
+                (import "f" (func $f (param "x" (own $t))))
+
+                (core module $m
+                    (import "" "f" (func $f (param i32)))
+                    (func (export "f") (param i32)
+                        local.get 0
+                        call $f
+                    )
+                )
+                (core func $f (canon lower (func $f)))
+                (core instance $i (instantiate $m
+                    (with "" (instance
+                        (export "f" (func $f))
+                    ))
+                ))
+
+                (func (export "f") (param "x" (borrow $t))
+                    (canon lift (core func $i "f")))
+            )
+        "#,
+    )?;
+
+    struct MyType;
+
+    let mut store = Store::new(&engine, ());
+    let mut linker = Linker::new(&engine);
+    linker
+        .root()
+        .resource("t", ResourceType::host::<MyType>(), |_, _| Ok(()))?;
+    linker
+        .root()
+        .func_wrap("f", |_, _: (Resource<MyType>,)| Ok(()))?;
+    let i = linker.instantiate(&mut store, &c)?;
+
+    let f = i.get_typed_func::<(&Resource<MyType>,), ()>(&mut store, "f")?;
+
+    let resource = Resource::new_own(100);
+    let err = f.call(&mut store, (&resource,)).unwrap_err();
+    assert!(
+        format!("{err:?}").contains("cannot lift own resource from a borrow"),
+        "bad error: {err:?}",
+    );
     Ok(())
 }