fix(module): preserve setup factory admission and release lifecycle captures

crates/tinybus-module/src/lib.rs

@@ -326,10 +326,25 @@ struct RuntimeState {
     inbound: StdMutex<Option<mpsc::Sender<IncomingFrame>>>,
     host: HostCalls,
     runtime: StdMutex<Option<tokio::runtime::Runtime>>,
-    reinitialize: Option<Reinitialize>,
+    reinitialize: Arc<StdMutex<Option<Reinitialize>>>,
 }
 
 type Reinitialize = Box<dyn Fn(*const u8, usize) -> i32 + Send + Sync>;
+async fn dispose_failed_setup_callback(
+    callback: Arc<StdMutex<Option<Reinitialize>>>,
+) -> Result<()> {
+    // Reinitialization holds exclusion while driving setup with block_on.
+    // Waiting for it must leave runtime workers free to advance that future
+    // and its timer. Observe completion before publishing the failure.
+    tokio::task::spawn_blocking(move || {
+        callback
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .take();
+    })
+    .await
+    .map_err(|_| Error::failed("module failed-setup callback disposal panicked"))
+}
 type ReinitializeFactory = Box<
     dyn FnOnce(tokio::runtime::Handle, Arc<StdMutex<Option<Connection>>>) -> Reinitialize + Send,
 >;
@@ -363,6 +378,15 @@ unsafe extern "C" fn shutdown(ctx: *mut c_void, deadline_ms: u64) -> i32 {
             return TB_BAD_ARGUMENT;
         }
         let state = unsafe { &*(ctx.cast::<RuntimeState>()) };
+        // An in-flight reinitializer finishes before disposal. Once the slot
+        // is empty, late calls see a tombstone. Release the lock before joining
+        // runtime tasks, whose failed-setup cleanup also uses this slot.
+        let mut callback = state
+            .reinitialize
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        callback.take();
+        drop(callback);
         state.inbound.lock().expect("module inbound lock").take();
         match state.runtime.lock().expect("module runtime lock").take() {
             Some(runtime) => {
@@ -383,7 +407,11 @@ unsafe extern "C" fn reinitialize(ctx: *mut c_void, ptr: *const u8, len: usize)
             return TB_BAD_ARGUMENT;
         }
         let state = unsafe { &*(ctx.cast::<RuntimeState>()) };
-        let Some(reinitialize) = &state.reinitialize else {
+        let callback = state
+            .reinitialize
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let Some(reinitialize) = callback.as_ref() else {
             return TB_CLOSED;
         };
         reinitialize(ptr, len)
@@ -496,6 +524,24 @@ where
     }
 }
 
+// Validate the frozen ABI prefix before configuration parsing or a factory can
+// allocate state. Never read or write beyond a caller's declared allocation.
+unsafe fn validate_start_arguments(
+    host: *const TbHostVtable,
+    out: *mut TbModuleVtable,
+    worker_threads: usize,
+) -> Option<(u32, TbHostVtable)> {
+    if host.is_null() || out.is_null() || worker_threads == 0 {
+        return None;
+    }
+    let out_capacity = unsafe { out.cast::<u32>().read() };
+    if out_capacity < TB_MODULE_VTABLE_BASE_SIZE {
+        return None;
+    }
+    let host = unsafe { TbHostVtable::read_compatible(host) }?;
+    Some((out_capacity, host))
+}
+
 unsafe fn start_module_runtime<F, Fut>(
     host: *const TbHostVtable,
     out: *mut TbModuleVtable,
@@ -510,19 +556,9 @@ where
     Fut: Future<Output = Result<()>> + Send + 'static,
 {
     match catch_unwind(AssertUnwindSafe(|| {
-        if host.is_null() || out.is_null() || worker_threads == 0 {
-            return TB_BAD_ARGUMENT;
-        }
-        // `out.size` is input as well as output: an older ABI-v1 host owns a
-        // shorter allocation. Read only the frozen first field until its
-        // capacity is known, then write no further than that boundary.
-        let out_capacity = unsafe { out.cast::<u32>().read() };
-        if out_capacity < TB_MODULE_VTABLE_BASE_SIZE {
-            return TB_BAD_ARGUMENT;
-        }
-        // `size` is the frozen prefix field; do not copy the full vtable until
-        // the host has proved that all v1 fields are present.
-        let Some(host) = (unsafe { TbHostVtable::read_compatible(host) }) else {
+        let Some((out_capacity, host)) =
+            (unsafe { validate_start_arguments(host, out, worker_threads) })
+        else {
             return TB_BAD_ARGUMENT;
         };
         let host = HostCalls(host);
@@ -587,9 +623,14 @@ where
         let connection_slot = Arc::new(StdMutex::new(None));
         let task_connection_slot = connection_slot.clone();
 
-        runtime.spawn(async move {
+        let reinitializer =
+            reinitialize_factory.map(|factory| factory(runtime.handle().clone(), connection_slot));
+        let supports_reinitialize = reinitializer.is_some();
+        let reinitializer = Arc::new(StdMutex::new(reinitializer));
+        let failed_setup_callback = reinitializer.clone();
+        let setup_task = runtime.spawn(async move {
             let connection = module_connection(transport);
-            let outcome = match connection.handshake().await.map(|()| connection) {
+            match connection.handshake().await.map(|()| connection) {
                 Ok(connection) => {
                     connection.__set_panic_handler(std::sync::Arc::new(move || {
                         let location = panic_location
@@ -618,17 +659,28 @@ where
                     }
                 }
                 Err(error) => Err(error),
-            };
-            if let Err(error) = outcome {
-                tracing::error!(error = %error, "tinybus module failed");
+            }
+        });
+        runtime.spawn(async move {
+            let outcome = setup_task.await;
+            if !matches!(outcome, Ok(Ok(()))) {
+                // A terminal setup failure may never receive ABI shutdown.
+                // Dispose its callback under the same lifecycle exclusion.
+                if dispose_failed_setup_callback(failed_setup_callback)
+                    .await
+                    .is_err()
+                {
+                    tracing::error!("tinybus module failed-setup callback disposal failed");
+                }
+                match outcome {
+                    Ok(Err(error)) => tracing::error!(error = %error, "tinybus module failed"),
+                    Err(_) => tracing::error!("tinybus module setup task failed"),
+                    Ok(Ok(())) => unreachable!(),
+                }
                 host.fault();
             }
         });
 
-        let reinitializer =
-            reinitialize_factory.map(|factory| factory(runtime.handle().clone(), connection_slot));
-        let supports_reinitialize = reinitializer.is_some();
-
         let state = Box::new(RuntimeState {
             host,
             inbound: StdMutex::new(Some(inbound_tx)),
@@ -775,6 +827,12 @@ where
 }
 
 /// Initialize one configured instance with fresh captured setup state.
+///
+/// # Safety
+/// Non-null pointers must address readable host ABI prefixes and writable
+/// output storage matching their declared sizes. Host callbacks and context
+/// must remain valid until runtime shutdown. Null pointers and invalid sizes
+/// or worker counts are rejected before configuration or factory evaluation.
 #[doc(hidden)]
 pub unsafe fn start_reconfigurable_module_with_factory<C, F, Fut>(
     host: *const TbHostVtable,
@@ -785,7 +843,7 @@ pub unsafe fn start_reconfigurable_module_with_factory<C, F, Fut>(
 ) -> i32
 where
     C: serde::de::DeserializeOwned + Send + 'static,
-    F: Fn(Connection, C) -> Fut + Clone + Send + Sync + 'static,
+    F: Fn(Connection, C) -> Fut + Send + Sync + 'static,
     Fut: Future<Output = Result<()>> + Send + 'static,
 {
     unsafe {
@@ -794,13 +852,19 @@ where
             out,
             worker_threads,
             detach_on_panic,
-            factory,
+            || shared_factory_setup(factory),
             false,
         )
     }
 }
 
 /// Initialize one linked configured instance with fresh captured setup state.
+///
+/// # Safety
+/// Non-null pointers must address readable host ABI prefixes and writable
+/// output storage matching their declared sizes. Host callbacks and context
+/// must remain valid until runtime shutdown. Null pointers and invalid sizes
+/// or worker counts are rejected before configuration or factory evaluation.
 #[doc(hidden)]
 pub unsafe fn start_linked_reconfigurable_module_with_factory<C, F, Fut>(
     host: *const TbHostVtable,
@@ -811,7 +875,7 @@ pub unsafe fn start_linked_reconfigurable_module_with_factory<C, F, Fut>(
 ) -> i32
 where
     C: serde::de::DeserializeOwned + Send + 'static,
-    F: Fn(Connection, C) -> Fut + Clone + Send + Sync + 'static,
+    F: Fn(Connection, C) -> Fut + Send + Sync + 'static,
     Fut: Future<Output = Result<()>> + Send + 'static,
 {
     unsafe {
@@ -820,12 +884,24 @@ where
             out,
             worker_threads,
             detach_on_panic,
-            factory,
+            || shared_factory_setup(factory),
             true,
         )
     }
 }
 
+// The factory callback itself is shared, rather than cloned: a user's Clone
+// may deep-copy interior state and detach initial setup from reinitialization.
+fn shared_factory_setup<C, F, Fut>(
+    factory: impl FnOnce() -> F,
+) -> impl Fn(Connection, C) -> Fut + Clone + Send + Sync + 'static
+where
+    F: Fn(Connection, C) -> Fut + Send + Sync + 'static,
+{
+    let setup = Arc::new(factory());
+    move |connection, config| setup(connection, config)
+}
+
 unsafe fn start_reconfigurable_module_factory_mode<C, F, Fut>(
     host: *const TbHostVtable,
     out: *mut TbModuleVtable,
@@ -839,6 +915,9 @@ where
     F: Fn(Connection, C) -> Fut + Clone + Send + Sync + 'static,
     Fut: Future<Output = Result<()>> + Send + 'static,
 {
+    if unsafe { validate_start_arguments(host, out, worker_threads) }.is_none() {
+        return TB_BAD_ARGUMENT;
+    }
     let parsed = catch_unwind(AssertUnwindSafe(|| unsafe { parse_config::<C>(host) }));
     let config = match parsed {
         Ok(Ok(config)) => config,
@@ -904,7 +983,9 @@ where
 /// Configured modules may use `setup_factory = factory` in place of `setup`.
 /// After valid configuration is parsed, the safe no-argument factory runs once
 /// per ABI initialization and returns a
-/// cloneable `Fn(Connection, Config) -> Future<Output = Result<()>>`. Its captured
+/// `Fn(Connection, Config) -> Future<Output = Result<()>>`. The callback is shared
+/// by reference for initial setup and reinitialization, regardless of its Clone
+/// behavior. Its captured
 /// state is reused for that instance's reinitializations, allowing admission to
 /// remain bound without process-global state. Separate instances receive fresh
 /// state. A factory panic returns `TB_PANICKED` without unwinding into the host.
@@ -1142,5 +1223,9 @@ mod tests;
 mod static_link_tests;
 
 #[cfg(test)]
-#[path = "setup_factory_tests.rs"]
+#[path = "lib_setup_factory_tests.rs"]
 mod setup_factory_test;
+
+#[cfg(test)]
+#[path = "lib_lifecycle_tests.rs"]
+mod lifecycle_test;