Remove an accessor of unforced_current_thread

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

@@ -2296,6 +2296,21 @@ impl StoreOpaque {
         }
         Ok(())
     }
+
+    /// Used by `ResourceTables` to record the scope of a borrow to get undone
+    /// in the future.
+    pub(crate) fn current_scope_id(&mut self) -> Result<Option<u32>> {
+        if !self.concurrency_support() {
+            return self.current_scope_id_not_concurrent();
+        }
+        let (bits, is_host) = match self.current_thread()? {
+            CurrentThread::Guest(id) => (id.task.rep(), false),
+            CurrentThread::Host(id) => (id.rep(), true),
+            CurrentThread::None => return Ok(None),
+        };
+        assert_eq!((bits << 1) >> 1, bits);
+        Ok(Some((bits << 1) | u32::from(is_host)))
+    }
 }
 
 enum CleanupTask {
@@ -5535,18 +5550,6 @@ impl ConcurrentState {
         }
     }
 
-    /// Used by `ResourceTables` to record the scope of a borrow to get undone
-    /// in the future.
-    pub fn current_call_context_scope_id(&self) -> Result<u32> {
-        let (bits, is_host) = match self.unforced_current_thread {
-            CurrentThread::Guest(id) => (id.task.rep(), false),
-            CurrentThread::Host(id) => (id.rep(), true),
-            CurrentThread::None => bail_bug!("current thread is not set"),
-        };
-        assert_eq!((bits << 1) >> 1, bits);
-        Ok((bits << 1) | u32::from(is_host))
-    }
-
     fn futures_mut(&mut self) -> Result<&mut FuturesUnordered<HostTaskFuture>> {
         match self.futures.get_mut().as_mut() {
             Some(f) => Ok(f),

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

@@ -908,7 +908,7 @@ impl<'a, T> Source<'a, T> {
                 bail_bug!("expected WriteState::GuestReady");
             };
 
-            let cx = &mut LiftContext::new(store.0.store_opaque_mut(), options, instance);
+            let cx = &mut LiftContext::new(store.0.store_opaque_mut(), options, instance)?;
             let ty = ty.payload(cx.types);
             let old_remaining = buffer.remaining_capacity();
             lift::<T, B>(
@@ -3327,7 +3327,7 @@ impl Instance {
 
                 let val = write_payload_ty
                     .map(|ty| {
-                        let lift = &mut LiftContext::new(store, write_options, write_instance);
+                        let lift = &mut LiftContext::new(store, write_options, write_instance)?;
                         let bytes = &lift.memory()[write_address..][..write_length_in_bytes];
                         Val::load(lift, *ty, bytes)
                     })
@@ -3406,7 +3406,7 @@ impl Instance {
                     }
                 } else {
                     let store_opaque = store.store_opaque_mut();
-                    let lift = &mut LiftContext::new(store_opaque, write_options, write_instance);
+                    let lift = &mut LiftContext::new(store_opaque, write_options, write_instance)?;
                     let bytes = &lift.memory()[write_address..][..write_length_in_bytes];
                     lift.consume_fuel_array(count, size_of::<Val>())?;
 
@@ -4251,7 +4251,7 @@ impl Instance {
         debug_msg_address: u32,
         debug_msg_len: u32,
     ) -> Result<u32> {
-        let lift_ctx = &mut LiftContext::new(store, options, self);
+        let lift_ctx = &mut LiftContext::new(store, options, self)?;
         let debug_msg = String::linear_lift_from_flat(
             lift_ctx,
             InterfaceType::String,

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

@@ -194,4 +194,8 @@ impl StoreOpaque {
     pub(crate) fn may_enter(&mut self, _instance: RuntimeInstance) -> Result<bool> {
         Ok(!self.trapped())
     }
+
+    pub(crate) fn current_scope_id(&mut self) -> Result<Option<u32>> {
+        self.current_scope_id_not_concurrent()
+    }
 }

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

@@ -696,7 +696,7 @@ impl Func {
         lift: impl FnOnce(&mut LiftContext, InterfaceType) -> Result<R>,
     ) -> Result<R> {
         let (options, _flags, ty, _) = self.abi_info(store);
-        let mut cx = LiftContext::new(store, options, self.instance);
+        let mut cx = LiftContext::new(store, options, self.instance)?;
         let ty = InterfaceType::Tuple(cx.types[ty].results);
         lift(&mut cx, ty)
     }

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

@@ -422,7 +422,7 @@ where
         options: OptionsIndex,
         storage: &mut [MaybeUninit<ValRaw>],
     ) -> Result<()> {
-        let mut lift = LiftContext::new(store.0.store_opaque_mut(), options, instance);
+        let mut lift = LiftContext::new(store.0.store_opaque_mut(), options, instance)?;
         let (params, rest) = self.load_params(&mut lift, ty, MAX_FLAT_PARAMS, storage)?;
 
         let ret = match self.run(store.as_context_mut(), params) {
@@ -473,7 +473,7 @@ where
 
         // Lift the parameters, either from flat storage or from linear
         // memory.
-        let mut lift = LiftContext::new(store.0.store_opaque_mut(), options, instance);
+        let mut lift = LiftContext::new(store.0.store_opaque_mut(), options, instance)?;
         let (params, rest) = self.load_params(&mut lift, ty, MAX_FLAT_ASYNC_PARAMS, storage)?;
 
         // Load/validate the return pointer, if present.

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

@@ -306,6 +306,7 @@ impl<'a, T: 'static> LowerContext<'a, T> {
 #[doc(hidden)]
 pub struct LiftContext<'a> {
     store_id: StoreId,
+    current_scope_id: Option<u32>,
     /// Like lowering, lifting always has options configured.
     options: OptionsIndex,
 
@@ -337,9 +338,10 @@ impl<'a> LiftContext<'a> {
         store: &'a mut StoreOpaque,
         options: OptionsIndex,
         instance_handle: Instance,
-    ) -> LiftContext<'a> {
+    ) -> Result<LiftContext<'a>> {
         let store_id = store.id();
         let hostcall_fuel = store.hostcall_fuel();
+        let current_scope_id = store.current_scope_id()?;
         // From `&mut StoreOpaque` provided the goal here is to project out
         // three different disjoint fields owned by the store: memory,
         // `CallContexts`, and `HandleTable`. There's no native API for that
@@ -352,8 +354,9 @@ impl<'a> LiftContext<'a> {
             store.lift_context_parts(instance_handle);
         let (component, instance) = instance.component_and_self();
 
-        LiftContext {
+        Ok(LiftContext {
             store_id,
+            current_scope_id,
             memory,
             options,
             types: component.types(),
@@ -363,7 +366,7 @@ impl<'a> LiftContext<'a> {
             host_table,
             host_resource_data,
             hostcall_fuel,
-        }
+        })
     }
 
     /// Returns the canonical options that are being used during lifting.
@@ -471,6 +474,7 @@ impl<'a> LiftContext<'a> {
                 host_table: self.host_table,
                 task_state: self.task_state,
                 guest: Some(self.instance.as_mut().instance_states()),
+                current_scope_id: self.current_scope_id,
             },
             self.host_resource_data,
         )

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

@@ -1809,7 +1809,10 @@ impl<T: Lift> WasmList<T> {
     // consumers should be validating through the iterator.
     pub fn get(&self, mut store: impl AsContextMut, index: usize) -> Option<Result<T>> {
         let store = store.as_context_mut().0;
-        let mut cx = LiftContext::new(store, self.options, self.instance);
+        let mut cx = match LiftContext::new(store, self.options, self.instance) {
+            Ok(cx) => cx,
+            Err(e) => return Some(Err(e)),
+        };
         self.get_from_store(&mut cx, index)
     }
 
@@ -1834,10 +1837,10 @@ impl<T: Lift> WasmList<T> {
     pub fn iter<'a, U: 'static>(
         &'a self,
         store: impl Into<StoreContextMut<'a, U>>,
-    ) -> impl ExactSizeIterator<Item = Result<T>> + 'a {
+    ) -> Result<impl ExactSizeIterator<Item = Result<T>> + 'a> {
         let store = store.into().0;
-        let mut cx = LiftContext::new(store, self.options, self.instance);
-        (0..self.len).map(move |i| self.get_from_store(&mut cx, i).unwrap())
+        let mut cx = LiftContext::new(store, self.options, self.instance)?;
+        Ok((0..self.len).map(move |i| self.get_from_store(&mut cx, i).unwrap()))
     }
 }
 

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

@@ -6,7 +6,7 @@ use crate::runtime::vm::component::{
     CallContext, ComponentInstance, HandleTable, OwnedComponentInstance,
 };
 use crate::store::{StoreData, StoreId, StoreOpaque};
-use crate::{AsContext, AsContextMut, Engine, Store, StoreContextMut};
+use crate::{AsContext, AsContextMut, Engine, Store, StoreContextMut, bail_bug};
 use core::pin::Pin;
 use wasmtime_environ::component::RuntimeComponentInstanceIndex;
 use wasmtime_environ::prelude::TryPrimaryMap;
@@ -403,10 +403,7 @@ impl StoreOpaque {
         vm::component::ResourceTables<'_>,
         &mut crate::component::HostResourceData,
     )> {
-        // NB: Force the current lazy deferred thread, if any, before handing
-        // out resource tables.
-        #[cfg(feature = "component-model-async")]
-        let _ = self.current_thread()?;
+        let current_scope_id = self.current_scope_id()?;
 
         let store_id = self.id();
         let data = self.component_data_mut();
@@ -425,6 +422,7 @@ impl StoreOpaque {
                 host_table: &mut data.component_host_table,
                 task_state: &mut data.task_state,
                 guest,
+                current_scope_id,
             },
             &mut data.host_resource_data,
         ))
@@ -466,6 +464,16 @@ impl StoreOpaque {
             None
         }
     }
+
+    pub(crate) fn current_scope_id_not_concurrent(&mut self) -> Result<Option<u32>> {
+        match &mut self.component_data_mut().task_state {
+            ComponentTaskState::NotConcurrent(state) => match state.scopes.len().checked_sub(1) {
+                Some(i) => Ok(Some(u32::try_from(i)?)),
+                None => Ok(None),
+            },
+            ComponentTaskState::Concurrent(_) => bail_bug!("should not be reachable"),
+        }
+    }
 }
 
 impl<T> Store<T> {
@@ -550,13 +558,6 @@ impl ComponentTaskState {
         }
     }
 
-    pub fn current_call_context_scope_id(&self) -> Result<u32> {
-        match self {
-            ComponentTaskState::NotConcurrent(state) => Ok(u32::try_from(state.scopes.len() - 1)?),
-            ComponentTaskState::Concurrent(state) => state.current_call_context_scope_id(),
-        }
-    }
-
     pub fn concurrent_state_mut(&mut self) -> &mut ConcurrentState {
         match self {
             ComponentTaskState::Concurrent(state) => state,

crates/wasmtime/src/runtime/vm/component/resources.rs

@@ -21,6 +21,7 @@
 //! namely in the `Resource<T>` and `ResourceAny` types.
 
 use super::{HandleTable, InstanceState, RemovedResource};
+use crate::bail_bug;
 use crate::component::store::ComponentTaskState;
 use crate::prelude::*;
 use core::error::Error;
@@ -71,6 +72,10 @@ pub struct ResourceTables<'a> {
     /// Task information about calls actively in use to track information such
     /// as borrow counts.
     pub task_state: &'a mut ComponentTaskState,
+
+    /// Identifier for the current "scope" which is used for various functions
+    /// on `task_state` above to mutate borrows/etc of the current scope.
+    pub current_scope_id: Option<u32>,
 }
 
 /// Typed representation of a "rep" for a resource.
@@ -261,6 +266,13 @@ impl ResourceTables<'_> {
         }
     }
 
+    fn current_scope_id(&self) -> Result<u32> {
+        match self.current_scope_id {
+            Some(id) => Ok(id),
+            None => bail_bug!("no current scope"),
+        }
+    }
+
     /// Extracts the underlying resource representation by lifting a "borrow"
     /// from the tables.
     ///
@@ -273,8 +285,8 @@ impl ResourceTables<'_> {
     pub fn resource_lift_borrow(&mut self, index: TypedResourceIndex) -> Result<u32> {
         let (rep, is_own) = self.table_for_index(&index).resource_lend(index)?;
         if is_own {
-            let scope = self.task_state.current_call_context_scope_id()?;
-            self.task_state.call_context(scope)?.lenders.push(index);
+            let current = self.current_scope_id()?;
+            self.task_state.call_context(current)?.lenders.push(index);
         }
         Ok(rep)
     }
@@ -292,7 +304,7 @@ impl ResourceTables<'_> {
     /// `VMComponentContext` which handles the special case of avoiding borrow
     /// tracking entirely.
     pub fn resource_lower_borrow(&mut self, resource: TypedResource) -> Result<u32> {
-        let scope = self.task_state.current_call_context_scope_id()?;
+        let scope = self.current_scope_id()?;
         let cx = self.task_state.call_context(scope)?;
         cx.borrow_count = cx.borrow_count.checked_add(1).unwrap();
         self.table_for_resource(&resource)
@@ -306,8 +318,7 @@ impl ResourceTables<'_> {
     /// resources that were originally passed in.
     #[inline]
     pub fn validate_scope_exit(&mut self) -> Result<()> {
-        let current = self.task_state.current_call_context_scope_id()?;
-        let cx = self.task_state.call_context(current)?;
+        let cx = self.task_state.call_context(self.current_scope_id()?)?;
         if cx.borrow_count > 0 {
             bail!("borrow handles still remain at the end of the call")
         }

tests/all/component_model/func.rs

@@ -773,12 +773,12 @@ fn strings() -> Result<()> {
 
         let ret = str_to_list8.call(&mut store, (x,))?.0;
         assert_eq!(
-            ret.iter(&mut store).collect::<Result<Vec<_>>>()?,
+            ret.iter(&mut store)?.collect::<Result<Vec<_>>>()?,
             x.as_bytes()
         );
 
         let ret = str_to_list16.call(&mut store, (x,))?.0;
-        assert_eq!(ret.iter(&mut store).collect::<Result<Vec<_>>>()?, utf16,);
+        assert_eq!(ret.iter(&mut store)?.collect::<Result<Vec<_>>>()?, utf16,);
 
         Ok(())
     };