Merge pull request #2659 from lann/factors-nested-data

factors: Allow instance state to be nested in store data

crates/factor-key-value/src/lib.rs

@@ -54,10 +54,7 @@ impl Factor for KeyValueFactor {
     type AppState = AppState;
     type InstanceBuilder = InstanceBuilder;
 
-    fn init<Factors: RuntimeFactors>(
-        &mut self,
-        mut ctx: InitContext<Factors, Self>,
-    ) -> anyhow::Result<()> {
+    fn init<T: Send + 'static>(&mut self, mut ctx: InitContext<T, Self>) -> anyhow::Result<()> {
         ctx.link_bindings(spin_world::v1::key_value::add_to_linker)?;
         ctx.link_bindings(spin_world::v2::key_value::add_to_linker)?;
         Ok(())

crates/factor-llm/src/lib.rs

@@ -34,7 +34,7 @@ impl Factor for LlmFactor {
     type AppState = AppState;
     type InstanceBuilder = InstanceState;
 
-    fn init<T: RuntimeFactors>(
+    fn init<T: Send + 'static>(
         &mut self,
         mut ctx: spin_factors::InitContext<T, Self>,
     ) -> anyhow::Result<()> {

crates/factor-outbound-http/src/lib.rs

@@ -17,7 +17,7 @@ impl Factor for OutboundHttpFactor {
     type AppState = ();
     type InstanceBuilder = InstanceState;
 
-    fn init<T: RuntimeFactors>(
+    fn init<T: Send + 'static>(
         &mut self,
         mut ctx: spin_factors::InitContext<T, Self>,
     ) -> anyhow::Result<()> {

crates/factor-outbound-http/src/wasi.rs

@@ -1,14 +1,12 @@
 use http::Request;
-use spin_factors::{
-    wasmtime::component::ResourceTable, RuntimeFactors, RuntimeFactorsInstanceState,
-};
+use spin_factors::{wasmtime::component::ResourceTable, RuntimeFactorsInstanceState};
 use wasmtime_wasi_http::{
     bindings::http::types::ErrorCode, WasiHttpCtx, WasiHttpImpl, WasiHttpView,
 };
 
 use crate::{wasi_2023_10_18, wasi_2023_11_10, OutboundHttpFactor};
 
-pub(crate) fn add_to_linker<T: RuntimeFactors>(
+pub(crate) fn add_to_linker<T: Send + 'static>(
     ctx: &mut spin_factors::InitContext<T, OutboundHttpFactor>,
 ) -> anyhow::Result<()> {
     fn type_annotate<T, F>(f: F) -> F

crates/factor-sqlite/src/lib.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 use host::InstanceState;
 
 use async_trait::async_trait;
-use spin_factors::{anyhow, Factor, RuntimeFactors};
+use spin_factors::{anyhow, Factor};
 use spin_locked_app::MetadataKey;
 use spin_world::v1::sqlite as v1;
 use spin_world::v2::sqlite as v2;
@@ -32,7 +32,7 @@ impl Factor for SqliteFactor {
     type AppState = AppState;
     type InstanceBuilder = InstanceState;
 
-    fn init<T: RuntimeFactors>(
+    fn init<T: Send + 'static>(
         &mut self,
         mut ctx: spin_factors::InitContext<T, Self>,
     ) -> anyhow::Result<()> {

crates/factor-variables/src/lib.rs

@@ -40,10 +40,7 @@ impl Factor for VariablesFactor {
     type AppState = AppState;
     type InstanceBuilder = InstanceState;
 
-    fn init<Factors: RuntimeFactors>(
-        &mut self,
-        mut ctx: InitContext<Factors, Self>,
-    ) -> anyhow::Result<()> {
+    fn init<T: Send + 'static>(&mut self, mut ctx: InitContext<T, Self>) -> anyhow::Result<()> {
         ctx.link_bindings(spin_world::v1::config::add_to_linker)?;
         ctx.link_bindings(spin_world::v2::variables::add_to_linker)?;
         Ok(())

crates/factor-wasi/src/lib.rs

@@ -43,9 +43,9 @@ impl Factor for WasiFactor {
     type AppState = ();
     type InstanceBuilder = InstanceBuilder;
 
-    fn init<Factors: RuntimeFactors>(
+    fn init<T: Send + 'static>(
         &mut self,
-        mut ctx: InitContext<Factors, Self>,
+        mut ctx: InitContext<T, Self>,
     ) -> anyhow::Result<()> {
         fn type_annotate<T, F>(f: F) -> F
         where

crates/factors-derive/src/lib.rs

@@ -77,18 +77,20 @@ fn expand_factors(input: &DeriveInput) -> syn::Result<TokenStream> {
             type InstanceBuilders = #builders_name;
             type InstanceState = #state_name;
 
-            #[allow(clippy::needless_option_as_deref)]
-            fn init(
+            fn init<T: AsMut<Self::InstanceState> + Send + 'static>(
                 &mut self,
-                linker: &mut #wasmtime::component::Linker<#state_name>,
+                linker: &mut #wasmtime::component::Linker<T>,
             ) -> #Result<()> {
                 #(
-                    #Factor::init::<Self>(
+                    #Factor::init::<T>(
                         &mut self.#factor_names,
-                        #factors_path::InitContext::<Self, #factor_types>::new(
+                        #factors_path::InitContext::<T, #factor_types>::new(
                             linker,
-                            |state| &mut state.#factor_names,
-                            |state| (&mut state.#factor_names, &mut state.__table),
+                            |data| &mut data.as_mut().#factor_names,
+                            |data| {
+                                let state = data.as_mut();
+                                (&mut state.#factor_names, &mut state.__table)
+                            },
                         )
                     ).map_err(#Error::factor_init_error::<#factor_types>)?;
                 )*
@@ -226,5 +228,11 @@ fn expand_factors(input: &DeriveInput) -> syn::Result<TokenStream> {
                 &mut self.__table
             }
         }
+
+        impl AsMut<#state_name> for #state_name {
+            fn as_mut(&mut self) -> &mut Self {
+                self
+            }
+        }
     })
 }

crates/factors-test/src/lib.rs

@@ -2,8 +2,8 @@ use spin_app::locked::LockedApp;
 use spin_factors::{
     anyhow::{self, Context},
     serde::de::DeserializeOwned,
-    wasmtime::{Config, Engine},
-    App, Linker, RuntimeConfigSource, RuntimeFactors,
+    wasmtime::{component::Linker, Config, Engine},
+    App, RuntimeConfigSource, RuntimeFactors,
 };
 use spin_loader::FilesMountStrategy;
 
@@ -57,8 +57,8 @@ impl TestEnvironment {
         &self,
         mut factors: T,
     ) -> anyhow::Result<T::InstanceState> {
-        let mut linker = Self::new_linker::<T>();
-        factors.init(&mut linker)?;
+        let mut linker = Self::new_linker::<T::InstanceState>();
+        factors.init::<T::InstanceState>(&mut linker)?;
 
         let locked_app = self
             .build_locked_app()
@@ -77,7 +77,7 @@ impl TestEnvironment {
         Ok(factors.build_instance_state(builders)?)
     }
 
-    pub fn new_linker<T: RuntimeFactors>() -> Linker<T> {
+    pub fn new_linker<T>() -> Linker<T> {
         let engine = Engine::new(Config::new().async_support(true))
             .expect("wasmtime engine failed to initialize");
         Linker::<T>::new(&engine)

crates/factors/src/factor.rs

@@ -1,11 +1,10 @@
 use std::any::Any;
 
-use wasmtime::component::ResourceTable;
+use wasmtime::component::{Linker, ResourceTable};
 
 use crate::{
     prepare::FactorInstanceBuilder, runtime_config::RuntimeConfigTracker, App, Error,
-    FactorRuntimeConfig, InstanceBuilders, Linker, PrepareContext, RuntimeConfigSource,
-    RuntimeFactors,
+    FactorRuntimeConfig, InstanceBuilders, PrepareContext, RuntimeConfigSource, RuntimeFactors,
 };
 
 /// A contained (i.e., "factored") piece of runtime functionality.
@@ -26,10 +25,13 @@ pub trait Factor: Any + Sized {
 
     /// Initializes this `Factor` for a runtime once at runtime startup.
     ///
-    /// This will be called at most once, before any call to [`FactorInstanceBuilder::new`].
-    /// `InitContext` provides access to a wasmtime `Linker`, so this is where any bindgen
-    /// `add_to_linker` calls go.
-    fn init<T: RuntimeFactors>(&mut self, mut ctx: InitContext<T, Self>) -> anyhow::Result<()> {
+    /// This will be called at most once, before any call to
+    /// [`Factor::prepare`]. `InitContext` provides access to a wasmtime
+    /// `Linker`, so this is where any bindgen `add_to_linker` calls go.
+    ///
+    /// The type parameter `T` here is the same as the [`wasmtime::Store`] type
+    /// parameter `T`, which will contain the [`RuntimeFactors::InstanceState`].
+    fn init<T: Send + 'static>(&mut self, mut ctx: InitContext<T, Self>) -> anyhow::Result<()> {
         _ = &mut ctx;
         Ok(())
     }
@@ -72,22 +74,20 @@ pub trait Factor: Any + Sized {
 pub type FactorInstanceState<F> =
     <<F as Factor>::InstanceBuilder as FactorInstanceBuilder>::InstanceState;
 
-pub(crate) type GetDataFn<T, U> =
-    fn(&mut <T as RuntimeFactors>::InstanceState) -> &mut FactorInstanceState<U>;
+pub(crate) type GetDataFn<T, U> = fn(&mut T) -> &mut FactorInstanceState<U>;
 
-pub(crate) type GetDataWithTableFn<T, U> = fn(
-    &mut <T as RuntimeFactors>::InstanceState,
-) -> (&mut FactorInstanceState<U>, &mut ResourceTable);
+pub(crate) type GetDataWithTableFn<T, U> =
+    fn(&mut T) -> (&mut FactorInstanceState<U>, &mut ResourceTable);
 
 /// An InitContext is passed to [`Factor::init`], giving access to the global
 /// common [`wasmtime::component::Linker`].
-pub struct InitContext<'a, T: RuntimeFactors, U: Factor> {
+pub struct InitContext<'a, T, U: Factor> {
     pub(crate) linker: &'a mut Linker<T>,
     pub(crate) get_data: GetDataFn<T, U>,
     pub(crate) get_data_with_table: GetDataWithTableFn<T, U>,
 }
 
-impl<'a, T: RuntimeFactors, U: Factor> InitContext<'a, T, U> {
+impl<'a, T, U: Factor> InitContext<'a, T, U> {
     #[doc(hidden)]
     pub fn new(
         linker: &'a mut Linker<T>,
@@ -122,7 +122,7 @@ impl<'a, T: RuntimeFactors, U: Factor> InitContext<'a, T, U> {
         &mut self,
         add_to_linker: impl Fn(
             &mut Linker<T>,
-            fn(&mut T::InstanceState) -> &mut FactorInstanceState<U>,
+            fn(&mut T) -> &mut FactorInstanceState<U>,
         ) -> anyhow::Result<()>,
     ) -> anyhow::Result<()> {
         add_to_linker(self.linker, self.get_data)

crates/factors/src/lib.rs

@@ -16,9 +16,6 @@ pub use crate::{
     runtime_factors::{RuntimeFactors, RuntimeFactorsInstanceState},
 };
 
-/// A [`wasmtime::component::Linker`] used for a [`RuntimeFactors`] collection.
-pub type Linker<T> = wasmtime::component::Linker<<T as RuntimeFactors>::InstanceState>;
-
 // Temporary wrappers while refactoring
 pub type App = spin_app::App<'static, spin_app::InertLoader>;
 pub type AppComponent<'a> = spin_app::AppComponent<'a, spin_app::InertLoader>;

crates/factors/src/runtime_factors.rs

@@ -1,6 +1,6 @@
-use wasmtime::component::ResourceTable;
+use wasmtime::component::{Linker, ResourceTable};
 
-use crate::{factor::FactorInstanceState, App, ConfiguredApp, Factor, Linker, RuntimeConfigSource};
+use crate::{factor::FactorInstanceState, App, ConfiguredApp, Factor, RuntimeConfigSource};
 
 /// A collection of `Factor`s that are initialized and configured together.
 ///
@@ -10,7 +10,7 @@ use crate::{factor::FactorInstanceState, App, ConfiguredApp, Factor, Linker, Run
 ///
 /// A typical usage of `RuntimeFactors` would look something like the following pseudo-code:
 ///
-/// ```no_run
+/// ```ignore
 /// #[derive(RuntimeFactors)]
 /// struct MyFactors {
 ///  // ...
@@ -21,8 +21,10 @@ use crate::{factor::FactorInstanceState, App, ConfiguredApp, Factor, Linker, Run
 /// factors.init(&mut linker)?;
 /// // Configure the factors with an app and runtime config
 /// let configured_app = factors.configure_app(app, runtime_config)?;
+/// // Prepare instance state builders
+/// let builders = factors.prepare(&configured_app, "component-id")?;
 /// // Build the instance state for the factors
-/// let data factors.build_instance_state(&configured_app, component.id())
+/// let data = factors.build_instance_state(builders)?;
 /// // Initialize a `wasmtime` store with the instance state
 /// let mut store = wasmtime::Store::new(&engine, data);
 /// // Instantiate the component
@@ -36,12 +38,16 @@ pub trait RuntimeFactors: Sized + 'static {
     /// The collection of all the `InstanceBuilder`s of the factors.
     type InstanceBuilders;
 
-    /// Initialize the factors with a linker.
+    /// Initialize the factors with the given linker.
     ///
-    /// Each factor's `init` is called in turn.
-    fn init(&mut self, linker: &mut Linker<Self>) -> crate::Result<()>;
+    /// Each factor's `init` is called in turn. Must be called once before
+    /// [`RuntimeFactors::prepare`].
+    fn init<T: AsMut<Self::InstanceState> + Send + 'static>(
+        &mut self,
+        linker: &mut Linker<T>,
+    ) -> crate::Result<()>;
 
-    /// Configure the factors with an app and runtime config.
+    /// Configure the factors with the given app and runtime config.
     fn configure_app(
         &self,
         app: App,
@@ -76,7 +82,7 @@ pub trait RuntimeFactors: Sized + 'static {
 /// Get the state of a particular Factor from the overall InstanceState
 ///
 /// Implemented by `#[derive(RuntimeFactors)]`
-pub trait RuntimeFactorsInstanceState: Send + 'static {
+pub trait RuntimeFactorsInstanceState: AsMut<Self> + Send + 'static {
     fn get_with_table<F: Factor>(
         &mut self,
     ) -> Option<(&mut FactorInstanceState<F>, &mut ResourceTable)>;

crates/factors/tests/smoke.rs

@@ -23,6 +23,17 @@ struct Factors {
     key_value: KeyValueFactor,
 }
 
+struct Data {
+    factors_instance_state: FactorsInstanceState,
+    _other_data: usize,
+}
+
+impl AsMut<FactorsInstanceState> for Data {
+    fn as_mut(&mut self) -> &mut FactorsInstanceState {
+        &mut self.factors_instance_state
+    }
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn smoke_test_works() -> anyhow::Result<()> {
     let mut factors = Factors {
@@ -50,21 +61,27 @@ async fn smoke_test_works() -> anyhow::Result<()> {
     let engine = wasmtime::Engine::new(wasmtime::Config::new().async_support(true))?;
     let mut linker = wasmtime::component::Linker::new(&engine);
 
-    factors.init(&mut linker).unwrap();
+    factors.init::<Data>(&mut linker).unwrap();
 
     let configured_app = factors.configure_app(app, TestSource)?;
     let builders = factors.prepare(&configured_app, "smoke-app")?;
-    let data = factors.build_instance_state(builders)?;
+    let state = factors.build_instance_state(builders)?;
 
     assert_eq!(
-        data.variables
+        state
+            .variables
             .resolver()
             .resolve("smoke-app", "other".try_into().unwrap())
             .await
             .unwrap(),
         "<other value>"
     );
 
+    let data = Data {
+        factors_instance_state: state,
+        _other_data: 1,
+    };
+
     let mut store = wasmtime::Store::new(&engine, data);
 
     let component = configured_app.app().components().next().unwrap();
@@ -84,7 +101,7 @@ async fn smoke_test_works() -> anyhow::Result<()> {
 
     // Invoke handler
     let req = http::Request::get("/").body(Default::default()).unwrap();
-    let mut wasi_http = OutboundHttpFactor::get_wasi_http_impl(store.data_mut()).unwrap();
+    let mut wasi_http = OutboundHttpFactor::get_wasi_http_impl(store.data_mut().as_mut()).unwrap();
     let request = wasi_http.new_incoming_request(req)?;
     let (response_tx, response_rx) = tokio::sync::oneshot::channel();
     let response = wasi_http.new_response_outparam(response_tx)?;