From 2ceca7e296619a5e0c696a9fdc973a6bad126619 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 01:24:41 -0300 Subject: [PATCH 01/46] feat(client): add canonical low-level builder --- src/client.rs | 3 + src/client/builder.rs | 333 ++++++++++++++++++++++++++++++++++++++++ src/client/lifecycle.rs | 44 ++++-- src/lib.rs | 2 +- 4 files changed, 371 insertions(+), 11 deletions(-) create mode 100644 src/client/builder.rs diff --git a/src/client.rs b/src/client.rs index 3562d9ab5..fcd1401b5 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,6 +1,7 @@ mod accessors; mod adapters; mod app_state; +mod builder; mod context_impl; mod device_registry; pub(crate) mod device_topology; @@ -13,6 +14,8 @@ pub(crate) mod offline_resume; mod sender_keys; mod sessions; mod voip; +use builder::ClientAssembly; +pub use builder::{ClientBuild, ClientBuilder, ClientBuilderError}; pub use voip::{CallError, Voip}; use crate::cache::Cache; diff --git a/src/client/builder.rs b/src/client/builder.rs new file mode 100644 index 000000000..5749bf042 --- /dev/null +++ b/src/client/builder.rs @@ -0,0 +1,333 @@ +use std::sync::Arc; + +use thiserror::Error; + +use super::Client; +use crate::cache_config::CacheConfig; +use crate::http::HttpClient; +use crate::store::persistence_manager::PersistenceManager; +use crate::sync_task::MajorSyncTask; +use crate::transport::TransportFactory; +use wacore::runtime::Runtime; + +/// Result of constructing a [`Client`]. +/// +/// The sync-task receiver has a single consumer and is therefore transferred +/// together with the client rather than hidden behind a cloneable handle. +pub struct ClientBuild { + client: Arc, + sync_task_receiver: async_channel::Receiver, +} + +impl ClientBuild { + pub(crate) fn new( + client: Arc, + sync_task_receiver: async_channel::Receiver, + ) -> Self { + Self { + client, + sync_task_receiver, + } + } + + /// Return the constructed client. + pub fn client(&self) -> Arc { + Arc::clone(&self.client) + } + + /// Transfer ownership of the client and its sync-task receiver. + pub fn into_parts(self) -> (Arc, async_channel::Receiver) { + (self.client, self.sync_task_receiver) + } +} + +/// A validated client-construction failure. +#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum ClientBuilderError { + #[error("missing async runtime")] + MissingRuntime, + #[error("missing persistence manager")] + MissingPersistenceManager, + #[error("missing transport factory")] + MissingTransportFactory, + #[error("missing HTTP client")] + MissingHttpClient, +} + +/// Runtime-validated, low-level builder for [`Client`]. +/// +/// Unlike [`crate::BotBuilder`], this builder deliberately does not use +/// typestate. FFI and embedded hosts can populate dependencies dynamically and +/// receive a typed error without encoding Rust generic state in their wrapper. +pub struct ClientBuilder { + runtime: Option>, + persistence_manager: Option>, + transport_factory: Option>, + http_client: Option>, + override_version: Option<(u32, u32, u32)>, + cache_config: CacheConfig, +} + +impl Default for ClientBuilder { + fn default() -> Self { + Self::new() + } +} + +impl ClientBuilder { + /// Create an empty builder. All four platform dependencies are required. + pub fn new() -> Self { + Self { + runtime: None, + persistence_manager: None, + transport_factory: None, + http_client: None, + override_version: None, + cache_config: CacheConfig::default(), + } + } + + pub fn with_runtime(mut self, runtime: R) -> Self + where + R: Runtime, + { + self.runtime = Some(Arc::new(runtime)); + self + } + + pub fn with_runtime_arc(mut self, runtime: Arc) -> Self { + self.runtime = Some(runtime); + self + } + + pub fn with_persistence_manager( + mut self, + persistence_manager: Arc, + ) -> Self { + self.persistence_manager = Some(persistence_manager); + self + } + + pub fn with_transport_factory(mut self, transport_factory: T) -> Self + where + T: TransportFactory + 'static, + { + self.transport_factory = Some(Arc::new(transport_factory)); + self + } + + pub fn with_transport_factory_arc( + mut self, + transport_factory: Arc, + ) -> Self { + self.transport_factory = Some(transport_factory); + self + } + + pub fn with_http_client(mut self, http_client: H) -> Self + where + H: HttpClient + 'static, + { + self.http_client = Some(Arc::new(http_client)); + self + } + + pub fn with_http_client_arc(mut self, http_client: Arc) -> Self { + self.http_client = Some(http_client); + self + } + + pub fn with_version_override(mut self, version: (u32, u32, u32)) -> Self { + self.override_version = Some(version); + self + } + + pub fn with_cache_config(mut self, cache_config: CacheConfig) -> Self { + self.cache_config = cache_config; + self + } + + /// Validate dependencies, assemble an inert client, then start its services. + pub async fn build(self) -> Result { + self.build_boxed().await + } + + #[inline(never)] + fn build_boxed( + self, + ) -> wacore::runtime::BoxFuture<'static, Result> { + Box::pin(async move { + let runtime = self.runtime.ok_or(ClientBuilderError::MissingRuntime)?; + let persistence_manager = self + .persistence_manager + .ok_or(ClientBuilderError::MissingPersistenceManager)?; + let transport_factory = self + .transport_factory + .ok_or(ClientBuilderError::MissingTransportFactory)?; + let http_client = self + .http_client + .ok_or(ClientBuilderError::MissingHttpClient)?; + + Ok(Self::build_required( + runtime, + persistence_manager, + transport_factory, + http_client, + self.override_version, + self.cache_config, + )) + }) + } + + pub(crate) fn build_required( + runtime: Arc, + persistence_manager: Arc, + transport_factory: Arc, + http_client: Arc, + override_version: Option<(u32, u32, u32)>, + cache_config: CacheConfig, + ) -> ClientBuild { + Client::assemble( + runtime, + persistence_manager, + transport_factory, + http_client, + override_version, + cache_config, + ) + .start() + } +} + +/// Owns the only path from a fully allocated client to started background +/// services, preventing callers from publishing a partially configured client. +pub(super) struct ClientAssembly { + client: Arc, + sync_task_receiver: async_channel::Receiver, +} + +impl ClientAssembly { + pub(super) fn new( + client: Arc, + sync_task_receiver: async_channel::Receiver, + ) -> Self { + Self { + client, + sync_task_receiver, + } + } + + pub(super) fn start(self) -> ClientBuild { + self.client.start_services(); + ClientBuild::new(self.client, self.sync_task_receiver) + } +} + +#[cfg(test)] +mod tests { + use std::future::Future; + use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use super::*; + use crate::runtime_impl::TokioRuntime; + use crate::test_utils::MockHttpClient; + use crate::transport::mock::MockTransportFactory; + use wacore::runtime::AbortHandle; + + struct CountingRuntime { + spawns: Arc, + } + + #[async_trait::async_trait] + impl Runtime for CountingRuntime { + fn spawn(&self, future: Pin + Send + 'static>>) -> AbortHandle { + self.spawns.fetch_add(1, Ordering::SeqCst); + TokioRuntime.spawn(future) + } + + fn sleep(&self, duration: Duration) -> Pin + Send>> { + TokioRuntime.sleep(duration) + } + + fn spawn_blocking( + &self, + f: Box, + ) -> Pin + Send>> { + TokioRuntime.spawn_blocking(f) + } + + fn yield_now(&self) -> Option + Send>>> { + TokioRuntime.yield_now() + } + } + + #[tokio::test] + async fn validates_required_dependencies_before_assembly() { + assert!(matches!( + ClientBuilder::new().build().await, + Err(ClientBuilderError::MissingRuntime) + )); + + assert!(matches!( + ClientBuilder::new() + .with_runtime(TokioRuntime) + .build() + .await, + Err(ClientBuilderError::MissingPersistenceManager) + )); + + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + assert!(matches!( + ClientBuilder::new() + .with_runtime(TokioRuntime) + .with_persistence_manager(Arc::clone(&persistence_manager)) + .build() + .await, + Err(ClientBuilderError::MissingTransportFactory) + )); + + assert!(matches!( + ClientBuilder::new() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .build() + .await, + Err(ClientBuilderError::MissingHttpClient) + )); + } + + #[tokio::test] + async fn assembly_is_inert_until_started() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let spawns = Arc::new(AtomicUsize::new(0)); + let runtime = Arc::new(CountingRuntime { + spawns: Arc::clone(&spawns), + }) as Arc; + + let assembly = Client::assemble( + runtime, + persistence_manager, + Arc::new(MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + CacheConfig::default(), + ); + assert_eq!(spawns.load(Ordering::SeqCst), 0); + + let build = assembly.start(); + assert_eq!(spawns.load(Ordering::SeqCst), 2); + build.client().signal_shutdown_sync(); + } +} diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 4d7de6912..c3ab2621c 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -11,6 +11,11 @@ impl Client { /// long is the reconnect backoff counter reset to its base. pub(crate) const STABLE_CONNECTION_RESET_MS: i64 = 30_000; + /// Create a runtime-validated low-level client builder. + pub fn builder() -> ClientBuilder { + ClientBuilder::new() + } + pub fn shutdown_signal(&self) -> wacore::runtime::ShutdownSignal { self.shutdown_notifier.subscribe() } @@ -87,7 +92,7 @@ impl Client { http_client: Arc, override_version: Option<(u32, u32, u32)>, ) -> (Arc, async_channel::Receiver) { - Self::new_with_cache_config( + ClientBuilder::build_required( runtime, persistence_manager, transport_factory, @@ -95,7 +100,7 @@ impl Client { override_version, CacheConfig::default(), ) - .await + .into_parts() } /// Create a new `Client` with a custom [`CacheConfig`]. @@ -107,6 +112,25 @@ impl Client { override_version: Option<(u32, u32, u32)>, cache_config: CacheConfig, ) -> (Arc, async_channel::Receiver) { + ClientBuilder::build_required( + runtime, + persistence_manager, + transport_factory, + http_client, + override_version, + cache_config, + ) + .into_parts() + } + + pub(super) fn assemble( + runtime: Arc, + persistence_manager: Arc, + transport_factory: Arc, + http_client: Arc, + override_version: Option<(u32, u32, u32)>, + cache_config: CacheConfig, + ) -> ClientAssembly { let mut unique_id_bytes = [0u8; 2]; rand::make_rng::().fill_bytes(&mut unique_id_bytes); @@ -303,9 +327,12 @@ impl Client { .attach_topology(Arc::clone(&arc.device_topology)); let _ = arc.self_weak.set(Arc::downgrade(&arc)); - // Warm up the LID-PN cache from persistent storage - let warm_up_arc = arc.clone(); - arc.runtime + ClientAssembly::new(arc, rx) + } + + pub(super) fn start_services(self: &Arc) { + let warm_up_arc = self.clone(); + self.runtime .spawn(Box::pin(async move { if let Err(e) = warm_up_arc.warm_up_lid_pn_cache().await { warn!("Failed to warm up LID-PN cache: {e}"); @@ -313,15 +340,12 @@ impl Client { })) .detach(); - // Start background task to clean up stale device registry entries - let cleanup_arc = arc.clone(); - arc.runtime + let cleanup_arc = self.clone(); + self.runtime .spawn(Box::pin(async move { cleanup_arc.device_registry_cleanup_loop().await; })) .detach(); - - (arc, rx) } // Deliberately NOT instrumented: this span would live for the entire client diff --git a/src/lib.rs b/src/lib.rs index 95de556bc..029c24e20 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -84,7 +84,6 @@ pub mod types; pub mod client; pub(crate) mod flush_scope; -pub use client::Client; /// Shared base error for transport/connection concerns; the per-domain error /// types embed it. pub use client::ClientError; @@ -94,6 +93,7 @@ pub use client::{ StatsSnapshot, StorageResourceReport, TransportResourceReport, }; pub use client::{CallError, Voip}; +pub use client::{Client, ClientBuild, ClientBuilder, ClientBuilderError}; pub use types::durability_hook::InboundDurabilityHook; pub use types::retry_admission::RetryAdmission; pub mod download; From a569b50b23d8136570ad0dcc9f487dc4f7e87482 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 01:31:50 -0300 Subject: [PATCH 02/46] refactor(bot): build clients through canonical pipeline --- src/bot.rs | 125 ++++------------- src/client/builder.rs | 281 ++++++++++++++++++++++++++++++++++++--- src/types/enc_handler.rs | 4 +- 3 files changed, 293 insertions(+), 117 deletions(-) diff --git a/src/bot.rs b/src/bot.rs index 29709ceae..5c8e9ffdb 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -1,5 +1,5 @@ use crate::cache_config::CacheConfig; -use crate::client::Client; +use crate::client::{Client, ClientBuilderError}; use crate::pair_code::PairCodeOptions; use crate::store::commands::DeviceCommand; use crate::store::error::StoreError; @@ -87,49 +87,8 @@ pub enum BotBuilderError { /// Initializing the device row in the storage backend failed. #[error("failed to initialize the device store: {0}")] Store(#[from] StoreError), - /// An inbound durability hook was registered with a backend that does not - /// implement the pending-inbound buffer it requires. - #[error("the configured backend does not support the inbound durability hook: {0}")] - UnsupportedDurabilityBackend(String), -} - -/// Verify the backend round-trips a pending-inbound buffer entry before we accept -/// an inbound durability hook. A backend relying on the no-op/`Err` trait -/// defaults fails here instead of silently looping every inbound message unacked. -async fn probe_durability_backend( - backend: &std::sync::Arc, -) -> std::result::Result<(), BotBuilderError> { - // A real JID (a backend may validate the format) and an id unique per probe - // invocation (pid + atomic counter) so concurrent builders on the same store - // never race on a shared probe row and false-fail. - use portable_atomic::{AtomicU64, Ordering}; - static PROBE_SEQ: AtomicU64 = AtomicU64::new(0); - const PROBE_JID: &str = "0@s.whatsapp.net"; - const PROBE_PAYLOAD: &[u8] = b"probe"; - let probe_id = format!( - "__wa_durability_probe_{}_{}__", - std::process::id(), - PROBE_SEQ.fetch_add(1, Ordering::Relaxed) - ); - let map_err = |e: StoreError| BotBuilderError::UnsupportedDurabilityBackend(e.to_string()); - backend - .store_pending_inbound(PROBE_JID, PROBE_JID, &probe_id, PROBE_PAYLOAD) - .await - .map_err(map_err)?; - let got = backend - .get_pending_inbound(PROBE_JID, PROBE_JID, &probe_id) - .await - .map_err(map_err)?; - backend - .delete_pending_inbound(PROBE_JID, PROBE_JID, &probe_id) - .await - .map_err(map_err)?; - if got.as_deref() != Some(PROBE_PAYLOAD) { - return Err(BotBuilderError::UnsupportedDurabilityBackend( - "pending-inbound buffer did not round-trip".to_string(), - )); - } - Ok(()) + #[error(transparent)] + Client(#[from] ClientBuilderError), } /// `message` is `Arc` so cloning the context across spawned tasks only bumps a @@ -1278,18 +1237,8 @@ impl BotBuilder { unreachable!("typestate guarantees all required fields are Provided") }; - // Instrument the runtime before anything spawns through it, so every - // internal task (noise sender, saver, workers) reports to the hook. - // Default (None): the original runtime is used untouched. The Bot - // keeps its own copy for the `run()` path (see the field doc). let task_instrument = self.task_instrument; let alloc_meter = self.alloc_meter; - let runtime: Arc = match task_instrument.clone() { - Some(instrument) => { - Arc::new(wacore::stats::InstrumentedRuntime::new(runtime, instrument)) - } - None => runtime, - }; // Note: For multi-account mode, create the backend with SqliteStore::new_for_device() // before passing it to with_backend_arc() @@ -1331,57 +1280,37 @@ impl BotBuilder { } info!("Creating client..."); - let (client, sync_task_receiver) = Client::new_with_cache_config( - runtime.clone(), - persistence_manager.clone(), - transport_factory, - http_client, - self.override_version, - self.cache_config, - ) - .await; - - let saver_handle = persistence_manager.run_background_saver( - runtime, - std::time::Duration::from_secs(30), - client.shutdown_signal(), - ); - // Tie the saver task to Arc so extracting client() and outliving - // Bot keeps periodic persistence alive. Client::drop on the last Arc - // drops the AbortHandle and aborts the task. - let _ = client.saver_handle.set(saver_handle); - - // Typed alloc-meter handle for resource_report (its poll hooks are - // already wired via task_instrument above). - if let Some(meter) = alloc_meter { - let _ = client.alloc_meter.set(meter); + let mut client_builder = Client::builder() + .with_runtime_arc(runtime) + .with_persistence_manager(persistence_manager) + .with_transport_factory_arc(transport_factory) + .with_http_client_arc(http_client) + .with_cache_config(self.cache_config) + .with_custom_enc_handlers(self.custom_enc_handlers) + .with_skip_history_sync(self.skip_history_sync) + .with_background_saver_interval(std::time::Duration::from_secs(30)); + + if let Some(version) = self.override_version { + client_builder = client_builder.with_version_override(version); } - - // Register custom enc handlers. Immutable after build, so set the whole - // map once; the receive hot path then reads it lock-free. - let _ = client.custom_enc_handlers.set(self.custom_enc_handlers); - - // Inbound durability hook (opt-in). Immutable after build; the receive - // path reads it lock-free. Probe the backend first: a backend that does - // not implement the pending-inbound buffer would otherwise leave every - // inbound message unacked and looping forever at runtime, so reject it - // here with a clear error instead. if let Some(hook) = self.inbound_durability_hook { - probe_durability_backend(&client.persistence_manager.backend()).await?; - let _ = client.inbound_durability_hook.set(hook); - } - - if self.skip_history_sync { - client.set_skip_history_sync(true); + client_builder = client_builder.with_inbound_durability_hook_arc(hook); } - if let Some(count) = self.wanted_pre_key_count { - client.set_wanted_pre_key_count(count); + client_builder = client_builder.with_wanted_pre_key_count(count); } - if let Some((burst, refill_per_min)) = self.resend_rate_limit { - client.set_resend_rate_limit(burst, refill_per_min); + client_builder = client_builder.with_resend_rate_limit(burst, refill_per_min); } + client_builder = match alloc_meter { + Some(meter) => client_builder.with_alloc_meter(meter), + None => match task_instrument.clone() { + Some(instrument) => client_builder.with_task_instrument(instrument), + None => client_builder, + }, + }; + + let (client, sync_task_receiver) = client_builder.build().await?.into_parts(); Ok(Bot { client, diff --git a/src/client/builder.rs b/src/client/builder.rs index 5749bf042..276a7bc88 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -1,13 +1,18 @@ +use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; use thiserror::Error; use super::Client; use crate::cache_config::CacheConfig; use crate::http::HttpClient; +use crate::store::error::StoreError; use crate::store::persistence_manager::PersistenceManager; use crate::sync_task::MajorSyncTask; use crate::transport::TransportFactory; +use crate::types::durability_hook::InboundDurabilityHook; +use crate::types::enc_handler::EncHandler; use wacore::runtime::Runtime; /// Result of constructing a [`Client`]. @@ -42,7 +47,7 @@ impl ClientBuild { } /// A validated client-construction failure. -#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Error, Clone, PartialEq, Eq)] #[non_exhaustive] pub enum ClientBuilderError { #[error("missing async runtime")] @@ -53,6 +58,8 @@ pub enum ClientBuilderError { MissingTransportFactory, #[error("missing HTTP client")] MissingHttpClient, + #[error("the configured backend does not support the inbound durability hook: {0}")] + UnsupportedDurabilityBackend(String), } /// Runtime-validated, low-level builder for [`Client`]. @@ -67,6 +74,14 @@ pub struct ClientBuilder { http_client: Option>, override_version: Option<(u32, u32, u32)>, cache_config: CacheConfig, + custom_enc_handlers: HashMap>, + inbound_durability_hook: Option>, + skip_history_sync: bool, + wanted_pre_key_count: Option, + resend_rate_limit: Option<(u32, u32)>, + task_instrument: Option>, + alloc_meter: Option>, + background_saver_interval: Option, } impl Default for ClientBuilder { @@ -85,6 +100,14 @@ impl ClientBuilder { http_client: None, override_version: None, cache_config: CacheConfig::default(), + custom_enc_handlers: HashMap::new(), + inbound_durability_hook: None, + skip_history_sync: false, + wanted_pre_key_count: None, + resend_rate_limit: None, + task_instrument: None, + alloc_meter: None, + background_saver_interval: None, } } @@ -148,6 +171,91 @@ impl ClientBuilder { self } + /// Register a handler for one encrypted payload type before the client starts. + pub fn with_enc_handler(mut self, payload_type: impl Into, handler: H) -> Self + where + H: EncHandler + 'static, + { + self.custom_enc_handlers + .insert(payload_type.into(), Arc::new(handler)); + self + } + + /// Register an already-shared encrypted payload handler. + pub fn with_enc_handler_arc( + mut self, + payload_type: impl Into, + handler: Arc, + ) -> Self { + self.custom_enc_handlers + .insert(payload_type.into(), handler); + self + } + + pub(crate) fn with_custom_enc_handlers( + mut self, + handlers: HashMap>, + ) -> Self { + self.custom_enc_handlers = handlers; + self + } + + /// Install the durable-inbound hook after verifying backend support. + pub fn with_inbound_durability_hook(mut self, hook: H) -> Self + where + H: InboundDurabilityHook + 'static, + { + self.inbound_durability_hook = Some(Arc::new(hook)); + self + } + + /// Install an already-shared durable-inbound hook. + pub fn with_inbound_durability_hook_arc( + mut self, + hook: Arc, + ) -> Self { + self.inbound_durability_hook = Some(hook); + self + } + + pub fn with_skip_history_sync(mut self, skip: bool) -> Self { + self.skip_history_sync = skip; + self + } + + pub fn with_wanted_pre_key_count(mut self, count: usize) -> Self { + self.wanted_pre_key_count = Some(count); + self + } + + pub fn with_resend_rate_limit(mut self, burst: u32, refill_per_min: u32) -> Self { + self.resend_rate_limit = Some((burst, refill_per_min)); + self + } + + /// Instrument every task spawned through the configured runtime. + pub fn with_task_instrument( + mut self, + instrument: Arc, + ) -> Self { + self.task_instrument = Some(instrument); + self.alloc_meter = None; + self + } + + /// Install allocation attribution as the task instrument. + pub fn with_alloc_meter(mut self, meter: Arc) -> Self { + self.task_instrument = Some(meter.clone()); + self.alloc_meter = Some(meter); + self + } + + /// Run periodic device persistence for the lifetime of the client. + pub fn with_background_saver_interval(mut self, interval: Duration) -> Self { + self.background_saver_interval = Some(interval); + self + } + /// Validate dependencies, assemble an inert client, then start its services. pub async fn build(self) -> Result { self.build_boxed().await @@ -158,25 +266,32 @@ impl ClientBuilder { self, ) -> wacore::runtime::BoxFuture<'static, Result> { Box::pin(async move { - let runtime = self.runtime.ok_or(ClientBuilderError::MissingRuntime)?; + let runtime = self + .runtime + .as_ref() + .cloned() + .ok_or(ClientBuilderError::MissingRuntime)?; let persistence_manager = self .persistence_manager + .as_ref() + .cloned() .ok_or(ClientBuilderError::MissingPersistenceManager)?; let transport_factory = self .transport_factory + .as_ref() + .cloned() .ok_or(ClientBuilderError::MissingTransportFactory)?; let http_client = self .http_client + .as_ref() + .cloned() .ok_or(ClientBuilderError::MissingHttpClient)?; - Ok(Self::build_required( - runtime, - persistence_manager, - transport_factory, - http_client, - self.override_version, - self.cache_config, - )) + if self.inbound_durability_hook.is_some() { + probe_durability_backend(&persistence_manager.backend()).await?; + } + + Ok(self.finish(runtime, persistence_manager, transport_factory, http_client)) }) } @@ -188,18 +303,107 @@ impl ClientBuilder { override_version: Option<(u32, u32, u32)>, cache_config: CacheConfig, ) -> ClientBuild { - Client::assemble( - runtime, - persistence_manager, - transport_factory, - http_client, + Self { override_version, cache_config, - ) - .start() + ..Self::new() + } + .finish(runtime, persistence_manager, transport_factory, http_client) + } + + fn finish( + self, + runtime: Arc, + persistence_manager: Arc, + transport_factory: Arc, + http_client: Arc, + ) -> ClientBuild { + let runtime: Arc = match self.task_instrument { + Some(instrument) => { + Arc::new(wacore::stats::InstrumentedRuntime::new(runtime, instrument)) + } + None => runtime, + }; + + let assembly = Client::assemble( + Arc::clone(&runtime), + Arc::clone(&persistence_manager), + transport_factory, + http_client, + self.override_version, + self.cache_config, + ); + let client = assembly.client(); + + if !self.custom_enc_handlers.is_empty() { + let _ = client.custom_enc_handlers.set(self.custom_enc_handlers); + } + if let Some(hook) = self.inbound_durability_hook { + let _ = client.inbound_durability_hook.set(hook); + } + if self.skip_history_sync { + client.set_skip_history_sync(true); + } + if let Some(count) = self.wanted_pre_key_count { + client.set_wanted_pre_key_count(count); + } + if let Some((burst, refill_per_min)) = self.resend_rate_limit { + client.set_resend_rate_limit(burst, refill_per_min); + } + if let Some(meter) = self.alloc_meter { + let _ = client.alloc_meter.set(meter); + } + + let build = assembly.start(); + if let Some(interval) = self.background_saver_interval { + let saver_handle = persistence_manager.run_background_saver( + runtime, + interval, + build.client.shutdown_signal(), + ); + let _ = build.client.saver_handle.set(saver_handle); + } + build } } +async fn probe_durability_backend( + backend: &Arc, +) -> Result<(), ClientBuilderError> { + use portable_atomic::{AtomicU64, Ordering}; + + static PROBE_SEQ: AtomicU64 = AtomicU64::new(0); + const PROBE_JID: &str = "0@s.whatsapp.net"; + const PROBE_PAYLOAD: &[u8] = b"probe"; + let probe_id = format!( + "__wa_durability_probe_{}_{}__", + std::process::id(), + PROBE_SEQ.fetch_add(1, Ordering::Relaxed) + ); + let map_err = + |error: StoreError| ClientBuilderError::UnsupportedDurabilityBackend(error.to_string()); + + backend + .store_pending_inbound(PROBE_JID, PROBE_JID, &probe_id, PROBE_PAYLOAD) + .await + .map_err(map_err)?; + let stored = backend + .get_pending_inbound(PROBE_JID, PROBE_JID, &probe_id) + .await + .map_err(map_err)?; + backend + .delete_pending_inbound(PROBE_JID, PROBE_JID, &probe_id) + .await + .map_err(map_err)?; + + if stored.as_deref() != Some(PROBE_PAYLOAD) { + return Err(ClientBuilderError::UnsupportedDurabilityBackend( + "pending-inbound buffer did not round-trip".to_string(), + )); + } + Ok(()) +} + /// Owns the only path from a fully allocated client to started background /// services, preventing callers from publishing a partially configured client. pub(super) struct ClientAssembly { @@ -222,6 +426,10 @@ impl ClientAssembly { self.client.start_services(); ClientBuild::new(self.client, self.sync_task_receiver) } + + fn client(&self) -> Arc { + Arc::clone(&self.client) + } } #[cfg(test)] @@ -330,4 +538,43 @@ mod tests { assert_eq!(spawns.load(Ordering::SeqCst), 2); build.client().signal_shutdown_sync(); } + + #[tokio::test] + async fn low_level_builder_installs_options_and_owned_services() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let spawns = Arc::new(AtomicUsize::new(0)); + let meter = Arc::new(wacore::stats::AllocMeter::new()); + + let build = ClientBuilder::new() + .with_runtime(CountingRuntime { + spawns: Arc::clone(&spawns), + }) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_skip_history_sync(true) + .with_wanted_pre_key_count(123) + .with_alloc_meter(Arc::clone(&meter)) + .with_background_saver_interval(Duration::from_secs(3600)) + .build() + .await + .expect("complete builder"); + let client = build.client(); + + assert!(client.skip_history_sync_enabled()); + assert_eq!(client.wanted_pre_key_count(), 123); + assert!( + client + .alloc_meter + .get() + .is_some_and(|installed| Arc::ptr_eq(installed, &meter)) + ); + assert!(client.saver_handle.get().is_some()); + assert_eq!(spawns.load(Ordering::SeqCst), 3); + client.signal_shutdown_sync(); + } } diff --git a/src/types/enc_handler.rs b/src/types/enc_handler.rs index d2a763512..9e048baad 100644 --- a/src/types/enc_handler.rs +++ b/src/types/enc_handler.rs @@ -151,7 +151,7 @@ mod tests { .await .expect("Failed to build bot"); - // Verify no custom handlers are registered (the map is set, just empty) - assert_eq!(bot.client().custom_enc_handlers.get().unwrap().len(), 0); + // Keep the hot-path map unallocated when no extension uses it. + assert!(bot.client().custom_enc_handlers.get().is_none()); } } From de5d5bc5b7f457476cdb0c4c83724817e210c203 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 01:51:31 -0300 Subject: [PATCH 03/46] feat(client): add generation-scoped extension lifecycle --- src/client.rs | 5 + src/client/builder.rs | 110 ++++- src/client/extension_lifecycle.rs | 673 ++++++++++++++++++++++++++++++ src/client/lifecycle.rs | 46 +- src/client/node_io.rs | 7 +- src/lib.rs | 5 +- 6 files changed, 826 insertions(+), 20 deletions(-) create mode 100644 src/client/extension_lifecycle.rs diff --git a/src/client.rs b/src/client.rs index fcd1401b5..e7d544ad7 100644 --- a/src/client.rs +++ b/src/client.rs @@ -5,6 +5,7 @@ mod builder; mod context_impl; mod device_registry; pub(crate) mod device_topology; +mod extension_lifecycle; mod iq_ops; mod lid_pn; mod lifecycle; @@ -16,6 +17,8 @@ mod sessions; mod voip; use builder::ClientAssembly; pub use builder::{ClientBuild, ClientBuilder, ClientBuilderError}; +use extension_lifecycle::LifecycleRegistration; +pub use extension_lifecycle::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; pub use voip::{CallError, Voip}; use crate::cache::Cache; @@ -665,6 +668,8 @@ pub struct Client { /// error / connect_failure / disconnect. Per-connection subscribers /// (keepalive, request waiters, read loop, offline flush) observe this. pub(crate) connection_shutdown: std::sync::Mutex, + /// Allocated only when an extension host installs lifecycle callbacks. + lifecycle: Option>, /// Per-session wire I/O and activity counters. Written at the transport /// chokepoints (noise sender task, read loop); the keepalive dead-socket /// watchdog reads its activity timestamps. Snapshot via [`Client::stats`]. diff --git a/src/client/builder.rs b/src/client/builder.rs index 276a7bc88..014c0d702 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -4,7 +4,7 @@ use std::time::Duration; use thiserror::Error; -use super::Client; +use super::{Client, ClientLifecycle, LifecycleRegistration}; use crate::cache_config::CacheConfig; use crate::http::HttpClient; use crate::store::error::StoreError; @@ -47,7 +47,7 @@ impl ClientBuild { } /// A validated client-construction failure. -#[derive(Debug, Error, Clone, PartialEq, Eq)] +#[derive(Debug, Error)] #[non_exhaustive] pub enum ClientBuilderError { #[error("missing async runtime")] @@ -60,6 +60,8 @@ pub enum ClientBuilderError { MissingHttpClient, #[error("the configured backend does not support the inbound durability hook: {0}")] UnsupportedDurabilityBackend(String), + #[error("client lifecycle installation failed: {0}")] + LifecycleInstall(#[source] anyhow::Error), } /// Runtime-validated, low-level builder for [`Client`]. @@ -82,6 +84,7 @@ pub struct ClientBuilder { task_instrument: Option>, alloc_meter: Option>, background_saver_interval: Option, + lifecycle: Option>, } impl Default for ClientBuilder { @@ -108,6 +111,7 @@ impl ClientBuilder { task_instrument: None, alloc_meter: None, background_saver_interval: None, + lifecycle: None, } } @@ -256,6 +260,21 @@ impl ClientBuilder { self } + /// Install the aggregate lifecycle used by extensions of this client. + pub fn with_lifecycle(mut self, lifecycle: L) -> Self + where + L: ClientLifecycle + 'static, + { + self.lifecycle = Some(Arc::new(lifecycle)); + self + } + + /// Install an already-shared aggregate lifecycle. + pub fn with_lifecycle_arc(mut self, lifecycle: Arc) -> Self { + self.lifecycle = Some(lifecycle); + self + } + /// Validate dependencies, assemble an inert client, then start its services. pub async fn build(self) -> Result { self.build_boxed().await @@ -291,11 +310,12 @@ impl ClientBuilder { probe_durability_backend(&persistence_manager.backend()).await?; } - Ok(self.finish(runtime, persistence_manager, transport_factory, http_client)) + self.finish(runtime, persistence_manager, transport_factory, http_client) + .await }) } - pub(crate) fn build_required( + pub(crate) async fn build_required( runtime: Arc, persistence_manager: Arc, transport_factory: Arc, @@ -303,21 +323,26 @@ impl ClientBuilder { override_version: Option<(u32, u32, u32)>, cache_config: CacheConfig, ) -> ClientBuild { - Self { + let result = Self { override_version, cache_config, ..Self::new() } .finish(runtime, persistence_manager, transport_factory, http_client) + .await; + match result { + Ok(build) => build, + Err(error) => unreachable!("default lifecycle-free build failed: {error}"), + } } - fn finish( + async fn finish( self, runtime: Arc, persistence_manager: Arc, transport_factory: Arc, http_client: Arc, - ) -> ClientBuild { + ) -> Result { let runtime: Arc = match self.task_instrument { Some(instrument) => { Arc::new(wacore::stats::InstrumentedRuntime::new(runtime, instrument)) @@ -325,6 +350,7 @@ impl ClientBuilder { None => runtime, }; + let lifecycle = self.lifecycle.map(LifecycleRegistration::new).map(Arc::new); let assembly = Client::assemble( Arc::clone(&runtime), Arc::clone(&persistence_manager), @@ -332,6 +358,7 @@ impl ClientBuilder { http_client, self.override_version, self.cache_config, + lifecycle, ); let client = assembly.client(); @@ -353,6 +380,12 @@ impl ClientBuilder { if let Some(meter) = self.alloc_meter { let _ = client.alloc_meter.set(meter); } + if let Some(lifecycle) = &client.lifecycle { + lifecycle + .install(Arc::downgrade(&client)) + .await + .map_err(ClientBuilderError::LifecycleInstall)?; + } let build = assembly.start(); if let Some(interval) = self.background_saver_interval { @@ -363,7 +396,7 @@ impl ClientBuilder { ); let _ = build.client.saver_handle.set(saver_handle); } - build + Ok(build) } } @@ -445,6 +478,27 @@ mod tests { use crate::transport::mock::MockTransportFactory; use wacore::runtime::AbortHandle; + struct FailingLifecycle { + spawns: Arc, + installed_client: std::sync::Mutex>>, + } + + impl ClientLifecycle for FailingLifecycle { + fn install<'a>( + &'a self, + client: std::sync::Weak, + ) -> wacore::runtime::BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + assert_eq!(self.spawns.load(Ordering::SeqCst), 0); + *self + .installed_client + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(client); + Err(anyhow::anyhow!("injected install failure")) + }) + } + } + struct CountingRuntime { spawns: Arc, } @@ -531,6 +585,7 @@ mod tests { Arc::new(MockHttpClient), None, CacheConfig::default(), + None, ); assert_eq!(spawns.load(Ordering::SeqCst), 0); @@ -577,4 +632,43 @@ mod tests { assert_eq!(spawns.load(Ordering::SeqCst), 3); client.signal_shutdown_sync(); } + + #[tokio::test] + async fn lifecycle_install_failure_publishes_nothing_and_starts_no_tasks() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let spawns = Arc::new(AtomicUsize::new(0)); + let lifecycle = Arc::new(FailingLifecycle { + spawns: Arc::clone(&spawns), + installed_client: std::sync::Mutex::new(None), + }); + + let result = ClientBuilder::new() + .with_runtime(CountingRuntime { + spawns: Arc::clone(&spawns), + }) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await; + + assert!(matches!( + result, + Err(ClientBuilderError::LifecycleInstall(_)) + )); + assert_eq!(spawns.load(Ordering::SeqCst), 0); + assert!( + lifecycle + .installed_client + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .as_ref() + .is_some_and(|client| client.upgrade().is_none()) + ); + } } diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs new file mode 100644 index 000000000..ea48221bf --- /dev/null +++ b/src/client/extension_lifecycle.rs @@ -0,0 +1,673 @@ +use std::fmt; +use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; +use std::sync::{Arc, Weak}; + +use super::Client; +use wacore::runtime::{BoxFuture, ShutdownNotifier, ShutdownSignal}; + +const SCOPE_OPEN: u8 = 0; +const SCOPE_READY: u8 = 1; +const SCOPE_CANCELLED: u8 = 2; +const SCOPE_CLOSED: u8 = 3; + +/// Observable state of one authenticated connection generation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum ConnectionScopeState { + Open, + Ready, + Cancelled, + Closed, +} + +struct ConnectionScopeInner { + generation: u64, + state: AtomicU8, + cancellation: ShutdownNotifier, +} + +/// Stable handle for work owned by one authenticated connection generation. +/// +/// A scope is cancelled synchronously when its generation is retired and is +/// marked closed only after the client's authoritative connection cleanup. +#[derive(Clone)] +pub struct ConnectionScope { + inner: Arc, +} + +impl ConnectionScope { + fn new(generation: u64) -> Self { + Self { + inner: Arc::new(ConnectionScopeInner { + generation, + state: AtomicU8::new(SCOPE_OPEN), + cancellation: ShutdownNotifier::new(), + }), + } + } + + pub fn generation(&self) -> u64 { + self.inner.generation + } + + pub fn state(&self) -> ConnectionScopeState { + match self.inner.state.load(Ordering::Acquire) { + SCOPE_OPEN => ConnectionScopeState::Open, + SCOPE_READY => ConnectionScopeState::Ready, + SCOPE_CANCELLED => ConnectionScopeState::Cancelled, + _ => ConnectionScopeState::Closed, + } + } + + pub fn cancellation_signal(&self) -> ShutdownSignal { + self.inner.cancellation.subscribe() + } + + pub fn is_cancelled(&self) -> bool { + self.inner.state.load(Ordering::Acquire) >= SCOPE_CANCELLED + } + + fn mark_ready(&self) -> bool { + self.inner + .state + .compare_exchange(SCOPE_OPEN, SCOPE_READY, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + } + + fn cancel(&self) { + let mut state = self.inner.state.load(Ordering::Acquire); + while state < SCOPE_CANCELLED { + match self.inner.state.compare_exchange_weak( + state, + SCOPE_CANCELLED, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => { + self.inner.cancellation.notify(); + return; + } + Err(actual) => state = actual, + } + } + } + + fn close(&self) { + let previous = self.inner.state.swap(SCOPE_CLOSED, Ordering::AcqRel); + if previous < SCOPE_CANCELLED { + self.inner.cancellation.notify(); + } + } +} + +impl fmt::Debug for ConnectionScope { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ConnectionScope") + .field("generation", &self.generation()) + .field("state", &self.state()) + .finish() + } +} + +/// Aggregate lifecycle seam installed during [`Client`](super::Client) construction. +/// +/// Implementations must make `install` transactional. The remaining callbacks +/// are serialized, awaited, and must not block indefinitely. A future plugin +/// host owns per-plugin ordering, timeout, and error isolation behind this one +/// client-level seam. `install` receives a weak client reference so retaining +/// the handle cannot create a client-to-extension reference cycle. +pub trait ClientLifecycle: wacore::sync_marker::MaybeSendSync { + fn install<'a>(&'a self, _client: Weak) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } + + fn on_ready<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } + + fn on_closed<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } +} + +pub(super) struct LifecycleRegistration { + handler: Arc, + active_scope: std::sync::Mutex>, + callback_gate: async_lock::Mutex<()>, + shutdown_gate: async_lock::Mutex, + scope_changed: event_listener::Event, + terminal: AtomicBool, +} + +impl LifecycleRegistration { + pub(super) fn new(handler: Arc) -> Self { + Self { + handler, + active_scope: std::sync::Mutex::new(None), + callback_gate: async_lock::Mutex::new(()), + shutdown_gate: async_lock::Mutex::new(false), + scope_changed: event_listener::Event::new(), + terminal: AtomicBool::new(false), + } + } + + pub(super) async fn install(&self, client: Weak) -> anyhow::Result<()> { + self.handler.install(client).await + } + + pub(super) fn begin_scope(&self, generation: u64) { + if self.terminal.load(Ordering::Acquire) { + return; + } + + let scope = ConnectionScope::new(generation); + let mut active = self + .active_scope + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.terminal.load(Ordering::Acquire) { + return; + } + let replaced = active.replace(scope); + drop(active); + if let Some(replaced) = replaced { + log::warn!( + "Replacing unclosed connection scope for generation {}", + replaced.generation() + ); + replaced.cancel(); + replaced.close(); + } + } + + pub(super) async fn ready(&self, generation: u64) -> bool { + let _callback_guard = self.callback_gate.lock().await; + if self.terminal.load(Ordering::Acquire) { + return false; + } + let scope = self.scope_for(generation); + let Some(scope) = scope.filter(ConnectionScope::mark_ready) else { + return false; + }; + + if let Err(error) = self.handler.on_ready(scope.clone()).await { + log::warn!("Client lifecycle on_ready failed: {error:#}"); + } + !scope.is_cancelled() + } + + pub(super) fn cancel_scope(&self, generation: u64) { + if let Some(scope) = self.scope_for(generation) { + scope.cancel(); + } + } + + pub(super) fn cancel_active_scope(&self) { + let scope = self + .active_scope + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone(); + if let Some(scope) = scope { + scope.cancel(); + } + } + + pub(super) async fn close_scope(&self, generation: u64) { + let _callback_guard = self.callback_gate.lock().await; + let scope = { + let mut active = self + .active_scope + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if active + .as_ref() + .is_some_and(|scope| scope.generation() == generation) + { + active.take() + } else { + None + } + }; + let Some(scope) = scope else { + return; + }; + + self.scope_changed.notify(usize::MAX); + scope.close(); + if let Err(error) = self.handler.on_closed(scope).await { + log::warn!("Client lifecycle on_closed failed: {error:#}"); + } + } + + pub(super) async fn shutdown(&self) { + self.terminal.store(true, Ordering::Release); + self.cancel_active_scope(); + + let mut started = self.shutdown_gate.lock().await; + if *started { + return; + } + *started = true; + + loop { + let scope_changed = self.scope_changed.listen(); + let callback_guard = self.callback_gate.lock().await; + if self.scope_for_active_generation().is_none() { + if let Err(error) = self.handler.shutdown().await { + log::warn!("Client lifecycle shutdown failed: {error:#}"); + } + return; + } + drop(callback_guard); + scope_changed.await; + } + } + + fn scope_for_active_generation(&self) -> Option { + self.active_scope + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + } + + fn scope_for(&self, generation: u64) -> Option { + self.active_scope + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .as_ref() + .filter(|scope| scope.generation() == generation) + .cloned() + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + + use async_trait::async_trait; + use bytes::Bytes; + + use super::*; + use crate::runtime_impl::TokioRuntime; + use crate::store::persistence_manager::PersistenceManager; + use crate::test_utils::MockHttpClient; + use crate::transport::mock::MockTransportFactory; + + #[derive(Default)] + struct RecordingLifecycle { + events: std::sync::Mutex>, + scopes: std::sync::Mutex>, + shutdowns: AtomicUsize, + } + + impl RecordingLifecycle { + fn events(&self) -> Vec { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + } + } + + impl ClientLifecycle for RecordingLifecycle { + fn install<'a>(&'a self, client: Weak) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + assert!(client.upgrade().is_some()); + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("install".to_string()); + Ok(()) + }) + } + + fn on_ready<'a>(&'a self, scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(format!("ready:{}", scope.generation())); + self.scopes + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(scope); + Ok(()) + }) + } + + fn on_closed<'a>(&'a self, scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(format!("closed:{}", scope.generation())); + Ok(()) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + self.shutdowns.fetch_add(1, Ordering::SeqCst); + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("shutdown".to_string()); + Ok(()) + }) + } + } + + struct BlockingDisconnect { + started: async_channel::Sender<()>, + release: async_channel::Receiver<()>, + } + + #[async_trait] + impl crate::transport::Transport for BlockingDisconnect { + async fn send(&self, _data: Bytes) -> anyhow::Result<()> { + Ok(()) + } + + async fn disconnect(&self) { + let _ = self.started.try_send(()); + let _ = self.release.recv().await; + } + } + + struct BlockingReadyLifecycle { + ready_started: async_channel::Sender<()>, + release_ready: async_channel::Receiver<()>, + scope: std::sync::Mutex>, + events: std::sync::Mutex>, + } + + struct LogoutOrderHandler { + lifecycle: Arc, + } + + impl wacore::types::events::EventHandler for LogoutOrderHandler { + fn handle_event(&self, event: Arc) { + if matches!(&*event, wacore::types::events::Event::LoggedOut(_)) { + self.lifecycle + .events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("logged-out".to_string()); + } + } + + fn interest(&self) -> wacore::types::events::EventInterest { + wacore::types::events::EventInterest::of(&[wacore::types::events::EventKind::LoggedOut]) + } + } + + impl ClientLifecycle for BlockingReadyLifecycle { + fn on_ready<'a>(&'a self, scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + *self + .scope + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(scope); + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("ready-started"); + let _ = self.ready_started.try_send(()); + let _ = self.release_ready.recv().await; + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("ready-finished"); + Ok(()) + }) + } + + fn on_closed<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("closed"); + Ok(()) + }) + } + } + + #[test] + fn scope_state_machine_is_sticky_and_cancellable() { + let scope = ConnectionScope::new(41); + let cancellation = scope.cancellation_signal(); + + assert_eq!(scope.state(), ConnectionScopeState::Open); + assert!(scope.mark_ready()); + assert_eq!(scope.state(), ConnectionScopeState::Ready); + scope.cancel(); + assert_eq!(scope.state(), ConnectionScopeState::Cancelled); + assert!(cancellation.is_fired()); + scope.cancel(); + scope.close(); + assert_eq!(scope.state(), ConnectionScopeState::Closed); + } + + #[tokio::test] + async fn cleanup_cancels_before_io_and_closes_after_authoritative_teardown() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let lifecycle = Arc::new(RecordingLifecycle::default()); + let build = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await + .expect("client build"); + let client = build.client(); + const GENERATION: u64 = 9; + client + .connection_generation + .store(GENERATION, Ordering::SeqCst); + let registration = client.lifecycle.as_ref().expect("lifecycle registration"); + registration.begin_scope(GENERATION); + client.dispatch_connected().await; + + let scope = lifecycle + .scopes + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .first() + .cloned() + .expect("ready scope"); + let cancelled = scope.cancellation_signal(); + let (started_tx, started_rx) = async_channel::bounded(1); + let (release_tx, release_rx) = async_channel::bounded(1); + *client.transport.lock().await = Some(Arc::new(BlockingDisconnect { + started: started_tx, + release: release_rx, + })); + + let cleanup_client = Arc::clone(&client); + let cleanup = tokio::spawn(async move { + cleanup_client.cleanup_connection_state().await; + }); + tokio::time::timeout(std::time::Duration::from_secs(2), started_rx.recv()) + .await + .expect("transport cleanup started") + .expect("transport remained alive"); + + assert_eq!(scope.state(), ConnectionScopeState::Cancelled); + assert!(cancelled.is_fired()); + assert_eq!(lifecycle.events(), vec!["install", "ready:9"]); + + release_tx.send(()).await.expect("release cleanup"); + tokio::time::timeout(std::time::Duration::from_secs(2), cleanup) + .await + .expect("cleanup completed") + .expect("cleanup did not panic"); + + assert_eq!(scope.state(), ConnectionScopeState::Closed); + assert_eq!(lifecycle.events(), vec!["install", "ready:9", "closed:9"]); + client.shutdown_lifecycle().await; + client.shutdown_lifecycle().await; + assert_eq!(lifecycle.shutdowns.load(Ordering::SeqCst), 1); + assert_eq!( + lifecycle.events(), + vec!["install", "ready:9", "closed:9", "shutdown"] + ); + client.signal_shutdown_sync(); + } + + #[tokio::test] + async fn cancellation_does_not_wait_for_a_running_callback() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let (ready_started_tx, ready_started_rx) = async_channel::bounded(1); + let (release_ready_tx, release_ready_rx) = async_channel::bounded(1); + let lifecycle = Arc::new(BlockingReadyLifecycle { + ready_started: ready_started_tx, + release_ready: release_ready_rx, + scope: std::sync::Mutex::new(None), + events: std::sync::Mutex::new(Vec::new()), + }); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await + .expect("client build") + .client(); + const GENERATION: u64 = 13; + client + .connection_generation + .store(GENERATION, Ordering::SeqCst); + client + .lifecycle + .as_ref() + .expect("lifecycle registration") + .begin_scope(GENERATION); + + let ready_client = Arc::clone(&client); + let ready_task = tokio::spawn(async move { + ready_client.dispatch_connected().await; + }); + ready_started_rx + .recv() + .await + .expect("ready callback started"); + let scope = lifecycle + .scope + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + .expect("ready scope"); + + let cleanup_client = Arc::clone(&client); + let cleanup_task = tokio::spawn(async move { + cleanup_client.cleanup_connection_state().await; + }); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while !scope.is_cancelled() { + tokio::task::yield_now().await; + } + }) + .await + .expect("scope cancellation"); + assert!(!cleanup_task.is_finished()); + + release_ready_tx.send(()).await.expect("release ready hook"); + ready_task.await.expect("ready task did not panic"); + cleanup_task.await.expect("cleanup task did not panic"); + assert_eq!(scope.state(), ConnectionScopeState::Closed); + assert_eq!( + *lifecycle + .events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["ready-started", "ready-finished", "closed"] + ); + client.signal_shutdown_sync(); + } + + #[tokio::test] + async fn shutdown_waits_for_the_active_scope_and_is_terminal() { + let lifecycle = Arc::new(RecordingLifecycle::default()); + let registration = Arc::new(LifecycleRegistration::new(lifecycle.clone())); + const GENERATION: u64 = 21; + registration.begin_scope(GENERATION); + let scope = registration + .scope_for(GENERATION) + .expect("active connection scope"); + + let shutdown_registration = Arc::clone(®istration); + let shutdown = tokio::spawn(async move { + shutdown_registration.shutdown().await; + }); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while !scope.is_cancelled() { + tokio::task::yield_now().await; + } + }) + .await + .expect("scope cancellation"); + assert!(!shutdown.is_finished()); + + registration.close_scope(GENERATION).await; + shutdown.await.expect("shutdown did not panic"); + assert_eq!( + lifecycle.events(), + vec!["closed:21".to_string(), "shutdown".to_string()] + ); + + registration.begin_scope(GENERATION + 1); + assert!(registration.scope_for(GENERATION + 1).is_none()); + assert!(!registration.ready(GENERATION + 1).await); + registration.shutdown().await; + assert_eq!(lifecycle.shutdowns.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn logout_event_precedes_terminal_lifecycle_shutdown() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let lifecycle = Arc::new(RecordingLifecycle::default()); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await + .expect("client build") + .client(); + client.register_handler(Arc::new(LogoutOrderHandler { + lifecycle: lifecycle.clone(), + })); + + client.logout().await.expect("client logout"); + + assert_eq!( + lifecycle.events(), + vec!["install", "logged-out", "shutdown"] + ); + } +} diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index c3ab2621c..a00cec185 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -30,6 +30,9 @@ impl Client { self.expected_disconnect.store(true, Ordering::Relaxed); self.is_running.store(false, Ordering::Relaxed); self.shutdown_notifier.notify(); + if let Some(lifecycle) = &self.lifecycle { + lifecycle.cancel_active_scope(); + } self.notify_connection_shutdown(); } @@ -72,7 +75,18 @@ impl Client { } /// Dispatch the Connected event and notify waiters. - pub(crate) fn dispatch_connected(&self) { + pub(crate) async fn dispatch_connected(&self) { + let generation = self.connection_generation.load(Ordering::SeqCst); + if let Some(lifecycle) = &self.lifecycle { + if !lifecycle.ready(generation).await { + debug!("Skipping Connected dispatch for retired generation {generation}"); + return; + } + if self.connection_generation.load(Ordering::SeqCst) != generation { + debug!("Skipping Connected dispatch after generation changed"); + return; + } + } self.is_ready.store(true, Ordering::Relaxed); wacore::telemetry::set_connected(true); self.core.event_bus.dispatch(Event::Connected( @@ -81,6 +95,12 @@ impl Client { self.connected_notifier.notify(usize::MAX); } + pub(super) async fn shutdown_lifecycle(&self) { + if let Some(lifecycle) = &self.lifecycle { + lifecycle.shutdown().await; + } + } + /// Create a new `Client` with default cache configuration. /// /// This is the standard constructor. Use [`Client::new_with_cache_config`] @@ -100,6 +120,7 @@ impl Client { override_version, CacheConfig::default(), ) + .await .into_parts() } @@ -120,6 +141,7 @@ impl Client { override_version, cache_config, ) + .await .into_parts() } @@ -130,6 +152,7 @@ impl Client { http_client: Arc, override_version: Option<(u32, u32, u32)>, cache_config: CacheConfig, + lifecycle: Option>, ) -> ClientAssembly { let mut unique_id_bytes = [0u8; 2]; rand::make_rng::().fill_bytes(&mut unique_id_bytes); @@ -157,6 +180,7 @@ impl Client { ik_handshake_failures: Arc::new(AtomicU32::new(0)), shutdown_notifier: wacore::runtime::ShutdownNotifier::new(), connection_shutdown: std::sync::Mutex::new(wacore::runtime::ShutdownNotifier::new()), + lifecycle, stats: Arc::new(wacore::stats::SessionStats::new()), transport: Arc::new(Mutex::new(None)), @@ -468,6 +492,7 @@ impl Client { ); self.runtime.sleep(delay).await; } + self.shutdown_lifecycle().await; info!("Client run loop has shut down."); } @@ -630,8 +655,6 @@ impl Client { warn!("Failed to send logout IQ: {e}"); } - self.disconnect().await; - self.core.event_bus.dispatch(Event::LoggedOut( crate::types::events::LoggedOut::builder() .on_connect(false) @@ -639,6 +662,8 @@ impl Client { .build(), )); + self.disconnect().await; + Ok(()) } @@ -698,6 +723,7 @@ impl Client { // final flush below and then be acked. self.msg_secret_buffer.seal(); self.msg_secret_buffer.flush().await; + self.shutdown_lifecycle().await; } /// Backoff step used by [`reconnect()`] to create an offline window. @@ -795,7 +821,11 @@ impl Client { // process_classified_message — no decrypt can START after the // permit-held cache settle below, so no rowless ratchet advances can // dirty the cache behind teardown's back. - self.connection_generation.fetch_add(1, Ordering::SeqCst); + let closed_generation = self.connection_generation.fetch_add(1, Ordering::SeqCst); + if let Some(lifecycle) = &self.lifecycle { + lifecycle.cancel_scope(closed_generation); + } + self.notify_connection_shutdown(); // The coalesced-flush scheduler needs no explicit reset: its state is // generation-scoped, so the bump above already hands ownership to the // next connection's first request and retires any stale worker. @@ -821,11 +851,6 @@ impl Client { // registry, so abort_all misses them. Drain them and notify `ended` so any waiter wakes. crate::voip::facade::drain_pending_outgoing_on_disconnect(self); } - // Signal the keepalive loop (and any other per-connection tasks) to - // exit promptly. Without this, a stale keepalive loop can overlap - // with the next one after reconnect. Uses the PER-CONNECTION signal - // so the terminal shutdown_notifier stays clean for reconnects. - self.notify_connection_shutdown(); // Close the socket as part of cleanup so this path is authoritative // even when reached via the run loop's graceful-exit flow (not just // `Client::disconnect()`). Transport impls make `disconnect()` @@ -955,6 +980,9 @@ impl Client { if let Some(proc) = self.app_state_processor.lock().await.as_ref() { proc.clear_key_cache().await; } + if let Some(lifecycle) = &self.lifecycle { + lifecycle.close_scope(closed_generation).await; + } } /// Waits for the noise socket to be established. diff --git a/src/client/node_io.rs b/src/client/node_io.rs index db6090a2e..cc1e6cb8f 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -697,6 +697,9 @@ impl Client { // Increment connection generation to invalidate any stale post-login tasks // from previous connections (e.g., during 515 reconnect cycles). let current_generation = self.connection_generation.fetch_add(1, Ordering::SeqCst) + 1; + if let Some(lifecycle) = &self.lifecycle { + lifecycle.begin_scope(current_generation); + } info!( "Successfully authenticated with WhatsApp servers! (gen={})", @@ -1104,7 +1107,7 @@ impl Client { // Presence is NOT sent here — WhatsApp Web sends presence from the // setting_pushName mutation handler (WAWebPushNameSync), not from // criticalSyncDone. Our setting_pushName handler already does this. - client_clone.dispatch_connected(); + client_clone.dispatch_connected().await; } Err(e) => { client_clone.log_sync_error("critical app state sync", &e); @@ -1166,7 +1169,7 @@ impl Client { // for an outdated connection that was replaced mid-await. check_generation!(); - client_clone.dispatch_connected(); + client_clone.dispatch_connected().await; } })).detach(); } diff --git a/src/lib.rs b/src/lib.rs index 029c24e20..de63783c1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -93,7 +93,10 @@ pub use client::{ StatsSnapshot, StorageResourceReport, TransportResourceReport, }; pub use client::{CallError, Voip}; -pub use client::{Client, ClientBuild, ClientBuilder, ClientBuilderError}; +pub use client::{ + Client, ClientBuild, ClientBuilder, ClientBuilderError, ClientLifecycle, ConnectionScope, + ConnectionScopeState, +}; pub use types::durability_hook::InboundDurabilityHook; pub use types::retry_admission::RetryAdmission; pub mod download; From 1f28c48b3446d170e664ecb0e2ca42283c2e5266 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 02:38:08 -0300 Subject: [PATCH 04/46] fix(client): retain superseded scopes until cleanup --- src/client/extension_lifecycle.rs | 97 +++++++++++++++++++++++-------- 1 file changed, 73 insertions(+), 24 deletions(-) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index ea48221bf..1acae7f7a 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -137,18 +137,24 @@ pub trait ClientLifecycle: wacore::sync_marker::MaybeSendSync { pub(super) struct LifecycleRegistration { handler: Arc, - active_scope: std::sync::Mutex>, + scopes: std::sync::Mutex, callback_gate: async_lock::Mutex<()>, shutdown_gate: async_lock::Mutex, scope_changed: event_listener::Event, terminal: AtomicBool, } +#[derive(Default)] +struct ScopeRegistry { + active: Option, + retired: Vec, +} + impl LifecycleRegistration { pub(super) fn new(handler: Arc) -> Self { Self { handler, - active_scope: std::sync::Mutex::new(None), + scopes: std::sync::Mutex::new(ScopeRegistry::default()), callback_gate: async_lock::Mutex::new(()), shutdown_gate: async_lock::Mutex::new(false), scope_changed: event_listener::Event::new(), @@ -166,22 +172,21 @@ impl LifecycleRegistration { } let scope = ConnectionScope::new(generation); - let mut active = self - .active_scope + let mut scopes = self + .scopes .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); if self.terminal.load(Ordering::Acquire) { return; } - let replaced = active.replace(scope); - drop(active); + let replaced = scopes.active.replace(scope); if let Some(replaced) = replaced { log::warn!( "Replacing unclosed connection scope for generation {}", replaced.generation() ); replaced.cancel(); - replaced.close(); + scopes.retired.push(replaced); } } @@ -208,12 +213,14 @@ impl LifecycleRegistration { } pub(super) fn cancel_active_scope(&self) { - let scope = self - .active_scope + let scopes = self + .scopes .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clone(); - if let Some(scope) = scope { + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(scope) = &scopes.active { + scope.cancel(); + } + for scope in &scopes.retired { scope.cancel(); } } @@ -221,17 +228,22 @@ impl LifecycleRegistration { pub(super) async fn close_scope(&self, generation: u64) { let _callback_guard = self.callback_gate.lock().await; let scope = { - let mut active = self - .active_scope + let mut scopes = self + .scopes .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); - if active + if scopes + .active .as_ref() .is_some_and(|scope| scope.generation() == generation) { - active.take() + scopes.active.take() } else { - None + scopes + .retired + .iter() + .position(|scope| scope.generation() == generation) + .map(|position| scopes.retired.remove(position)) } }; let Some(scope) = scope else { @@ -258,7 +270,7 @@ impl LifecycleRegistration { loop { let scope_changed = self.scope_changed.listen(); let callback_guard = self.callback_gate.lock().await; - if self.scope_for_active_generation().is_none() { + if !self.has_open_scopes() { if let Err(error) = self.handler.shutdown().await { log::warn!("Client lifecycle shutdown failed: {error:#}"); } @@ -269,19 +281,29 @@ impl LifecycleRegistration { } } - fn scope_for_active_generation(&self) -> Option { - self.active_scope + fn has_open_scopes(&self) -> bool { + let scopes = self + .scopes .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clone() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + scopes.active.is_some() || !scopes.retired.is_empty() } fn scope_for(&self, generation: u64) -> Option { - self.active_scope + let scopes = self + .scopes .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) + .unwrap_or_else(|poisoned| poisoned.into_inner()); + scopes + .active .as_ref() .filter(|scope| scope.generation() == generation) + .or_else(|| { + scopes + .retired + .iter() + .find(|scope| scope.generation() == generation) + }) .cloned() } } @@ -604,6 +626,33 @@ mod tests { client.signal_shutdown_sync(); } + #[tokio::test] + async fn replaced_scope_stays_closeable_by_its_generation() { + let lifecycle = Arc::new(RecordingLifecycle::default()); + let registration = LifecycleRegistration::new(lifecycle.clone()); + + registration.begin_scope(31); + assert!(registration.ready(31).await); + let first = registration.scope_for(31).expect("first scope"); + + registration.begin_scope(32); + assert_eq!(first.state(), ConnectionScopeState::Cancelled); + assert!(registration.scope_for(31).is_some()); + assert!(registration.ready(32).await); + + registration.close_scope(31).await; + assert_eq!(first.state(), ConnectionScopeState::Closed); + assert!(registration.scope_for(31).is_none()); + assert!(registration.scope_for(32).is_some()); + + registration.close_scope(32).await; + registration.shutdown().await; + assert_eq!( + lifecycle.events(), + vec!["ready:31", "ready:32", "closed:31", "closed:32", "shutdown",] + ); + } + #[tokio::test] async fn shutdown_waits_for_the_active_scope_and_is_terminal() { let lifecycle = Arc::new(RecordingLifecycle::default()); From a0da3e2e9637c3261b5340fd18315eac0de3fa5b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 03:05:36 -0300 Subject: [PATCH 05/46] fix(client): harden extension lifecycle callbacks --- src/client/builder.rs | 5 +- src/client/extension_lifecycle.rs | 587 +++++++++++++++++++++++++----- src/client/lifecycle.rs | 2 +- src/client/node_io.rs | 9 +- src/lib.rs | 5 +- 5 files changed, 520 insertions(+), 88 deletions(-) diff --git a/src/client/builder.rs b/src/client/builder.rs index 014c0d702..3fc8e56de 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -350,7 +350,10 @@ impl ClientBuilder { None => runtime, }; - let lifecycle = self.lifecycle.map(LifecycleRegistration::new).map(Arc::new); + let lifecycle = self + .lifecycle + .map(|handler| LifecycleRegistration::new(handler, Arc::clone(&runtime))) + .map(Arc::new); let assembly = Client::assemble( Arc::clone(&runtime), Arc::clone(&persistence_manager), diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 1acae7f7a..a2c09da7a 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -1,14 +1,25 @@ +use std::cell::Cell; +use std::collections::VecDeque; use std::fmt; use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; use std::sync::{Arc, Weak}; +use std::time::Duration; use super::Client; -use wacore::runtime::{BoxFuture, ShutdownNotifier, ShutdownSignal}; +use futures::FutureExt; +use wacore::runtime::{ + BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, timeout as rt_timeout, +}; const SCOPE_OPEN: u8 = 0; const SCOPE_READY: u8 = 1; const SCOPE_CANCELLED: u8 = 2; const SCOPE_CLOSED: u8 = 3; +const CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); + +std::thread_local! { + static ACTIVE_CALLBACK: Cell<*const LifecycleRegistration> = const { Cell::new(std::ptr::null()) }; +} /// Observable state of one authenticated connection generation. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -112,11 +123,11 @@ impl fmt::Debug for ConnectionScope { /// Aggregate lifecycle seam installed during [`Client`](super::Client) construction. /// -/// Implementations must make `install` transactional. The remaining callbacks -/// are serialized, awaited, and must not block indefinitely. A future plugin -/// host owns per-plugin ordering, timeout, and error isolation behind this one -/// client-level seam. `install` receives a weak client reference so retaining -/// the handle cannot create a client-to-extension reference cycle. +/// Implementations must make `install` transactional. Connection callbacks are +/// serialized and bounded; connection cleanup only schedules `on_closed` so a +/// stalled extension cannot block reconnect. A future plugin host owns +/// per-plugin ordering and isolation behind this client-level seam. `install` +/// receives a weak client reference so retaining it cannot create a cycle. pub trait ClientLifecycle: wacore::sync_marker::MaybeSendSync { fn install<'a>(&'a self, _client: Weak) -> BoxFuture<'a, anyhow::Result<()>> { Box::pin(async { Ok(()) }) @@ -137,10 +148,12 @@ pub trait ClientLifecycle: wacore::sync_marker::MaybeSendSync { pub(super) struct LifecycleRegistration { handler: Arc, + runtime: Arc, scopes: std::sync::Mutex, - callback_gate: async_lock::Mutex<()>, - shutdown_gate: async_lock::Mutex, - scope_changed: event_listener::Event, + callback_queue: std::sync::Mutex, + shutdown_complete: AtomicBool, + shutdown_notifier: ShutdownNotifier, + callback_timeout: Duration, terminal: AtomicBool, } @@ -150,14 +163,62 @@ struct ScopeRegistry { retired: Vec, } +enum LifecycleCallback { + Ready { + scope: ConnectionScope, + done: async_channel::Sender, + }, + Closed(ConnectionScope), + Shutdown, +} + +#[derive(Default)] +struct CallbackQueue { + pending: VecDeque, + shutdown_requested: bool, + shutdown_enqueued: bool, + drain_scheduled: bool, +} + +struct CallbackContextGuard { + previous: *const LifecycleRegistration, +} + +impl CallbackContextGuard { + fn enter(registration: &LifecycleRegistration) -> Self { + let previous = ACTIVE_CALLBACK.replace(registration); + Self { previous } + } +} + +impl Drop for CallbackContextGuard { + fn drop(&mut self) { + ACTIVE_CALLBACK.set(self.previous); + } +} + +fn callback_context_active(registration: &LifecycleRegistration) -> bool { + ACTIVE_CALLBACK.with(|active| std::ptr::eq(active.get(), registration)) +} + impl LifecycleRegistration { - pub(super) fn new(handler: Arc) -> Self { + pub(super) fn new(handler: Arc, runtime: Arc) -> Self { + Self::new_with_timeout(handler, runtime, CALLBACK_TIMEOUT) + } + + fn new_with_timeout( + handler: Arc, + runtime: Arc, + callback_timeout: Duration, + ) -> Self { Self { handler, + runtime, scopes: std::sync::Mutex::new(ScopeRegistry::default()), - callback_gate: async_lock::Mutex::new(()), - shutdown_gate: async_lock::Mutex::new(false), - scope_changed: event_listener::Event::new(), + callback_queue: std::sync::Mutex::new(CallbackQueue::default()), + shutdown_complete: AtomicBool::new(false), + shutdown_notifier: ShutdownNotifier::new(), + callback_timeout, terminal: AtomicBool::new(false), } } @@ -166,18 +227,19 @@ impl LifecycleRegistration { self.handler.install(client).await } - pub(super) fn begin_scope(&self, generation: u64) { + pub(super) fn begin_scope_if_current( + &self, + generation: u64, + is_current: impl FnOnce() -> bool, + ) -> bool { if self.terminal.load(Ordering::Acquire) { - return; + return false; } let scope = ConnectionScope::new(generation); - let mut scopes = self - .scopes - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if self.terminal.load(Ordering::Acquire) { - return; + let mut scopes = self.scopes(); + if self.terminal.load(Ordering::Acquire) || !is_current() { + return false; } let replaced = scopes.active.replace(scope); if let Some(replaced) = replaced { @@ -188,22 +250,38 @@ impl LifecycleRegistration { replaced.cancel(); scopes.retired.push(replaced); } + true } - pub(super) async fn ready(&self, generation: u64) -> bool { - let _callback_guard = self.callback_gate.lock().await; + pub(super) async fn ready(self: &Arc, generation: u64) -> bool { if self.terminal.load(Ordering::Acquire) { return false; } - let scope = self.scope_for(generation); - let Some(scope) = scope.filter(ConnectionScope::mark_ready) else { - return false; + let (done_tx, done_rx) = async_channel::bounded(1); + let scope = { + let scopes = self.scopes(); + let scope = scopes + .active + .as_ref() + .filter(|scope| scope.generation() == generation) + .or_else(|| { + scopes + .retired + .iter() + .find(|scope| scope.generation() == generation) + }) + .cloned(); + let Some(scope) = scope.filter(ConnectionScope::mark_ready) else { + return false; + }; + self.enqueue_callback(LifecycleCallback::Ready { + scope: scope.clone(), + done: done_tx, + }); + scope }; - if let Err(error) = self.handler.on_ready(scope.clone()).await { - log::warn!("Client lifecycle on_ready failed: {error:#}"); - } - !scope.is_cancelled() + done_rx.recv().await.unwrap_or(false) && !scope.is_cancelled() } pub(super) fn cancel_scope(&self, generation: u64) { @@ -213,10 +291,7 @@ impl LifecycleRegistration { } pub(super) fn cancel_active_scope(&self) { - let scopes = self - .scopes - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let scopes = self.scopes(); if let Some(scope) = &scopes.active { scope.cancel(); } @@ -225,13 +300,9 @@ impl LifecycleRegistration { } } - pub(super) async fn close_scope(&self, generation: u64) { - let _callback_guard = self.callback_gate.lock().await; + pub(super) fn close_scope(self: &Arc, generation: u64) { let scope = { - let mut scopes = self - .scopes - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut scopes = self.scopes(); if scopes .active .as_ref() @@ -250,50 +321,153 @@ impl LifecycleRegistration { return; }; - self.scope_changed.notify(usize::MAX); scope.close(); - if let Err(error) = self.handler.on_closed(scope).await { - log::warn!("Client lifecycle on_closed failed: {error:#}"); - } + self.enqueue_callback(LifecycleCallback::Closed(scope)); + self.enqueue_shutdown_if_ready(); } - pub(super) async fn shutdown(&self) { + pub(super) async fn shutdown(self: &Arc) { self.terminal.store(true, Ordering::Release); self.cancel_active_scope(); + { + let mut queue = self.callback_queue(); + queue.shutdown_requested = true; + } + self.enqueue_shutdown_if_ready(); + + if self.shutdown_complete.load(Ordering::Acquire) || callback_context_active(self) { + return; + } + + let completed = self.shutdown_notifier.subscribe(); + if self.shutdown_complete.load(Ordering::Acquire) { + return; + } + wacore::runtime::wait_for_shutdown(&completed).await; + } - let mut started = self.shutdown_gate.lock().await; - if *started { + fn enqueue_callback(self: &Arc, callback: LifecycleCallback) { + let should_spawn = { + let mut queue = self.callback_queue(); + queue.pending.push_back(callback); + if queue.drain_scheduled { + false + } else { + queue.drain_scheduled = true; + true + } + }; + self.spawn_callback_driver(should_spawn); + } + + fn enqueue_shutdown_if_ready(self: &Arc) { + let no_open_scopes = { + let scopes = self.scopes(); + scopes.active.is_none() && scopes.retired.is_empty() + }; + if !no_open_scopes { return; } - *started = true; + let should_spawn = { + let mut queue = self.callback_queue(); + if !queue.shutdown_requested || queue.shutdown_enqueued { + return; + } + queue.shutdown_enqueued = true; + queue.pending.push_back(LifecycleCallback::Shutdown); + if queue.drain_scheduled { + false + } else { + queue.drain_scheduled = true; + true + } + }; + self.spawn_callback_driver(should_spawn); + } + + fn spawn_callback_driver(self: &Arc, should_spawn: bool) { + if !should_spawn { + return; + } + let registration = Arc::clone(self); + self.runtime + .spawn(Box::pin(async move { + registration.drive_callbacks().await; + })) + .detach(); + } + + async fn drive_callbacks(self: Arc) { loop { - let scope_changed = self.scope_changed.listen(); - let callback_guard = self.callback_gate.lock().await; - if !self.has_open_scopes() { - if let Err(error) = self.handler.shutdown().await { - log::warn!("Client lifecycle shutdown failed: {error:#}"); + let callback = { + let mut queue = self.callback_queue(); + match queue.pending.pop_front() { + Some(callback) => callback, + None => { + queue.drain_scheduled = false; + return; + } + } + }; + + match callback { + LifecycleCallback::Ready { scope, done } => { + self.run_callback("on_ready", self.handler.on_ready(scope.clone())) + .await; + let _ = done.try_send(!scope.is_cancelled()); + } + LifecycleCallback::Closed(scope) => { + self.run_callback("on_closed", self.handler.on_closed(scope)) + .await; + } + LifecycleCallback::Shutdown => { + self.run_callback("shutdown", self.handler.shutdown()).await; + self.shutdown_complete.store(true, Ordering::Release); + self.shutdown_notifier.notify(); } - return; } - drop(callback_guard); - scope_changed.await; } } - fn has_open_scopes(&self) -> bool { - let scopes = self - .scopes + async fn run_callback( + &self, + name: &'static str, + mut callback: BoxFuture<'_, anyhow::Result<()>>, + ) { + let callback = std::future::poll_fn(|context| { + let _callback_context = CallbackContextGuard::enter(self); + callback.as_mut().poll(context) + }); + let result = std::panic::AssertUnwindSafe(rt_timeout( + &*self.runtime, + self.callback_timeout, + callback, + )) + .catch_unwind() + .await; + match result { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(error))) => log::warn!("Client lifecycle {name} failed: {error:#}"), + Ok(Err(_)) => log::warn!("Client lifecycle {name} timed out"), + Err(_) => log::warn!("Client lifecycle {name} panicked"), + } + } + + fn scopes(&self) -> std::sync::MutexGuard<'_, ScopeRegistry> { + self.scopes + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn callback_queue(&self) -> std::sync::MutexGuard<'_, CallbackQueue> { + self.callback_queue .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - scopes.active.is_some() || !scopes.retired.is_empty() + .unwrap_or_else(|poisoned| poisoned.into_inner()) } fn scope_for(&self, generation: u64) -> Option { - let scopes = self - .scopes - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let scopes = self.scopes(); scopes .active .as_ref() @@ -409,6 +583,85 @@ mod tests { events: std::sync::Mutex>, } + #[derive(Default)] + struct ReentrantDisconnectLifecycle { + client: std::sync::Mutex>>, + events: std::sync::Mutex>, + } + + impl ClientLifecycle for ReentrantDisconnectLifecycle { + fn install<'a>(&'a self, client: Weak) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + *self + .client + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(client); + Ok(()) + }) + } + + fn on_ready<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("ready-started"); + let client = self + .client + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .as_ref() + .and_then(Weak::upgrade) + .expect("installed client"); + client.disconnect().await; + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("ready-finished"); + Ok(()) + }) + } + + fn on_closed<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("closed"); + Ok(()) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("shutdown"); + Ok(()) + }) + } + } + + struct BlockingShutdownLifecycle { + started: async_channel::Sender<()>, + release: async_channel::Receiver<()>, + calls: AtomicUsize, + completed: AtomicBool, + } + + impl ClientLifecycle for BlockingShutdownLifecycle { + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + self.calls.fetch_add(1, Ordering::SeqCst); + let _ = self.started.try_send(()); + let _ = self.release.recv().await; + self.completed.store(true, Ordering::Release); + Ok(()) + }) + } + } + struct LogoutOrderHandler { lifecycle: Arc, } @@ -500,7 +753,7 @@ mod tests { .connection_generation .store(GENERATION, Ordering::SeqCst); let registration = client.lifecycle.as_ref().expect("lifecycle registration"); - registration.begin_scope(GENERATION); + assert!(registration.begin_scope_if_current(GENERATION, || true)); client.dispatch_connected().await; let scope = lifecycle @@ -538,7 +791,6 @@ mod tests { .expect("cleanup did not panic"); assert_eq!(scope.state(), ConnectionScopeState::Closed); - assert_eq!(lifecycle.events(), vec!["install", "ready:9", "closed:9"]); client.shutdown_lifecycle().await; client.shutdown_lifecycle().await; assert_eq!(lifecycle.shutdowns.load(Ordering::SeqCst), 1); @@ -578,11 +830,13 @@ mod tests { client .connection_generation .store(GENERATION, Ordering::SeqCst); - client - .lifecycle - .as_ref() - .expect("lifecycle registration") - .begin_scope(GENERATION); + assert!( + client + .lifecycle + .as_ref() + .expect("lifecycle registration") + .begin_scope_if_current(GENERATION, || true) + ); let ready_client = Arc::clone(&client); let ready_task = tokio::spawn(async move { @@ -610,12 +864,15 @@ mod tests { }) .await .expect("scope cancellation"); - assert!(!cleanup_task.is_finished()); + tokio::time::timeout(std::time::Duration::from_secs(2), cleanup_task) + .await + .expect("cleanup completed while callback was blocked") + .expect("cleanup task did not panic"); + assert_eq!(scope.state(), ConnectionScopeState::Closed); release_ready_tx.send(()).await.expect("release ready hook"); ready_task.await.expect("ready task did not panic"); - cleanup_task.await.expect("cleanup task did not panic"); - assert_eq!(scope.state(), ConnectionScopeState::Closed); + client.shutdown_lifecycle().await; assert_eq!( *lifecycle .events @@ -626,26 +883,185 @@ mod tests { client.signal_shutdown_sync(); } + #[tokio::test] + async fn ready_callback_can_disconnect_its_client() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let lifecycle = Arc::new(ReentrantDisconnectLifecycle::default()); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await + .expect("client build") + .client(); + const GENERATION: u64 = 17; + client + .connection_generation + .store(GENERATION, Ordering::SeqCst); + assert!( + client + .lifecycle + .as_ref() + .expect("lifecycle registration") + .begin_scope_if_current(GENERATION, || true) + ); + + tokio::time::timeout( + std::time::Duration::from_secs(2), + client.dispatch_connected(), + ) + .await + .expect("reentrant disconnect completed"); + client.shutdown_lifecycle().await; + + assert_eq!( + *lifecycle + .events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["ready-started", "ready-finished", "closed", "shutdown"] + ); + assert!(!client.is_logged_in()); + assert!(!client.is_ready.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn callback_timeout_does_not_hold_connection_cleanup() { + let (ready_started_tx, ready_started_rx) = async_channel::bounded(1); + let (_release_ready_tx, release_ready_rx) = async_channel::bounded(1); + let lifecycle = Arc::new(BlockingReadyLifecycle { + ready_started: ready_started_tx, + release_ready: release_ready_rx, + scope: std::sync::Mutex::new(None), + events: std::sync::Mutex::new(Vec::new()), + }); + let registration = Arc::new(LifecycleRegistration::new_with_timeout( + lifecycle.clone(), + Arc::new(TokioRuntime), + std::time::Duration::from_millis(20), + )); + const GENERATION: u64 = 19; + assert!(registration.begin_scope_if_current(GENERATION, || true)); + + let ready_registration = Arc::clone(®istration); + let ready = tokio::spawn(async move { ready_registration.ready(GENERATION).await }); + ready_started_rx + .recv() + .await + .expect("ready callback started"); + let scope = lifecycle + .scope + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + .expect("ready scope"); + + registration.cancel_scope(GENERATION); + registration.close_scope(GENERATION); + assert_eq!(scope.state(), ConnectionScopeState::Closed); + assert!( + !tokio::time::timeout(std::time::Duration::from_secs(1), ready) + .await + .expect("ready callback was bounded") + .expect("ready task did not panic") + ); + tokio::time::timeout(std::time::Duration::from_secs(1), registration.shutdown()) + .await + .expect("lifecycle shutdown completed"); + + assert_eq!( + *lifecycle + .events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["ready-started", "closed"] + ); + } + + #[tokio::test] + async fn cancelled_shutdown_waiter_does_not_cancel_shutdown() { + let (started_tx, started_rx) = async_channel::bounded(1); + let (release_tx, release_rx) = async_channel::bounded(1); + let lifecycle = Arc::new(BlockingShutdownLifecycle { + started: started_tx, + release: release_rx, + calls: AtomicUsize::new(0), + completed: AtomicBool::new(false), + }); + let registration = Arc::new(LifecycleRegistration::new( + lifecycle.clone(), + Arc::new(TokioRuntime), + )); + + let first_registration = Arc::clone(®istration); + let first = tokio::spawn(async move { first_registration.shutdown().await }); + started_rx.recv().await.expect("shutdown callback started"); + first.abort(); + let _ = first.await; + + release_tx + .send(()) + .await + .expect("release shutdown callback"); + tokio::time::timeout(std::time::Duration::from_secs(2), registration.shutdown()) + .await + .expect("later shutdown waiter observed completion"); + + assert!(lifecycle.completed.load(Ordering::Acquire)); + assert_eq!(lifecycle.calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn stale_generation_is_rejected_before_scope_publication() { + let lifecycle = Arc::new(RecordingLifecycle::default()); + let registration = Arc::new(LifecycleRegistration::new( + lifecycle.clone(), + Arc::new(TokioRuntime), + )); + let generation = portable_atomic::AtomicU64::new(24); + generation.store(25, Ordering::SeqCst); + + assert!( + !registration + .begin_scope_if_current(24, || { generation.load(Ordering::SeqCst) == 24 }) + ); + assert!(registration.scope_for(24).is_none()); + tokio::time::timeout(std::time::Duration::from_secs(2), registration.shutdown()) + .await + .expect("shutdown did not wait for a rejected scope"); + assert_eq!(lifecycle.shutdowns.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn replaced_scope_stays_closeable_by_its_generation() { let lifecycle = Arc::new(RecordingLifecycle::default()); - let registration = LifecycleRegistration::new(lifecycle.clone()); + let registration = Arc::new(LifecycleRegistration::new( + lifecycle.clone(), + Arc::new(TokioRuntime), + )); - registration.begin_scope(31); + assert!(registration.begin_scope_if_current(31, || true)); assert!(registration.ready(31).await); let first = registration.scope_for(31).expect("first scope"); - registration.begin_scope(32); + assert!(registration.begin_scope_if_current(32, || true)); assert_eq!(first.state(), ConnectionScopeState::Cancelled); assert!(registration.scope_for(31).is_some()); assert!(registration.ready(32).await); - registration.close_scope(31).await; + registration.close_scope(31); assert_eq!(first.state(), ConnectionScopeState::Closed); assert!(registration.scope_for(31).is_none()); assert!(registration.scope_for(32).is_some()); - registration.close_scope(32).await; + registration.close_scope(32); registration.shutdown().await; assert_eq!( lifecycle.events(), @@ -656,9 +1072,12 @@ mod tests { #[tokio::test] async fn shutdown_waits_for_the_active_scope_and_is_terminal() { let lifecycle = Arc::new(RecordingLifecycle::default()); - let registration = Arc::new(LifecycleRegistration::new(lifecycle.clone())); + let registration = Arc::new(LifecycleRegistration::new( + lifecycle.clone(), + Arc::new(TokioRuntime), + )); const GENERATION: u64 = 21; - registration.begin_scope(GENERATION); + assert!(registration.begin_scope_if_current(GENERATION, || true)); let scope = registration .scope_for(GENERATION) .expect("active connection scope"); @@ -676,14 +1095,14 @@ mod tests { .expect("scope cancellation"); assert!(!shutdown.is_finished()); - registration.close_scope(GENERATION).await; + registration.close_scope(GENERATION); shutdown.await.expect("shutdown did not panic"); assert_eq!( lifecycle.events(), vec!["closed:21".to_string(), "shutdown".to_string()] ); - registration.begin_scope(GENERATION + 1); + assert!(!registration.begin_scope_if_current(GENERATION + 1, || true)); assert!(registration.scope_for(GENERATION + 1).is_none()); assert!(!registration.ready(GENERATION + 1).await); registration.shutdown().await; diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index a00cec185..61b5f8e42 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -981,7 +981,7 @@ impl Client { proc.clear_key_cache().await; } if let Some(lifecycle) = &self.lifecycle { - lifecycle.close_scope(closed_generation).await; + lifecycle.close_scope(closed_generation); } } diff --git a/src/client/node_io.rs b/src/client/node_io.rs index cc1e6cb8f..11d6ab7fa 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -698,7 +698,14 @@ impl Client { // from previous connections (e.g., during 515 reconnect cycles). let current_generation = self.connection_generation.fetch_add(1, Ordering::SeqCst) + 1; if let Some(lifecycle) = &self.lifecycle { - lifecycle.begin_scope(current_generation); + let opened = lifecycle.begin_scope_if_current(current_generation, || { + self.connection_generation.load(Ordering::SeqCst) == current_generation + && !self.expected_disconnect.load(Ordering::Acquire) + }); + if !opened { + debug!("Ignoring stanza retired during lifecycle publication"); + return; + } } info!( diff --git a/src/lib.rs b/src/lib.rs index de63783c1..fa2799a89 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -179,7 +179,10 @@ pub mod version; /// `use whatsapp_rust::prelude::*;`. pub mod prelude { pub use crate::bot::{Bot, BotBuilder, BotHandle, EventDelivery, MessageContext}; - pub use crate::client::{Client, ClientError}; + pub use crate::client::{ + Client, ClientBuilder, ClientBuilderError, ClientError, ClientLifecycle, ConnectionScope, + ConnectionScopeState, + }; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] pub use crate::runtime_impl::TokioRuntime; From b42a2dfceb193c1afeafecaa57b637622776f5f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 03:20:30 -0300 Subject: [PATCH 06/46] feat(events): add removable filtered subscriptions --- examples/voip-cli/src/main.rs | 2 +- src/bot.rs | 7 +- src/client/accessors.rs | 18 +- src/client/extension_lifecycle.rs | 8 +- src/client/tests.rs | 6 +- src/features/labels.rs | 2 +- src/handlers/call.rs | 30 +-- src/handlers/notification/mod.rs | 40 ++-- src/history_sync.rs | 2 +- src/lib.rs | 5 +- src/message/commit_batch.rs | 10 +- src/message/durability.rs | 2 +- src/message/tests.rs | 86 ++++---- src/pair_code.rs | 6 +- src/passkey/flow.rs | 16 +- src/receipt.rs | 2 +- storages/chat-store/src/lib.rs | 2 +- storages/chat-store/src/store.rs | 2 +- tests/e2e/src/lib.rs | 4 +- wacore/src/types/events.rs | 353 +++++++++++++++++++++++++----- 20 files changed, 434 insertions(+), 169 deletions(-) diff --git a/examples/voip-cli/src/main.rs b/examples/voip-cli/src/main.rs index b35eba802..01044f75a 100644 --- a/examples/voip-cli/src/main.rs +++ b/examples/voip-cli/src/main.rs @@ -1660,7 +1660,7 @@ async fn run_bot(mode: Mode) -> Result<()> { // accept flow, so the raw-node-forwarding crutch the old hand-rolled inbound path needed is gone. let manages_media = accept || target.is_some(); let observer = Arc::new(CallObserver::new(client.clone(), accept, video, audio)); - client.register_handler(observer.clone()); + let _observer_subscription = client.subscribe_handler(observer.clone()); if let Some(peer) = target { let client2 = client.clone(); diff --git a/src/bot.rs b/src/bot.rs index 5c8e9ffdb..e3436421c 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -590,14 +590,15 @@ impl Bot { client .core .event_bus - .add_handler(Arc::new(CallbackBusAdapter::new( + .subscribe_handler(Arc::new(CallbackBusAdapter::new( client.clone(), event_handlers, event_delivery, - ))); + ))) + .detach(); } for handler in raw_handlers { - client.core.event_bus.add_handler(handler); + client.core.event_bus.subscribe_handler(handler).detach(); } // If pair code options are set, spawn a task to request pair code after socket is ready diff --git a/src/client/accessors.rs b/src/client/accessors.rs index 46426ac13..c99e81e29 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -27,9 +27,21 @@ impl Client { cache } - /// Registers an external event handler to the core event bus. - pub fn register_handler(&self, handler: Arc) { - self.core.event_bus.add_handler(handler); + /// Subscribe an external event handler with an explicit event filter. + pub fn subscribe( + &self, + interest: wacore::types::events::EventInterest, + handler: Arc, + ) -> wacore::types::events::Subscription { + self.core.event_bus.subscribe(interest, handler) + } + + /// Subscribe using the handler's current registration-time interest hint. + pub fn subscribe_handler( + &self, + handler: Arc, + ) -> wacore::types::events::Subscription { + self.core.event_bus.subscribe_handler(handler) } /// Enable or disable raw node forwarding. diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index a2c09da7a..e7193c68d 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -1127,9 +1127,11 @@ mod tests { .await .expect("client build") .client(); - client.register_handler(Arc::new(LogoutOrderHandler { - lifecycle: lifecycle.clone(), - })); + client + .subscribe_handler(Arc::new(LogoutOrderHandler { + lifecycle: lifecycle.clone(), + })) + .detach(); client.logout().await.expect("client logout"); diff --git a/src/client/tests.rs b/src/client/tests.rs index 2098a64cc..42eb8b504 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -293,7 +293,9 @@ async fn test_ack_dispatches_server_ack_event() { let client = crate::test_utils::create_test_client().await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone() as Arc); + client + .subscribe_handler(collector.clone() as Arc) + .detach(); // Plain message ack (no waiter registered): event fires with the ack's // class, from and server timestamp; error is None. @@ -1597,7 +1599,7 @@ async fn connect_failure_403_dispatches_account_locked_logout() { use wacore::types::events::ChannelEventHandler; let client = create_offline_sync_test_client().await; let (handler, events) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); // location="rva" is a region routing token and must not change the verdict. let failure = NodeBuilder::new("failure") diff --git a/src/features/labels.rs b/src/features/labels.rs index 7d175c355..684158ee5 100644 --- a/src/features/labels.rs +++ b/src/features/labels.rs @@ -252,7 +252,7 @@ mod tests { fn run(m: &Mutation) -> (bool, Vec>) { let bus = CoreEventBus::new(); let rec = Arc::new(Recorder::default()); - bus.add_handler(rec.clone()); + bus.subscribe_handler(rec.clone()).detach(); let handled = dispatch_label_mutation(&bus, m, false); let events = rec.events.lock().unwrap().clone(); (handled, events) diff --git a/src/handlers/call.rs b/src/handlers/call.rs index f295e5269..780b34550 100644 --- a/src/handlers/call.rs +++ b/src/handlers/call.rs @@ -653,7 +653,7 @@ mod tests { let (client, sends) = make_sending_client_with_failure_after(None).await; let event_rx = register_native_opus_call(&client, Vec::new()); let (handler, global_rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let mut cancelled = false; assert!( @@ -829,7 +829,7 @@ mod tests { let (client, sends) = make_sending_client_with_failure_after(Some(1)).await; let (global_handler, global_rx) = ChannelEventHandler::new(); - client.register_handler(global_handler); + client.subscribe_handler(global_handler).detach(); let registry = client.call_registry(); let generation = registry.insert(wacore::voip::CallSession::new_outgoing( "CALL-ID-0001", @@ -875,7 +875,7 @@ mod tests { let (client, send_started, release_send) = make_blocking_sending_client().await; let (global_handler, global_rx) = ChannelEventHandler::new(); - client.register_handler(global_handler); + client.subscribe_handler(global_handler).detach(); let registry = client.call_registry(); let stale_generation = registry.insert(wacore::voip::CallSession::new_incoming( "CALL-ID-0001", @@ -966,7 +966,7 @@ mod tests { let client = make_sending_client().await; let (global_handler, global_rx) = ChannelEventHandler::new(); - client.register_handler(global_handler); + client.subscribe_handler(global_handler).detach(); let registry = client.call_registry(); let session = wacore::voip::CallSession::new_incoming( "CALL-ID-0001", @@ -1128,7 +1128,7 @@ mod tests { let client = make_client().await; let (global_handler, global_rx) = ChannelEventHandler::new(); - client.register_handler(global_handler); + client.subscribe_handler(global_handler).detach(); let registry = client.call_registry(); let generation = registry.insert(wacore::voip::CallSession::new_incoming( "CALL-ID-0001", @@ -1201,7 +1201,7 @@ mod tests { async fn offer_dispatches_event() { let client = make_client().await; let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let node = node_to_owned_ref(&offer_stanza()); let mut cancelled = false; @@ -1222,7 +1222,7 @@ mod tests { async fn unrecognized_action_does_not_dispatch() { let client = make_client().await; let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let node = node_to_owned_ref( &NodeBuilder::new("call") @@ -1296,7 +1296,7 @@ mod tests { async fn malformed_stanza_does_not_error_or_dispatch() { let client = make_client().await; let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let node = node_to_owned_ref( &NodeBuilder::new("call") @@ -1479,7 +1479,7 @@ mod tests { async fn unanswered_incoming_terminate_surfaces_missed_call() { let client = make_client().await; let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let mut cancelled = false; // The offer rings (marks the call ringing). @@ -1517,7 +1517,7 @@ mod tests { async fn duplicate_terminate_does_not_refire_missed_call() { let client = make_client().await; let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let mut cancelled = false; assert!( @@ -1556,7 +1556,7 @@ mod tests { async fn outgoing_call_terminate_does_not_surface_missed_call() { let client = make_client().await; let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let peer = Jid::new("222222222222222", Server::Lid); let creator = Jid::new("111111111111111", Server::Lid); // us, the caller @@ -1602,7 +1602,7 @@ mod tests { ] { let client = make_client().await; let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let mut cancelled = false; assert!( @@ -1660,7 +1660,7 @@ mod tests { async fn timeout_terminate_surfaces_missed_call() { let client = make_client().await; let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let mut cancelled = false; assert!( @@ -1725,7 +1725,7 @@ mod tests { *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let mut cancelled = false; // The offer rings. @@ -1776,7 +1776,7 @@ mod tests { async fn answered_call_then_caller_terminate_is_not_missed() { let client = make_client().await; let (handler, rx) = ChannelEventHandler::new(); - client.register_handler(handler); + client.subscribe_handler(handler).detach(); let mut cancelled = false; // The offer rings. diff --git a/src/handlers/notification/mod.rs b/src/handlers/notification/mod.rs index b279948b0..92f924051 100644 --- a/src/handlers/notification/mod.rs +++ b/src/handlers/notification/mod.rs @@ -469,7 +469,7 @@ mod tests { async fn test_contacts_update_dispatches_contact_updated_event() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let node = NodeBuilder::new("notification") .attr("type", "contacts") @@ -500,7 +500,7 @@ mod tests { // Creates two mappings: old_lid→old_pn AND new_lid→new_pn. let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let node = NodeBuilder::new("notification") .attr("type", "contacts") @@ -554,7 +554,7 @@ mod tests { async fn test_contacts_modify_without_lid_skips_mapping() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let node = NodeBuilder::new("notification") .attr("type", "contacts") @@ -576,7 +576,7 @@ mod tests { async fn test_contacts_sync_dispatches_contact_sync_requested_event() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let node = NodeBuilder::new("notification") .attr("type", "contacts") @@ -602,7 +602,7 @@ mod tests { async fn test_contacts_add_remove_do_not_dispatch_events() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); for tag in ["add", "remove"] { let node = NodeBuilder::new("notification") @@ -624,7 +624,7 @@ mod tests { async fn test_contacts_empty_notification_ignored() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); // No child element let node = NodeBuilder::new("notification") @@ -648,7 +648,7 @@ mod tests { async fn test_contacts_modify_same_jid_still_dispatches() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let node = NodeBuilder::new("notification") .attr("type", "contacts") @@ -704,7 +704,7 @@ mod tests { async fn test_contacts_modify_missing_new_attr_drops_event() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let node = NodeBuilder::new("notification") .attr("type", "contacts") @@ -728,7 +728,7 @@ mod tests { async fn test_group_change_number_dispatches_with_new_owner_and_suggestions() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let node = NodeBuilder::new("notification") .attr("type", "w:gp2") @@ -858,7 +858,7 @@ mod tests { // We don't maintain a userhash index, so this should be a no-op. let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let node = NodeBuilder::new("notification") .attr("type", "contacts") @@ -879,7 +879,7 @@ mod tests { async fn test_identity_change_dispatches_event_and_invalidates_cache() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); // Pre-populate device registry so clear_device_record has something to clear let record = wacore::store::traits::DeviceListRecord { @@ -942,7 +942,7 @@ mod tests { async fn test_identity_change_ignores_self_primary() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); // Set our own JID so the self-check works client @@ -971,7 +971,7 @@ mod tests { async fn test_identity_change_ignores_companion_device() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let node = NodeBuilder::new("notification") .attr("type", "encrypt") @@ -991,7 +991,7 @@ mod tests { async fn test_local_identity_change_dispatches_implicit_event() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let sender: Jid = "5511777777777@s.whatsapp.net".parse().unwrap(); handle_local_identity_change(&client, sender).await; @@ -1018,7 +1018,7 @@ mod tests { async fn test_local_identity_change_skips_self() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); client .persistence_manager @@ -1040,7 +1040,7 @@ mod tests { async fn test_local_identity_change_skips_companion_device() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let sender: Jid = "5511777777777:5@s.whatsapp.net".parse().unwrap(); handle_local_identity_change(&client, sender).await; @@ -1151,7 +1151,7 @@ mod tests { use wacore::types::jid::JidExt; let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); // Prior identity present so the gate runs (the offline attr only defers the // eager session re-establishment, not the change notification). @@ -1189,7 +1189,7 @@ mod tests { use wacore::types::jid::JidExt; let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let target: Jid = "5511666666666@s.whatsapp.net".parse().unwrap(); let addr = target.to_protocol_address(); @@ -1280,7 +1280,7 @@ mod tests { use wacore::types::jid::JidExt; let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let pn = "5511555555555"; let lid = "100000000000055"; @@ -1338,7 +1338,7 @@ mod tests { use wacore::types::jid::JidExt; let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); // Cold cache: no PN->LID mapping, so resolve_encryption_jid(PN) returns PN. let pn_jid: Jid = "5511444444444@s.whatsapp.net".parse().unwrap(); diff --git a/src/history_sync.rs b/src/history_sync.rs index c472857f5..4e7969bf6 100644 --- a/src/history_sync.rs +++ b/src/history_sync.rs @@ -883,7 +883,7 @@ mod tests { // Register a handler BEFORE the task so retain_blob is true. let (handler, event_rx) = wacore::types::events::ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); + client.core.event_bus.subscribe_handler(handler).detach(); client .process_history_sync_task("HIST_LAZY_EVENT".to_string(), notification.into()) diff --git a/src/lib.rs b/src/lib.rs index fa2799a89..b2b083cda 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -191,7 +191,10 @@ pub mod prelude { pub use crate::shutdown::shutdown_signal; #[cfg(feature = "sqlite-storage")] pub use crate::store::SqliteStore; - pub use crate::types::events::{BatchOrigin, Event, EventKind, InboundMessage, MessageBatch}; + pub use crate::types::events::{ + BatchOrigin, ChannelEventHandler, Event, EventHandler, EventInterest, EventKind, + InboundMessage, MessageBatch, Subscription, + }; pub use crate::types::message::MessageInfo; pub use crate::{Jid, Server}; pub use wacore::proto_helpers::{MessageBuilderExt, MessageExt}; diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 566172ccc..5aaca3d06 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -915,7 +915,7 @@ mod tests { }); let _ = client.inbound_durability_hook.set(hook.clone()); let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); + client.core.event_bus.subscribe_handler(handler).detach(); client.inbound_commit_batch.reset(); for id in ["B1", "B2", "B3"] { @@ -962,7 +962,7 @@ mod tests { }); let _ = client.inbound_durability_hook.set(hook.clone()); let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); + client.core.event_bus.subscribe_handler(handler).detach(); client.commit_or_batch_inbound(item("L1"), false).await; assert_eq!( @@ -1008,7 +1008,7 @@ mod tests { let client = create_test_client_with_failing_http("batch_no_hook").await; client.inbound_commit_batch.reset(); let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); + client.core.event_bus.subscribe_handler(handler).detach(); client.commit_or_batch_inbound(item("N1"), false).await; client.commit_or_batch_inbound(item("N2"), false).await; @@ -1038,7 +1038,7 @@ mod tests { }); let _ = client.inbound_durability_hook.set(hook.clone()); let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); + client.core.event_bus.subscribe_handler(handler).detach(); client.commit_or_batch_inbound(item("T1"), false).await; client.commit_or_batch_inbound(item("T2"), false).await; @@ -1077,7 +1077,7 @@ mod tests { }); let _ = client.inbound_durability_hook.set(hook.clone()); let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); + client.core.event_bus.subscribe_handler(handler).detach(); client.commit_or_batch_inbound(item("C1"), false).await; client.inbound_commit_batch.reset(); diff --git a/src/message/durability.rs b/src/message/durability.rs index 949353f64..c4307eaef 100644 --- a/src/message/durability.rs +++ b/src/message/durability.rs @@ -234,7 +234,7 @@ mod tests { // Redelivery once the commit succeeds clears the buffer AND finally // dispatches the event (the original batch never did). let (handler, rx) = wacore::types::events::ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); + client.core.event_bus.subscribe_handler(handler).detach(); hook.succeed.store(true, Ordering::SeqCst); client.ack_or_replay_to_hook(&info).await; assert_eq!(hook.calls.load(Ordering::SeqCst), 3); diff --git a/src/message/tests.rs b/src/message/tests.rs index efc8f71b0..1d31f7c7d 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -378,7 +378,7 @@ async fn batch_accumulates_undecryptable_and_dispatches_once() { .await; let recorder = Arc::new(EventRecorder::default()); - client.register_handler(recorder.clone()); + client.subscribe_handler(recorder.clone()).detach(); let sender_jid: Jid = "1234567890@s.whatsapp.net" .parse() @@ -5010,7 +5010,7 @@ fn build_unavailable_stanza(sender: &str, msg_id: &str, with_enc: bool) -> Arc bool { let (client, _transport) = capturing_client(test_id).await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "5511777776666@s.whatsapp.net"; let parent_id = "WINDOW_PARENT"; @@ -8689,7 +8689,7 @@ async fn secret_encrypted_edit_decrypts_via_resolver_when_store_empty() { seed_test_pn(&client).await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); assert!( client @@ -8753,7 +8753,7 @@ async fn secret_encrypted_edit_decrypts_via_resolver_when_store_empty() { async fn decrypted_message_edit_recaptures_secret_for_next_edit() { let (client, _transport) = capturing_client("secret_edit_chain").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "5511777776666@s.whatsapp.net"; let parent_id = "PARENT_EDIT"; @@ -8846,7 +8846,7 @@ async fn secret_encrypted_message_edit_uses_lid_pn_fallback_in_group() { let (client, _transport) = capturing_client("secret_edit_alt_group").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "120363021033254949@g.us"; let parent_id = "GROUP_PARENT_EDIT"; @@ -8926,7 +8926,7 @@ async fn decrypted_message_edit_refreshes_alternate_secret_alias() { let (client, _transport) = capturing_client("secret_edit_alt_refresh").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "120363021033254949@g.us"; let parent_id = "GROUP_PARENT_EDIT"; @@ -9054,7 +9054,7 @@ async fn msmsg_decrypts_when_secret_is_stored() { use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; let (client, _transport) = capturing_client("msmsg_ok").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "867051314767696@bot"; let our_pn = "5511000000001@s.whatsapp.net"; @@ -9130,7 +9130,7 @@ async fn msmsg_decrypts_when_secret_is_stored() { async fn msmsg_without_stored_secret_nacks_495() { let (client, transport) = capturing_client("msmsg_nosecret").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let bot_reply_id = "BOT_REPLY_NS"; let outbound_id = "OUTBOUND_NS"; @@ -9246,7 +9246,7 @@ async fn msmsg_bot_edit_uses_edit_target_id_for_hkdf() { use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; let (client, _transport) = capturing_client("msmsg_edit").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "867051314767696@bot"; let our_pn = "5511000000001@s.whatsapp.net"; @@ -9395,7 +9395,7 @@ async fn msmsg_bot_edit_first_keeps_info_id() { use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; let (client, _transport) = capturing_client("msmsg_first").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "867051314767696@bot"; let our_pn = "5511000000001@s.whatsapp.net"; @@ -9467,7 +9467,7 @@ async fn msmsg_falls_back_to_info_id_when_primary_uses_edit_target() { use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; let (client, _transport) = capturing_client("msmsg_fb_to_info").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "867051314767696@bot"; let our_pn = "5511000000001@s.whatsapp.net"; @@ -9553,7 +9553,7 @@ async fn msmsg_falls_back_to_edit_target_when_primary_uses_info_id() { use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; let (client, _transport) = capturing_client("msmsg_fb_to_edit").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "867051314767696@bot"; let our_pn = "5511000000001@s.whatsapp.net"; @@ -10233,7 +10233,7 @@ async fn mixed_msmsg_and_unknown_enc_still_decrypts_msmsg() { ))) .await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "867051314767696@bot"; let our_lid = "999888777666555@lid"; @@ -10320,7 +10320,7 @@ async fn msmsg_alternate_lookup_resolves_lid_to_stored_pn() { let (client, _transport) = capturing_client("msmsg_alt_lookup").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "867051314767696@bot"; let our_lid_user = "999888777666555"; @@ -10425,7 +10425,7 @@ async fn fanout_capture_lets_subsequent_msmsg_decrypt() { ))) .await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let bot_chat: Jid = "867051314767696@bot".parse().unwrap(); let our_lid_str = "999888777666555@lid"; @@ -10547,7 +10547,7 @@ async fn msmsg_outbound_put_and_inbound_get_match_for_lid_bot() { ))) .await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let bot_chat: Jid = "867051314767696@bot".parse().unwrap(); let outbound_id = "OUT_LID"; @@ -10635,7 +10635,7 @@ async fn msmsg_with_bot_device_suffix_round_trips() { use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; let (client, _transport) = capturing_client("msmsg_bot_device").await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let chat = "867051314767696@bot"; let our_pn = "5511000000001@s.whatsapp.net"; @@ -10833,7 +10833,7 @@ async fn enc_comment_inbound_dispatches_body_with_parent_link() { let (client, _transport) = capturing_client("enc_comment_inbound").await; ensure_bob_paired(&client).await; let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); + client.core.event_bus.subscribe_handler(handler).detach(); let group: Jid = "120363400000000002@g.us".parse().expect("group"); let author: Jid = "5511888887777@s.whatsapp.net".parse().expect("author"); diff --git a/src/pair_code.rs b/src/pair_code.rs index 84065e790..b2c502e5f 100644 --- a/src/pair_code.rs +++ b/src/pair_code.rs @@ -842,7 +842,7 @@ mod tests { async fn refresh_code_matching_ref_dispatches_event() { let client = create_test_client().await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let pairing_ref = vec![5, 6, 7, 8]; set_waiting(&client, pairing_ref.clone(), wacore::time::now_secs(), 0).await; @@ -867,7 +867,7 @@ mod tests { async fn refresh_code_without_force_manual_defaults_false() { let client = create_test_client().await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); let pairing_ref = vec![5, 6, 7, 8]; set_waiting(&client, pairing_ref.clone(), wacore::time::now_secs(), 0).await; @@ -891,7 +891,7 @@ mod tests { async fn refresh_code_mismatched_ref_is_ignored() { let client = create_test_client().await; let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); set_waiting(&client, vec![5, 6, 7, 8], wacore::time::now_secs(), 0).await; diff --git a/src/passkey/flow.rs b/src/passkey/flow.rs index b1f18ca38..ac37287db 100644 --- a/src/passkey/flow.rs +++ b/src/passkey/flow.rs @@ -889,7 +889,9 @@ mod tests { async fn passkey_prologue_request_emits_event_without_committing_rotation() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone() as Arc); + client + .subscribe_handler(collector.clone() as Arc) + .detach(); let before = client .persistence_manager @@ -926,7 +928,9 @@ mod tests { async fn passkey_prologue_request_from_non_server_is_ignored() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone() as Arc); + client + .subscribe_handler(collector.clone() as Arc) + .detach(); let child = NodeBuilder::new(TAG_PASSKEY_REQUEST_OPTIONS) .bytes(b"{}".to_vec()) @@ -951,7 +955,9 @@ mod tests { async fn passkey_prologue_request_without_inline_options_falls_back_to_fetch() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone() as Arc); + client + .subscribe_handler(collector.clone() as Arc) + .detach(); // No inline options: the handler falls back to an IQ fetch. The test client // isn't connected, so the fetch fails and surfaces a non-continuation error. @@ -973,7 +979,9 @@ mod tests { async fn passkey_continuation_without_session_emits_error() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone() as Arc); + client + .subscribe_handler(collector.clone() as Arc) + .detach(); let primary = wa::PrimaryEphemeralIdentity { public_key: Some(vec![0xAB; 32]), diff --git a/src/receipt.rs b/src/receipt.rs index 711a885e2..afb5d45f6 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -1708,7 +1708,7 @@ mod tests { .await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + client.subscribe_handler(collector.clone()).detach(); (client, collector) } diff --git a/storages/chat-store/src/lib.rs b/storages/chat-store/src/lib.rs index 383c30ffc..d78f9ac1a 100644 --- a/storages/chat-store/src/lib.rs +++ b/storages/chat-store/src/lib.rs @@ -21,7 +21,7 @@ //! //! ```ignore //! let chat_store = ChatStore::new(&sqlite_store).await?; -//! client.register_handler(chat_store.handler()); +//! let _chat_subscription = client.subscribe_handler(chat_store.handler()); //! //! let chats = chat_store.chats(false, 50).await?; //! let page = chat_store.messages(&chats[0].jid, None, 40).await?; diff --git a/storages/chat-store/src/store.rs b/storages/chat-store/src/store.rs index c195fc6af..c947f70ed 100644 --- a/storages/chat-store/src/store.rs +++ b/storages/chat-store/src/store.rs @@ -58,7 +58,7 @@ pub(crate) enum WriterMsg { /// Wire-up: /// ```ignore /// let chat_store = ChatStore::new(&sqlite_store).await?; -/// client.register_handler(chat_store.handler()); +/// let _chat_subscription = client.subscribe_handler(chat_store.handler()); /// let mut changes = chat_store.subscribe(); /// ``` pub struct ChatStore { diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index d075f703c..af5761580 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -213,7 +213,7 @@ impl TestClient { if push_name_pre_seeded { client.set_force_active_delivery_receipts(true); } - client.register_handler(event_handler); + client.subscribe_handler(event_handler).detach(); // The mock server no longer auto-pairs (legacy timer is off by // default). Spawn an out-of-process "phone" that POSTs the first @@ -222,7 +222,7 @@ impl TestClient { // Uses its own ChannelEventHandler because async_channel is MPMC: // sharing event_rx would steal events from wait_for_event below. let (qr_handler, qr_rx) = ChannelEventHandler::new(); - client.register_handler(qr_handler); + client.subscribe_handler(qr_handler).detach(); let _qr_responder = spawn_qr_autoresponder_http(qr_rx); let run_handle = bot.spawn(); diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index fae040c67..89ebb0aec 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -4,6 +4,7 @@ use crate::types::message::MessageInfo; use crate::types::presence::{ChatPresence, ChatPresenceMedia, ReceiptType}; use bytes::Bytes; use chrono::{DateTime, Duration, Utc}; +use portable_atomic::{AtomicU64, Ordering}; use serde::Serialize; use std::fmt; use std::sync::{Arc, OnceLock, RwLock}; @@ -204,7 +205,7 @@ impl Serialize for LazyHistorySync { /// Discriminant for each [`Event`] variant, used to express handler interest /// without materializing the event. One per `Event` variant, in declaration /// order; the value doubles as a bit index in [`EventInterest`], so there can -/// be at most 64 kinds. +/// be at most 128 kinds. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[repr(u8)] #[non_exhaustive] @@ -281,9 +282,9 @@ impl EventKind { // fails compilation instead of silently corrupting the mask at runtime. const _: () = assert!((EventKind::ServerAck as u8) < EventKind::CAPACITY); -/// A set of [`EventKind`]s a handler wants delivered. The event bus skips -/// materializing and dispatching events whose kind no handler wants, so a -/// handler that subscribes to a few kinds never pays for boxing the others. +/// A set of [`EventKind`]s a handler wants delivered. Producers can query the +/// aggregate interest before building expensive payloads, and dispatch avoids +/// allocating an `Arc` when no handler wants the kind. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct EventInterest(u128); @@ -323,14 +324,18 @@ impl EventInterest { pub const fn union(self, other: Self) -> Self { EventInterest(self.0 | other.0) } + + const fn words(self) -> (u64, u64) { + (self.0 as u64, (self.0 >> 64) as u64) + } } pub trait EventHandler: crate::sync_marker::MaybeSendSync { fn handle_event(&self, event: Arc); - /// Which event kinds this handler wants. Defaults to all kinds, so the bus - /// keeps delivering everything to handlers that don't opt into a narrower - /// set. Override to let the bus skip materializing unwanted events. + /// Registration-time interest hint used by + /// [`CoreEventBus::subscribe_handler`]. The bus captures it once; use + /// [`Subscription::update_interest`] for later changes. fn interest(&self) -> EventInterest { EventInterest::ALL } @@ -341,7 +346,7 @@ pub trait EventHandler: crate::sync_marker::MaybeSendSync { /// # Example /// ```ignore /// let (handler, rx) = ChannelEventHandler::new(); -/// client.register_handler(handler); +/// let _subscription = client.subscribe_handler(handler); /// while let Ok(event) = rx.recv().await { /// if matches!(&*event, Event::Connected(_)) { break; } /// } @@ -363,22 +368,174 @@ impl EventHandler for ChannelEventHandler { } } +#[derive(Clone)] +struct HandlerEntry { + id: u64, + interest: EventInterest, + handler: Arc, +} + /// Immutable snapshot of the registered handlers. `dispatch` clones only the /// outer `Arc` (one refcount bump, no `Vec` allocation), then drops the lock and -/// iterates the snapshot. Handler interest is re-evaluated per dispatch so a -/// handler whose `interest()` widens at runtime still receives the new kinds. +/// iterates it. Interests change only through the subscription that owns an +/// entry, so one snapshot always contains a coherent handler/filter pair. #[derive(Default)] struct HandlerSnapshot { - handlers: Vec>, + handlers: Vec, +} + +struct CoreEventBusInner { + handlers: RwLock>, + next_id: AtomicU64, + interest_low: AtomicU64, + interest_high: AtomicU64, +} + +impl Default for CoreEventBusInner { + fn default() -> Self { + Self { + handlers: RwLock::new(Arc::new(HandlerSnapshot::default())), + next_id: AtomicU64::new(1), + interest_low: AtomicU64::new(0), + interest_high: AtomicU64::new(0), + } + } +} + +impl CoreEventBusInner { + fn snapshot(&self) -> Arc { + self.handlers + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + } + + fn publish_interest(&self, interest: EventInterest) { + let (low, high) = interest.words(); + if low != 0 { + self.interest_low.fetch_or(low, Ordering::Release); + } + if high != 0 { + self.interest_high.fetch_or(high, Ordering::Release); + } + } + + fn store_aggregate(&self, snapshot: &HandlerSnapshot) { + let aggregate = snapshot + .handlers + .iter() + .fold(EventInterest::none(), |all, entry| { + all.union(entry.interest) + }); + let (low, high) = aggregate.words(); + self.interest_low.store(low, Ordering::Release); + self.interest_high.store(high, Ordering::Release); + } + + fn remove(&self, id: u64) -> bool { + let mut guard = self + .handlers + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let current = &**guard; + let Some(position) = current.handlers.iter().position(|entry| entry.id == id) else { + return false; + }; + let mut handlers = Vec::with_capacity(current.handlers.len() - 1); + handlers.extend(current.handlers[..position].iter().cloned()); + handlers.extend(current.handlers[position + 1..].iter().cloned()); + let snapshot = Arc::new(HandlerSnapshot { handlers }); + // Retire the entry before clearing bits; an early read may only be a + // harmless false positive. + *guard = Arc::clone(&snapshot); + self.store_aggregate(&snapshot); + true + } + + fn update_interest(&self, id: u64, interest: EventInterest) -> bool { + let mut guard = self + .handlers + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let current = &**guard; + let Some(position) = current.handlers.iter().position(|entry| entry.id == id) else { + return false; + }; + if current.handlers[position].interest == interest { + return true; + } + + // Publish additions first so a completed snapshot update can never be + // hidden by the lock-free producer filter. + self.publish_interest(interest); + let mut handlers = current.handlers.clone(); + handlers[position].interest = interest; + let snapshot = Arc::new(HandlerSnapshot { handlers }); + *guard = Arc::clone(&snapshot); + self.store_aggregate(&snapshot); + true + } + + fn has_handler_for(&self, kind: EventKind) -> bool { + let bit = kind as u8; + if bit < 64 { + self.interest_low.load(Ordering::Acquire) & (1u64 << bit) != 0 + } else { + self.interest_high.load(Ordering::Acquire) & (1u64 << (bit - 64)) != 0 + } + } +} + +/// Removal token for one event-handler registration. +/// +/// Dropping it removes the handler. A dispatch that already cloned the old +/// snapshot may still complete once, while later dispatches cannot see it. +#[must_use = "dropping the subscription immediately unregisters the event handler"] +pub struct Subscription { + bus: std::sync::Weak, + id: u64, + active: bool, +} + +impl Subscription { + /// Replace this registration's filter without re-registering its handler. + /// Returns `false` if the bus no longer exists or the entry was removed. + pub fn update_interest(&self, interest: EventInterest) -> bool { + self.active + && self + .bus + .upgrade() + .is_some_and(|bus| bus.update_interest(self.id, interest)) + } + + /// Remove the handler now instead of waiting for `Drop`. + pub fn unsubscribe(mut self) -> bool { + let removed = self.remove(); + self.active = false; + removed + } + + /// Keep this registration for the remaining lifetime of the event bus. + pub fn detach(mut self) { + self.active = false; + } + + fn remove(&self) -> bool { + self.bus.upgrade().is_some_and(|bus| bus.remove(self.id)) + } +} + +impl Drop for Subscription { + fn drop(&mut self) { + if self.active { + self.remove(); + } + } } #[derive(Default, Clone)] pub struct CoreEventBus { - // Copy-on-write: the snapshot is only swapped (under the lock) when a - // handler is added, which happens at startup. `dispatch` takes a cheap - // outer-Arc clone and then drops the lock, so a concurrent `add_handler` - // can never invalidate a snapshot a dispatch is iterating. - handlers: Arc>>, + inner: Arc, } impl CoreEventBus { @@ -387,22 +544,43 @@ impl CoreEventBus { } fn snapshot(&self) -> Arc { - self.handlers - .read() - .expect("RwLock should not be poisoned") - .clone() + self.inner.snapshot() } - pub fn add_handler(&self, handler: Arc) { + /// Register `handler` with an explicit, stable filter. + pub fn subscribe( + &self, + interest: EventInterest, + handler: Arc, + ) -> Subscription { let mut guard = self + .inner .handlers .write() - .expect("RwLock should not be poisoned"); + .unwrap_or_else(|poisoned| poisoned.into_inner()); let current = &**guard; + let id = self.inner.next_id.fetch_add(1, Ordering::Relaxed); let mut handlers = Vec::with_capacity(current.handlers.len() + 1); handlers.extend(current.handlers.iter().cloned()); - handlers.push(handler); + handlers.push(HandlerEntry { + id, + interest, + handler, + }); + // An early bit only takes the slow path against the previous snapshot. + self.inner.publish_interest(interest); *guard = Arc::new(HandlerSnapshot { handlers }); + Subscription { + bus: Arc::downgrade(&self.inner), + id, + active: true, + } + } + + /// Register using the handler's current [`EventHandler::interest`] hint. + pub fn subscribe_handler(&self, handler: Arc) -> Subscription { + let interest = handler.interest(); + self.subscribe(interest, handler) } /// Returns true if there are any event handlers registered. @@ -415,25 +593,19 @@ impl CoreEventBus { /// skip producing an event nobody would receive (e.g. retaining a large /// `HistorySync` blob when only message-only handlers are registered). pub fn has_handler_for(&self, kind: EventKind) -> bool { - self.snapshot() - .handlers - .iter() - .any(|h| h.interest().wants(kind)) + self.inner.has_handler_for(kind) } pub fn dispatch(&self, event: Event) { - let snapshot = self.snapshot(); - // Skip materializing the event (Arc) when no handler wants this kind. The - // interest is re-evaluated here (not read from a cached aggregate) so a - // handler whose interest() widens at runtime is never short-circuited out. let kind = event.kind(); - if !snapshot.handlers.iter().any(|h| h.interest().wants(kind)) { + if !self.has_handler_for(kind) { return; } + let snapshot = self.snapshot(); let event = Arc::new(event); - for handler in &snapshot.handlers { - if handler.interest().wants(kind) { - handler.handle_event(Arc::clone(&event)); + for entry in &snapshot.handlers { + if entry.interest.wants(kind) { + entry.handler.handle_event(Arc::clone(&event)); } } } @@ -1920,8 +2092,8 @@ mod tests { kinds: Mutex::new(Vec::new()), interest: EventInterest::ALL, }); - bus.add_handler(only_msg.clone()); - bus.add_handler(all.clone()); + let _only_msg = bus.subscribe_handler(only_msg.clone()); + let _all = bus.subscribe_handler(all.clone()); bus.dispatch(Event::Connected(Connected::builder().build())); @@ -1942,53 +2114,48 @@ mod tests { } } let bus2 = CoreEventBus::new(); - bus2.add_handler(Arc::new(Counter)); + let _counter = bus2.subscribe_handler(Arc::new(Counter)); bus2.dispatch(Event::Connected(Connected::builder().build())); assert_eq!(CALLS.load(Ordering::SeqCst), 0); } #[test] - fn dispatch_respects_dynamically_widened_interest() { - use std::sync::Mutex; + fn subscription_updates_interest_explicitly() { use std::sync::atomic::{AtomicUsize, Ordering}; - // A handler whose interest() widens after registration. dispatch must - // re-read interest each time (never a stale cached aggregate), so the - // newly-wanted kind is delivered. struct Dynamic { - interest: Mutex, hits: AtomicUsize, } impl EventHandler for Dynamic { fn handle_event(&self, _: Arc) { self.hits.fetch_add(1, Ordering::SeqCst); } - fn interest(&self) -> EventInterest { - *self.interest.lock().unwrap() - } } let bus = CoreEventBus::new(); let h = Arc::new(Dynamic { - interest: Mutex::new(EventInterest::of(&[EventKind::Messages])), hits: AtomicUsize::new(0), }); - bus.add_handler(h.clone()); + let subscription = bus.subscribe(EventInterest::of(&[EventKind::Messages]), h.clone()); // Not yet interested in Connected: dropped before materialization. bus.dispatch(Event::Connected(Connected::builder().build())); assert_eq!(h.hits.load(Ordering::SeqCst), 0); assert!(!bus.has_handler_for(EventKind::Connected)); - // Widen interest at runtime. - *h.interest.lock().unwrap() = EventInterest::ALL; + assert!(subscription.update_interest(EventInterest::ALL)); assert!(bus.has_handler_for(EventKind::Connected)); bus.dispatch(Event::Connected(Connected::builder().build())); assert_eq!( h.hits.load(Ordering::SeqCst), 1, - "a handler whose interest widened at runtime must receive the newly-wanted kind" + "the updated subscription must receive the newly-wanted kind" ); + + assert!(subscription.update_interest(EventInterest::none())); + assert!(!bus.has_handler_for(EventKind::Connected)); + bus.dispatch(Event::Connected(Connected::builder().build())); + assert_eq!(h.hits.load(Ordering::SeqCst), 1); } #[test] @@ -2007,13 +2174,15 @@ mod tests { assert!(!bus.has_handler_for(EventKind::Messages)); assert!(!bus.has_handler_for(EventKind::Receipt)); - bus.add_handler(Arc::new(Narrow(EventInterest::of(&[EventKind::Messages])))); + let _messages = + bus.subscribe_handler(Arc::new(Narrow(EventInterest::of(&[EventKind::Messages])))); assert!(bus.has_handlers()); assert!(bus.has_handler_for(EventKind::Messages)); assert!(!bus.has_handler_for(EventKind::Receipt)); // has_handler_for is true once any registered handler wants the kind. - bus.add_handler(Arc::new(Narrow(EventInterest::of(&[EventKind::Receipt])))); + let _receipt = + bus.subscribe_handler(Arc::new(Narrow(EventInterest::of(&[EventKind::Receipt])))); assert!(bus.has_handler_for(EventKind::Messages)); assert!(bus.has_handler_for(EventKind::Receipt)); assert!(!bus.has_handler_for(EventKind::Connected)); @@ -2035,11 +2204,12 @@ mod tests { let bus = CoreEventBus::new(); let log = Arc::new(Mutex::new(Vec::new())); + let mut subscriptions = Vec::new(); for id in 0..5u32 { - bus.add_handler(Arc::new(Tagged { + subscriptions.push(bus.subscribe_handler(Arc::new(Tagged { id, log: log.clone(), - })); + }))); } bus.dispatch(Event::Connected(Connected::builder().build())); // Copy-on-write rebuilds must keep registration order intact. @@ -2073,14 +2243,15 @@ mod tests { } } self.bus - .add_handler(Arc::new(Late(self.invocations.clone()))); + .subscribe_handler(Arc::new(Late(self.invocations.clone()))) + .detach(); } } } let bus = CoreEventBus::new(); let invocations = Arc::new(AtomicUsize::new(0)); - bus.add_handler(Arc::new(AddsDuringDispatch { + let _registration = bus.subscribe_handler(Arc::new(AddsDuringDispatch { bus: bus.clone(), invocations: invocations.clone(), added: Mutex::new(false), @@ -2096,4 +2267,70 @@ mod tests { bus.dispatch(Event::Connected(Connected::builder().build())); assert_eq!(invocations.load(Ordering::SeqCst), 3); } + + #[test] + fn dropping_subscription_unregisters_handler() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct Counter(Arc); + impl EventHandler for Counter { + fn handle_event(&self, _: Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + let bus = CoreEventBus::new(); + let calls = Arc::new(AtomicUsize::new(0)); + let subscription = bus.subscribe_handler(Arc::new(Counter(Arc::clone(&calls)))); + bus.dispatch(Event::Connected(Connected::builder().build())); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + drop(subscription); + assert!(!bus.has_handlers()); + assert!(!bus.has_handler_for(EventKind::Connected)); + bus.dispatch(Event::Connected(Connected::builder().build())); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[test] + fn in_flight_dispatch_can_finish_after_unsubscribe() { + use std::sync::Barrier; + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct Blocking { + started: Arc, + release: Arc, + calls: Arc, + } + impl EventHandler for Blocking { + fn handle_event(&self, _: Arc) { + self.calls.fetch_add(1, Ordering::SeqCst); + self.started.wait(); + self.release.wait(); + } + } + + let bus = CoreEventBus::new(); + let calls = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Barrier::new(2)); + let release = Arc::new(Barrier::new(2)); + let subscription = bus.subscribe_handler(Arc::new(Blocking { + started: Arc::clone(&started), + release: Arc::clone(&release), + calls: Arc::clone(&calls), + })); + let dispatch_bus = bus.clone(); + let dispatch = std::thread::spawn(move || { + dispatch_bus.dispatch(Event::Connected(Connected::builder().build())); + }); + + started.wait(); + drop(subscription); + release.wait(); + dispatch.join().expect("dispatch thread"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + bus.dispatch(Event::Connected(Connected::builder().build())); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } } From c2ead6c4348a32387470a7af83e2fe4af1ad75ba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 03:21:42 -0300 Subject: [PATCH 07/46] docs(client): clarify lifecycle termination semantics --- src/client/builder.rs | 3 +-- src/client/extension_lifecycle.rs | 3 +++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/client/builder.rs b/src/client/builder.rs index 3fc8e56de..36599b7cf 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -352,8 +352,7 @@ impl ClientBuilder { let lifecycle = self .lifecycle - .map(|handler| LifecycleRegistration::new(handler, Arc::clone(&runtime))) - .map(Arc::new); + .map(|handler| Arc::new(LifecycleRegistration::new(handler, Arc::clone(&runtime)))); let assembly = Client::assemble( Arc::clone(&runtime), Arc::clone(&persistence_manager), diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index e7193c68d..3e9ff5e80 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -70,10 +70,13 @@ impl ConnectionScope { } } + /// Fires when this scope stops owning connection work, whether it is + /// cancelled during retirement or reaches final closure after cleanup. pub fn cancellation_signal(&self) -> ShutdownSignal { self.inner.cancellation.subscribe() } + /// Returns `true` after either cancellation or final closure. pub fn is_cancelled(&self) -> bool { self.inner.state.load(Ordering::Acquire) >= SCOPE_CANCELLED } From bc16830b43e3d69d21d43ce97db1e17be09b71b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 08:55:26 -0300 Subject: [PATCH 08/46] fix(client): reference-count raw node forwarding --- src/client.rs | 24 +++++++++++++++++++--- src/client/accessors.rs | 42 ++++++++++++++++++++++++++++++++------ src/client/lifecycle.rs | 2 +- src/client/node_io.rs | 8 ++++---- src/lib.rs | 4 ++-- wacore/src/types/events.rs | 2 +- 6 files changed, 65 insertions(+), 17 deletions(-) diff --git a/src/client.rs b/src/client.rs index e7d544ad7..d82c1d585 100644 --- a/src/client.rs +++ b/src/client.rs @@ -57,6 +57,25 @@ use portable_atomic::{AtomicI64, AtomicU64}; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; +/// Lease that keeps raw decoded stanza events enabled for one consumer. +/// +/// Dropping the final lease disables forwarding. The lease holds only a weak +/// client reference, so it cannot keep the client alive. +#[must_use = "dropping the lease immediately releases raw-node forwarding"] +pub struct RawNodeLease { + client: std::sync::Weak, +} + +impl Drop for RawNodeLease { + fn drop(&mut self) { + let Some(client) = self.client.upgrade() else { + return; + }; + let previous = client.raw_node_forwarding.fetch_sub(1, Ordering::Relaxed); + debug_assert!(previous > 0, "raw-node forwarding lease underflow"); + } +} + /// Filter for matching incoming stanzas (nodes) by tag and attributes. /// /// Used with [`Client::wait_for_node`] to wait for specific stanzas. @@ -1020,9 +1039,8 @@ pub struct Client { /// its allocation-churn snapshot. Unset unless that builder method was used. pub(crate) alloc_meter: std::sync::OnceLock>, - /// When true, emit `Event::RawNode` for every decoded stanza before router dispatch. - /// Default false — only enable when external consumers need raw protocol access. - raw_node_forwarding: AtomicBool, + /// Number of consumers currently requesting `Event::RawNode` forwarding. + raw_node_forwarding: AtomicUsize, /// Active VoIP calls and their media-task abort handles. `abort_all` runs from the /// connection-cleanup path so a disconnect/reconnect tears down every in-flight call. Behind the diff --git a/src/client/accessors.rs b/src/client/accessors.rs index c99e81e29..ac28e17a0 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -44,12 +44,24 @@ impl Client { self.core.event_bus.subscribe_handler(handler) } - /// Enable or disable raw node forwarding. - /// When enabled, `Event::RawNode` is emitted for every decoded stanza before - /// the stanza router dispatches it. Only enable when external consumers need - /// raw protocol access (e.g. voice call stanzas). - pub fn set_raw_node_forwarding(&self, enabled: bool) { - self.raw_node_forwarding.store(enabled, Ordering::Relaxed); + /// Acquire raw decoded stanza forwarding for one consumer. + /// + /// `Event::RawNode` remains enabled until every acquired lease is dropped. + pub fn acquire_raw_node_forwarding(self: &Arc) -> RawNodeLease { + let incremented = self + .raw_node_forwarding + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| { + count.checked_add(1) + }) + .is_ok(); + assert!(incremented, "raw-node forwarding lease counter overflow"); + RawNodeLease { + client: Arc::downgrade(self), + } + } + + pub(crate) fn raw_node_forwarding_enabled(&self) -> bool { + self.raw_node_forwarding.load(Ordering::Relaxed) != 0 } /// Enable or disable skipping of history sync notifications at runtime. @@ -528,6 +540,24 @@ impl Client { } } +#[cfg(test)] +mod raw_node_tests { + #[tokio::test] + async fn raw_node_forwarding_stays_enabled_until_the_last_lease_drops() { + let client = crate::test_utils::create_test_client().await; + assert!(!client.raw_node_forwarding_enabled()); + + let first = client.acquire_raw_node_forwarding(); + let second = client.acquire_raw_node_forwarding(); + assert!(client.raw_node_forwarding_enabled()); + + drop(first); + assert!(client.raw_node_forwarding_enabled()); + drop(second); + assert!(!client.raw_node_forwarding_enabled()); + } +} + #[cfg(test)] mod send_checks { fn assert_send(_: &T) {} diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 61b5f8e42..48b740a38 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -335,7 +335,7 @@ impl Client { self_weak: std::sync::OnceLock::new(), saver_handle: std::sync::OnceLock::new(), alloc_meter: std::sync::OnceLock::new(), - raw_node_forwarding: AtomicBool::new(false), + raw_node_forwarding: AtomicUsize::new(0), #[cfg(feature = "voip-runtime")] call_registry: std::sync::Arc::new(wacore::voip::CallRegistry::new()), #[cfg(feature = "voip-runtime")] diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 11d6ab7fa..3036354cd 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -192,13 +192,13 @@ impl Client { // observes these events or every raw node. "receipt" => { !self.synchronous_ack - && !self.raw_node_forwarding.load(Ordering::Relaxed) + && !self.raw_node_forwarding_enabled() && !self.core.event_bus.has_handler_for( wacore::types::events::EventKind::Receipt, ) } "ack" => { - !self.raw_node_forwarding.load(Ordering::Relaxed) + !self.raw_node_forwarding_enabled() && !self.core.event_bus.has_handler_for( wacore::types::events::EventKind::ServerAck, ) @@ -326,7 +326,7 @@ impl Client { // ACKs need shared ownership only for opt-in raw/node observers. The // usual response-waiter path borrows the node and can skip the Arc. if node.tag() == "ack" - && !self.raw_node_forwarding.load(Ordering::Relaxed) + && !self.raw_node_forwarding_enabled() && self.node_waiter_count.load(Ordering::Acquire) == 0 && !self.offline_sync_metrics.active.load(Ordering::Acquire) { @@ -457,7 +457,7 @@ impl Client { // Emit raw node before any early returns so all decoded stanzas // (including IQ responses and xmlstreamend) reach external observers - if self.raw_node_forwarding.load(Ordering::Relaxed) { + if self.raw_node_forwarding_enabled() { self.core .event_bus .dispatch(Event::RawNode(Arc::clone(&node))); diff --git a/src/lib.rs b/src/lib.rs index b2b083cda..7d6d264bc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -95,7 +95,7 @@ pub use client::{ pub use client::{CallError, Voip}; pub use client::{ Client, ClientBuild, ClientBuilder, ClientBuilderError, ClientLifecycle, ConnectionScope, - ConnectionScopeState, + ConnectionScopeState, RawNodeLease, }; pub use types::durability_hook::InboundDurabilityHook; pub use types::retry_admission::RetryAdmission; @@ -181,7 +181,7 @@ pub mod prelude { pub use crate::bot::{Bot, BotBuilder, BotHandle, EventDelivery, MessageContext}; pub use crate::client::{ Client, ClientBuilder, ClientBuilderError, ClientError, ClientLifecycle, ConnectionScope, - ConnectionScopeState, + ConnectionScopeState, RawNodeLease, }; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index 89ebb0aec..2eca302a2 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -906,7 +906,7 @@ pub enum Event { /// Raw decoded stanza, emitted before router dispatch. /// Library extension — no WA Web equivalent (WA Web has no raw stanza observer). - /// Gated by `Client::set_raw_node_forwarding(true)` to avoid overhead when unused. + /// Gated by `Client::acquire_raw_node_forwarding()` to avoid overhead when unused. #[serde(skip)] RawNode(Arc), From 3bd4bd58130bc9457e952f09191dbdd406bf967c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:07:49 -0300 Subject: [PATCH 09/46] fix(client): make lifecycle teardown cancellation-safe --- src/client/extension_lifecycle.rs | 204 ++++++++++++++++++++++++++++-- src/client/lifecycle.rs | 21 ++- 2 files changed, 210 insertions(+), 15 deletions(-) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 3e9ff5e80..ef888f9bd 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -304,9 +304,13 @@ impl LifecycleRegistration { } pub(super) fn close_scope(self: &Arc, generation: u64) { - let scope = { + self.close_scope_with(generation, || {}); + } + + fn close_scope_with(self: &Arc, generation: u64, after_remove: impl FnOnce()) { + let should_spawn = { let mut scopes = self.scopes(); - if scopes + let scope = if scopes .active .as_ref() .is_some_and(|scope| scope.generation() == generation) @@ -318,15 +322,32 @@ impl LifecycleRegistration { .iter() .position(|scope| scope.generation() == generation) .map(|position| scopes.retired.remove(position)) + }; + let Some(scope) = scope else { + return; + }; + + scope.close(); + after_remove(); + + // Publish closure before exposing an empty registry so terminal + // shutdown cannot overtake the final on_closed callback. + let no_open_scopes = scopes.active.is_none() && scopes.retired.is_empty(); + let mut queue = self.callback_queue(); + queue.pending.push_back(LifecycleCallback::Closed(scope)); + if no_open_scopes && queue.shutdown_requested && !queue.shutdown_enqueued { + queue.shutdown_enqueued = true; + queue.pending.push_back(LifecycleCallback::Shutdown); + } + if queue.drain_scheduled { + false + } else { + queue.drain_scheduled = true; + true } - }; - let Some(scope) = scope else { - return; }; - scope.close(); - self.enqueue_callback(LifecycleCallback::Closed(scope)); - self.enqueue_shutdown_if_ready(); + self.spawn_callback_driver(should_spawn); } pub(super) async fn shutdown(self: &Arc) { @@ -416,16 +437,18 @@ impl LifecycleRegistration { match callback { LifecycleCallback::Ready { scope, done } => { - self.run_callback("on_ready", self.handler.on_ready(scope.clone())) + let callback_scope = scope.clone(); + self.run_callback("on_ready", move |handler| handler.on_ready(callback_scope)) .await; let _ = done.try_send(!scope.is_cancelled()); } LifecycleCallback::Closed(scope) => { - self.run_callback("on_closed", self.handler.on_closed(scope)) + self.run_callback("on_closed", move |handler| handler.on_closed(scope)) .await; } LifecycleCallback::Shutdown => { - self.run_callback("shutdown", self.handler.shutdown()).await; + self.run_callback("shutdown", |handler| handler.shutdown()) + .await; self.shutdown_complete.store(true, Ordering::Release); self.shutdown_notifier.notify(); } @@ -433,11 +456,19 @@ impl LifecycleRegistration { } } - async fn run_callback( - &self, + async fn run_callback<'a>( + &'a self, name: &'static str, - mut callback: BoxFuture<'_, anyhow::Result<()>>, + create: impl FnOnce(&'a dyn ClientLifecycle) -> BoxFuture<'a, anyhow::Result<()>>, ) { + let callback = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _callback_context = CallbackContextGuard::enter(self); + create(&*self.handler) + })); + let Ok(mut callback) = callback else { + log::warn!("Client lifecycle {name} panicked"); + return; + }; let callback = std::future::poll_fn(|context| { let _callback_context = CallbackContextGuard::enter(self); callback.as_mut().poll(context) @@ -665,6 +696,30 @@ mod tests { } } + #[derive(Default)] + struct SynchronousPanicLifecycle { + ready_calls: AtomicUsize, + closed_calls: AtomicUsize, + shutdown_calls: AtomicUsize, + } + + impl ClientLifecycle for SynchronousPanicLifecycle { + fn on_ready<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + self.ready_calls.fetch_add(1, Ordering::SeqCst); + panic!("synchronous on_ready panic"); + } + + fn on_closed<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + self.closed_calls.fetch_add(1, Ordering::SeqCst); + panic!("synchronous on_closed panic"); + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + self.shutdown_calls.fetch_add(1, Ordering::SeqCst); + panic!("synchronous shutdown panic"); + } + } + struct LogoutOrderHandler { lifecycle: Arc, } @@ -1021,6 +1076,127 @@ mod tests { assert_eq!(lifecycle.calls.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn synchronous_callback_panics_do_not_strand_the_driver() { + let lifecycle = Arc::new(SynchronousPanicLifecycle::default()); + let registration = Arc::new(LifecycleRegistration::new( + lifecycle.clone(), + Arc::new(TokioRuntime), + )); + const GENERATION: u64 = 23; + assert!(registration.begin_scope_if_current(GENERATION, || true)); + + assert!(registration.ready(GENERATION).await); + registration.close_scope(GENERATION); + tokio::time::timeout(std::time::Duration::from_secs(2), registration.shutdown()) + .await + .expect("callback driver recovered from synchronous panics"); + + assert_eq!(lifecycle.ready_calls.load(Ordering::SeqCst), 1); + assert_eq!(lifecycle.closed_calls.load(Ordering::SeqCst), 1); + assert_eq!(lifecycle.shutdown_calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn final_scope_closure_is_published_before_shutdown() { + let lifecycle = Arc::new(RecordingLifecycle::default()); + let registration = Arc::new(LifecycleRegistration::new( + lifecycle.clone(), + Arc::new(TokioRuntime), + )); + const GENERATION: u64 = 29; + assert!(registration.begin_scope_if_current(GENERATION, || true)); + + let removed = Arc::new(std::sync::Barrier::new(2)); + let release = Arc::new(std::sync::Barrier::new(2)); + let close_registration = Arc::clone(®istration); + let close_removed = Arc::clone(&removed); + let close_release = Arc::clone(&release); + let close = tokio::task::spawn_blocking(move || { + close_registration.close_scope_with(GENERATION, || { + close_removed.wait(); + close_release.wait(); + }); + }); + + removed.wait(); + let shutdown_registration = Arc::clone(®istration); + let shutdown = tokio::spawn(async move { shutdown_registration.shutdown().await }); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + assert!(!shutdown.is_finished()); + + release.wait(); + close.await.expect("scope close task"); + shutdown.await.expect("shutdown task"); + assert_eq!( + lifecycle.events(), + vec!["closed:29".to_string(), "shutdown".to_string()] + ); + } + + #[tokio::test] + async fn cancelled_cleanup_waiter_does_not_strand_its_scope() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let lifecycle = Arc::new(RecordingLifecycle::default()); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await + .expect("client build") + .client(); + const GENERATION: u64 = 37; + client + .connection_generation + .store(GENERATION, Ordering::SeqCst); + let registration = client.lifecycle.as_ref().expect("lifecycle registration"); + assert!(registration.begin_scope_if_current(GENERATION, || true)); + client.dispatch_connected().await; + let scope = registration + .scope_for(GENERATION) + .expect("connection scope"); + + let (started_tx, started_rx) = async_channel::bounded(1); + let (release_tx, release_rx) = async_channel::bounded(1); + *client.transport.lock().await = Some(Arc::new(BlockingDisconnect { + started: started_tx, + release: release_rx, + })); + + let cleanup_client = Arc::clone(&client); + let cleanup = tokio::spawn(async move { + cleanup_client.cleanup_connection_state().await; + }); + started_rx.recv().await.expect("cleanup reached transport"); + cleanup.abort(); + let _ = cleanup.await; + assert_eq!(scope.state(), ConnectionScopeState::Cancelled); + + release_tx.send(()).await.expect("release cleanup"); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while scope.state() != ConnectionScopeState::Closed { + tokio::task::yield_now().await; + } + }) + .await + .expect("detached cleanup closed the scope"); + assert!(registration.scope_for(GENERATION).is_none()); + + registration.shutdown().await; + assert_eq!( + lifecycle.events(), + vec!["install", "ready:37", "closed:37", "shutdown"] + ); + client.signal_shutdown_sync(); + } + #[tokio::test] async fn stale_generation_is_rejected_before_scope_publication() { let lifecycle = Arc::new(RecordingLifecycle::default()); diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 48b740a38..4e89f558e 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -812,7 +812,26 @@ impl Client { feature = "tracing", tracing::instrument(name = "wa.conn.cleanup", level = "debug", skip_all) )] - pub(crate) async fn cleanup_connection_state(&self) { + pub(crate) async fn cleanup_connection_state(self: &Arc) { + if self.lifecycle.is_none() { + self.cleanup_connection_state_inner().await; + return; + } + + // Scope closure must survive a caller dropping its cleanup waiter. + let completed = wacore::runtime::ShutdownNotifier::new(); + let completion = completed.subscribe(); + let client = Arc::clone(self); + self.runtime + .spawn(Box::pin(async move { + client.cleanup_connection_state_inner().await; + completed.notify(); + })) + .detach(); + wacore::runtime::wait_for_shutdown(&completion).await; + } + + async fn cleanup_connection_state_inner(&self) { // Bump the generation FIRST: it is the "this connection is over" // signal every per-connection loop already polls. Chat-lane workers // stop draining their queues (their remaining stanzas were never From 66f464a9370b52dc8c2eea2af0e5b0b37f2849dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:21:00 -0300 Subject: [PATCH 10/46] docs(client): clarify lifecycle test hook constraints --- src/client/extension_lifecycle.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index ef888f9bd..fa3391b66 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -307,6 +307,7 @@ impl LifecycleRegistration { self.close_scope_with(generation, || {}); } + /// Non-noop hooks are test-only, must run off-executor, and must not re-enter lifecycle APIs. fn close_scope_with(self: &Arc, generation: u64, after_remove: impl FnOnce()) { let should_spawn = { let mut scopes = self.scopes(); From 835d61d2772cf5033812aacd15a848d645a9521b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 10:13:59 -0300 Subject: [PATCH 11/46] feat(plugins): add transactional native plugin host --- src/bot.rs | 52 + src/client.rs | 4 +- src/client/builder.rs | 66 +- src/client/extension_lifecycle.rs | 4 +- src/client/lifecycle.rs | 7 +- src/lib.rs | 10 + src/plugins/mod.rs | 1826 +++++++++++++++++++++++++++++ 7 files changed, 1956 insertions(+), 13 deletions(-) create mode 100644 src/plugins/mod.rs diff --git a/src/bot.rs b/src/bot.rs index e3436421c..c6e2990e7 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -1,6 +1,7 @@ use crate::cache_config::CacheConfig; use crate::client::{Client, ClientBuilderError}; use crate::pair_code::PairCodeOptions; +use crate::plugins::{ClientPlugin, PluginRegistration}; use crate::store::commands::DeviceCommand; use crate::store::error::StoreError; use crate::store::persistence_manager::PersistenceManager; @@ -671,6 +672,7 @@ pub struct BotBuilder< resend_rate_limit: Option<(u32, u32)>, task_instrument: Option>, alloc_meter: Option>, + plugins: Vec, _marker: PhantomData<(B, T, H, R)>, } @@ -696,6 +698,7 @@ impl BotBuilder BotBuilder { resend_rate_limit: self.resend_rate_limit, task_instrument: self.task_instrument, alloc_meter: self.alloc_meter, + plugins: self.plugins, _marker: PhantomData, } } @@ -844,6 +848,18 @@ impl BotBuilder { self } + /// Register a native plugin without changing the builder's typestate. + pub fn with_plugin(mut self, plugin: P) -> Self { + self.plugins.push(PluginRegistration::new(plugin)); + self + } + + /// Register an already-shared native plugin without changing its marker type. + pub fn with_plugin_arc(mut self, plugin: Arc

) -> Self { + self.plugins.push(PluginRegistration::new_arc(plugin)); + self + } + // ── Event handler registration (additive; order of registration is kept, // but handlers run on their own tasks, so cross-event ordering is not // guaranteed) ────────────────────────────────────────────────────── @@ -1288,6 +1304,7 @@ impl BotBuilder { .with_http_client_arc(http_client) .with_cache_config(self.cache_config) .with_custom_enc_handlers(self.custom_enc_handlers) + .with_plugin_registrations(self.plugins) .with_skip_history_sync(self.skip_history_sync) .with_background_saver_interval(std::time::Duration::from_secs(30)); @@ -1385,6 +1402,41 @@ mod tests { .client() } + struct BotBuilderPlugin; + + impl ClientPlugin for BotBuilderPlugin { + type Api = &'static str; + + fn manifest(&self) -> crate::plugins::PluginManifest { + crate::plugins::PluginManifest::new("bot-builder-test", "0.1.0") + } + + fn install( + &self, + _context: crate::plugins::PluginContext, + ) -> wacore::runtime::BoxFuture<'_, anyhow::Result>> { + Box::pin(async { Ok(Arc::new("installed")) }) + } + } + + #[tokio::test] + async fn typestate_builder_preserves_registered_plugins() { + let bot = Bot::builder() + .with_plugin(BotBuilderPlugin) + .with_backend_arc(create_test_sqlite_backend().await) + .with_transport_factory(TokioWebSocketTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_runtime(TokioRuntime) + .build() + .await + .expect("bot plugin build"); + assert_eq!( + bot.client().plugin::().as_deref(), + Some(&"installed") + ); + bot.client().disconnect().await; + } + fn pairing_code_event(code: &str) -> Arc { Arc::new(Event::PairingCode( crate::types::events::PairingCode::builder() diff --git a/src/client.rs b/src/client.rs index d82c1d585..83e39b3ae 100644 --- a/src/client.rs +++ b/src/client.rs @@ -15,7 +15,7 @@ pub(crate) mod offline_resume; mod sender_keys; mod sessions; mod voip; -use builder::ClientAssembly; +use builder::{ClientAssembly, ClientExtensions}; pub use builder::{ClientBuild, ClientBuilder, ClientBuilderError}; use extension_lifecycle::LifecycleRegistration; pub use extension_lifecycle::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; @@ -689,6 +689,8 @@ pub struct Client { pub(crate) connection_shutdown: std::sync::Mutex, /// Allocated only when an extension host installs lifecycle callbacks. lifecycle: Option>, + /// Allocated only when at least one build-time plugin is registered. + pub(crate) plugin_host: Option>, /// Per-session wire I/O and activity counters. Written at the transport /// chokepoints (noise sender task, read loop); the keepalive dead-socket /// watchdog reads its activity timestamps. Snapshot via [`Client::stats`]. diff --git a/src/client/builder.rs b/src/client/builder.rs index 36599b7cf..a975621ff 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -7,6 +7,7 @@ use thiserror::Error; use super::{Client, ClientLifecycle, LifecycleRegistration}; use crate::cache_config::CacheConfig; use crate::http::HttpClient; +use crate::plugins::{ClientPlugin, PluginHost, PluginPlan, PluginPlanError, PluginRegistration}; use crate::store::error::StoreError; use crate::store::persistence_manager::PersistenceManager; use crate::sync_task::MajorSyncTask; @@ -62,6 +63,10 @@ pub enum ClientBuilderError { UnsupportedDurabilityBackend(String), #[error("client lifecycle installation failed: {0}")] LifecycleInstall(#[source] anyhow::Error), + #[error("plugin host installation failed: {0}")] + PluginInstall(#[source] anyhow::Error), + #[error("invalid plugin plan: {0}")] + PluginPlan(#[from] PluginPlanError), } /// Runtime-validated, low-level builder for [`Client`]. @@ -85,6 +90,7 @@ pub struct ClientBuilder { alloc_meter: Option>, background_saver_interval: Option, lifecycle: Option>, + plugins: Vec, } impl Default for ClientBuilder { @@ -112,6 +118,7 @@ impl ClientBuilder { alloc_meter: None, background_saver_interval: None, lifecycle: None, + plugins: Vec::new(), } } @@ -275,6 +282,26 @@ impl ClientBuilder { self } + /// Register a native plugin for transactional installation before services start. + pub fn with_plugin(mut self, plugin: P) -> Self { + self.plugins.push(PluginRegistration::new(plugin)); + self + } + + /// Register an already-shared native plugin without changing its marker type. + pub fn with_plugin_arc(mut self, plugin: Arc

) -> Self { + self.plugins.push(PluginRegistration::new_arc(plugin)); + self + } + + pub(crate) fn with_plugin_registrations( + mut self, + registrations: Vec, + ) -> Self { + self.plugins = registrations; + self + } + /// Validate dependencies, assemble an inert client, then start its services. pub async fn build(self) -> Result { self.build_boxed().await @@ -343,6 +370,7 @@ impl ClientBuilder { transport_factory: Arc, http_client: Arc, ) -> Result { + let plugin_plan = PluginPlan::prepare(self.plugins)?; let runtime: Arc = match self.task_instrument { Some(instrument) => { Arc::new(wacore::stats::InstrumentedRuntime::new(runtime, instrument)) @@ -350,8 +378,13 @@ impl ClientBuilder { None => runtime, }; - let lifecycle = self - .lifecycle + let mut lifecycle_handler = self.lifecycle; + let plugin_host = plugin_plan.map(|plan| { + let host = PluginHost::new(plan, lifecycle_handler.take()); + lifecycle_handler = Some(host.clone()); + host + }); + let lifecycle = lifecycle_handler .map(|handler| Arc::new(LifecycleRegistration::new(handler, Arc::clone(&runtime)))); let assembly = Client::assemble( Arc::clone(&runtime), @@ -360,7 +393,10 @@ impl ClientBuilder { http_client, self.override_version, self.cache_config, - lifecycle, + ClientExtensions { + lifecycle, + plugin_host, + }, ); let client = assembly.client(); @@ -382,11 +418,14 @@ impl ClientBuilder { if let Some(meter) = self.alloc_meter { let _ = client.alloc_meter.set(meter); } - if let Some(lifecycle) = &client.lifecycle { - lifecycle - .install(Arc::downgrade(&client)) - .await - .map_err(ClientBuilderError::LifecycleInstall)?; + if let Some(lifecycle) = &client.lifecycle + && let Err(error) = lifecycle.install(Arc::downgrade(&client)).await + { + return Err(if client.plugin_host.is_some() { + ClientBuilderError::PluginInstall(error) + } else { + ClientBuilderError::LifecycleInstall(error) + }); } let build = assembly.start(); @@ -398,6 +437,9 @@ impl ClientBuilder { ); let _ = build.client.saver_handle.set(saver_handle); } + if let Some(plugin_host) = &client.plugin_host { + plugin_host.activate(); + } Ok(build) } } @@ -446,6 +488,12 @@ pub(super) struct ClientAssembly { sync_task_receiver: async_channel::Receiver, } +#[derive(Default)] +pub(super) struct ClientExtensions { + pub(super) lifecycle: Option>, + pub(super) plugin_host: Option>, +} + impl ClientAssembly { pub(super) fn new( client: Arc, @@ -587,7 +635,7 @@ mod tests { Arc::new(MockHttpClient), None, CacheConfig::default(), - None, + ClientExtensions::default(), ); assert_eq!(spawns.load(Ordering::SeqCst), 0); diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index fa3391b66..35e4e31ac 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -47,7 +47,7 @@ pub struct ConnectionScope { } impl ConnectionScope { - fn new(generation: u64) -> Self { + pub(crate) fn new(generation: u64) -> Self { Self { inner: Arc::new(ConnectionScopeInner { generation, @@ -88,7 +88,7 @@ impl ConnectionScope { .is_ok() } - fn cancel(&self) { + pub(crate) fn cancel(&self) { let mut state = self.inner.state.load(Ordering::Acquire); while state < SCOPE_CANCELLED { match self.inner.state.compare_exchange_weak( diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 4e89f558e..31aa2556f 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -152,8 +152,12 @@ impl Client { http_client: Arc, override_version: Option<(u32, u32, u32)>, cache_config: CacheConfig, - lifecycle: Option>, + extensions: ClientExtensions, ) -> ClientAssembly { + let ClientExtensions { + lifecycle, + plugin_host, + } = extensions; let mut unique_id_bytes = [0u8; 2]; rand::make_rng::().fill_bytes(&mut unique_id_bytes); @@ -181,6 +185,7 @@ impl Client { shutdown_notifier: wacore::runtime::ShutdownNotifier::new(), connection_shutdown: std::sync::Mutex::new(wacore::runtime::ShutdownNotifier::new()), lifecycle, + plugin_host, stats: Arc::new(wacore::stats::SessionStats::new()), transport: Arc::new(Mutex::new(None)), diff --git a/src/lib.rs b/src/lib.rs index 7d6d264bc..544d464d1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -111,6 +111,13 @@ pub(crate) mod msg_secret_buffer; pub mod pair; pub mod pair_code; pub mod passkey; +pub mod plugins; +pub use plugins::{ + ClientPlugin, PluginCapabilities, PluginCapability, PluginConnectionScope, + PluginConnectionTasks, PluginContext, PluginCoreEvents, PluginIq, PluginIqError, + PluginManifest, PluginMessaging, PluginMessagingError, PluginPlanError, PluginResourceError, + PluginTasks, +}; pub mod request; pub(crate) mod signal_flush; pub use request::IqError; @@ -183,6 +190,9 @@ pub mod prelude { Client, ClientBuilder, ClientBuilderError, ClientError, ClientLifecycle, ConnectionScope, ConnectionScopeState, RawNodeLease, }; + pub use crate::plugins::{ + ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, PluginManifest, + }; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] pub use crate::runtime_impl::TokioRuntime; diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs new file mode 100644 index 000000000..83fdf9eae --- /dev/null +++ b/src/plugins/mod.rs @@ -0,0 +1,1826 @@ +//! Build-time client plugins and their capability-scoped host. + +use std::any::{Any, TypeId}; +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::future::Future; +use std::panic::AssertUnwindSafe; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, OnceLock, Weak}; + +use futures::FutureExt; +use thiserror::Error; +use wacore::iq::spec::IqSpec; +use wacore::runtime::{ + BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, Spawnable, wait_for_shutdown, +}; +use wacore::sync_marker::MaybeSendSync; +use wacore::types::events::{EventHandler, EventInterest, Subscription}; +use wacore_binary::Jid; +use waproto::whatsapp::Message; + +use crate::Client; +use crate::client::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; +use crate::request::IqError; +use crate::send::{SendError, SendResult}; + +const CAP_CORE_EVENTS: u8 = 1 << 0; +const CAP_TASKS: u8 = 1 << 1; +const CAP_MESSAGING: u8 = 1 << 2; +const CAP_IQ: u8 = 1 << 3; + +/// A capability a plugin asks the host to expose during installation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[non_exhaustive] +pub enum PluginCapability { + CoreEvents, + Tasks, + Messaging, + Iq, +} + +impl PluginCapability { + pub const fn identifier(self) -> &'static str { + match self { + Self::CoreEvents => "events.core.observe", + Self::Tasks => "tasks.spawn", + Self::Messaging => "messaging.send", + Self::Iq => "iq.execute", + } + } + + const fn bit(self) -> u8 { + match self { + Self::CoreEvents => CAP_CORE_EVENTS, + Self::Tasks => CAP_TASKS, + Self::Messaging => CAP_MESSAGING, + Self::Iq => CAP_IQ, + } + } +} + +/// Compact set of capabilities requested by one plugin. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct PluginCapabilities(u8); + +impl PluginCapabilities { + pub const NONE: Self = Self(0); + + pub const fn with(self, capability: PluginCapability) -> Self { + Self(self.0 | capability.bit()) + } + + pub const fn contains(self, capability: PluginCapability) -> bool { + self.0 & capability.bit() != 0 + } +} + +/// Build-time declaration used for validation, ordering, and future foreign adapters. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub struct PluginManifest { + id: String, + version: String, + dependencies: Vec, + capabilities: PluginCapabilities, +} + +impl PluginManifest { + pub fn new(id: impl Into, version: impl Into) -> Self { + Self { + id: id.into(), + version: version.into(), + dependencies: Vec::new(), + capabilities: PluginCapabilities::NONE, + } + } + + pub fn with_dependency(mut self, plugin_id: impl Into) -> Self { + self.dependencies.push(plugin_id.into()); + self + } + + pub const fn with_capability(mut self, capability: PluginCapability) -> Self { + self.capabilities = self.capabilities.with(capability); + self + } + + pub fn id(&self) -> &str { + &self.id + } + + pub fn version(&self) -> &str { + &self.version + } + + pub fn dependencies(&self) -> &[String] { + &self.dependencies + } + + pub const fn capabilities(&self) -> PluginCapabilities { + self.capabilities + } +} + +/// A trusted native plugin installed exactly once while the client is still inert. +/// Capabilities shape the handles it receives; they are not an in-process sandbox. +pub trait ClientPlugin: MaybeSendSync + 'static { + type Api: MaybeSendSync + 'static; + + fn manifest(&self) -> PluginManifest; + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>>; + + fn on_ready(&self, _scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } + + fn on_closed(&self, _scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } + + /// Release plugin-owned state. This may run after `install` began but returned an error. + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } +} + +/// Manifest validation or dependency-ordering failure. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum PluginPlanError { + #[error("plugin {plugin_type} panicked while producing its manifest")] + ManifestPanicked { plugin_type: &'static str }, + #[error("invalid plugin id `{id}`")] + InvalidId { id: String }, + #[error("plugin `{plugin_id}` has an invalid version `{version}`")] + InvalidVersion { plugin_id: String, version: String }, + #[error("plugin id `{id}` is registered more than once")] + DuplicateId { id: String }, + #[error("plugin marker type `{plugin_type}` is registered more than once")] + DuplicateType { plugin_type: &'static str }, + #[error("plugin `{plugin_id}` lists dependency `{dependency}` more than once")] + DuplicateDependency { + plugin_id: String, + dependency: String, + }, + #[error("plugin `{plugin_id}` requires missing plugin `{dependency}`")] + MissingDependency { + plugin_id: String, + dependency: String, + }, + #[error("plugin dependency cycle involves: {plugins:?}")] + DependencyCycle { plugins: Vec }, +} + +/// Capability use after the client or plugin scope has ended. +#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum PluginResourceError { + #[error("the client is no longer available")] + ClientUnavailable, + #[error("the plugin host has not started yet")] + NotActive, + #[error("the plugin scope is shutting down")] + ShuttingDown, +} + +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum PluginMessagingError { + #[error(transparent)] + Resource(#[from] PluginResourceError), + #[error(transparent)] + Send(#[from] SendError), +} + +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum PluginIqError { + #[error(transparent)] + Resource(#[from] PluginResourceError), + #[error(transparent)] + Iq(#[from] IqError), +} + +struct PluginResources { + active: AtomicBool, + closed: AtomicBool, + activation: ShutdownNotifier, + shutdown: ShutdownNotifier, + subscriptions: Mutex>, +} + +impl PluginResources { + fn new() -> Arc { + Arc::new(Self { + active: AtomicBool::new(false), + closed: AtomicBool::new(false), + activation: ShutdownNotifier::new(), + shutdown: ShutdownNotifier::new(), + subscriptions: Mutex::new(Vec::new()), + }) + } + + fn activate(&self) { + if self.closed.load(Ordering::Acquire) { + return; + } + self.active.store(true, Ordering::Release); + self.activation.notify(); + } + + fn ensure_active(&self) -> Result<(), PluginResourceError> { + if self.closed.load(Ordering::Acquire) { + Err(PluginResourceError::ShuttingDown) + } else if !self.active.load(Ordering::Acquire) { + Err(PluginResourceError::NotActive) + } else { + Ok(()) + } + } + + fn retain_subscription(&self, subscription: Subscription) -> Result<(), PluginResourceError> { + let mut subscriptions = self + .subscriptions + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.closed.load(Ordering::Acquire) { + drop(subscription); + return Err(PluginResourceError::ShuttingDown); + } + subscriptions.push(subscription); + Ok(()) + } + + fn close(&self) { + if self.closed.swap(true, Ordering::AcqRel) { + return; + } + self.shutdown.notify(); + self.subscriptions + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clear(); + } +} + +/// Install-scoped task capability. Work starts after the complete plugin set is published and +/// stops during rollback or shutdown. +#[derive(Clone)] +pub struct PluginTasks { + runtime: Arc, + resources: Arc, +} + +impl PluginTasks { + pub fn spawn(&self, future: F) -> Result<(), PluginResourceError> + where + F: Future + Spawnable, + { + if self.resources.closed.load(Ordering::Acquire) { + return Err(PluginResourceError::ShuttingDown); + } + spawn_after_activation(&self.runtime, Arc::clone(&self.resources), future); + Ok(()) + } + + pub fn shutdown_signal(&self) -> ShutdownSignal { + self.resources.shutdown.subscribe() + } +} + +/// Selective subscription access to the sealed core event bus. +/// Handlers run inline and must hand slow work to a task capability. +#[derive(Clone)] +pub struct PluginCoreEvents { + client: Weak, + resources: Arc, +} + +impl PluginCoreEvents { + pub fn subscribe( + &self, + interest: EventInterest, + handler: Arc, + ) -> Result<(), PluginResourceError> { + let client = self + .client + .upgrade() + .ok_or(PluginResourceError::ClientUnavailable)?; + let subscription = client.subscribe(interest, handler); + self.resources.retain_subscription(subscription) + } +} + +/// High-level message sending without exposing the raw client or backend. +#[derive(Clone)] +pub struct PluginMessaging { + client: Weak, + resources: Arc, +} + +impl PluginMessaging { + pub async fn send_message( + &self, + to: Jid, + message: Message, + ) -> Result { + self.resources.ensure_active()?; + let client = self + .client + .upgrade() + .ok_or(PluginResourceError::ClientUnavailable)?; + Ok(client.send_message(to, message).await?) + } + + pub async fn send_text( + &self, + to: Jid, + text: String, + ) -> Result { + self.resources.ensure_active()?; + let client = self + .client + .upgrade() + .ok_or(PluginResourceError::ClientUnavailable)?; + Ok(client.send_text(to, text).await?) + } +} + +/// Typed IQ execution without exposing the raw client or stores. +#[derive(Clone)] +pub struct PluginIq { + client: Weak, + resources: Arc, +} + +impl PluginIq { + pub async fn execute(&self, spec: S) -> Result + where + S: IqSpec, + { + self.resources.ensure_active()?; + let client = self + .client + .upgrade() + .ok_or(PluginResourceError::ClientUnavailable)?; + Ok(client.execute(spec).await?) + } +} + +/// Capabilities and already-installed dependencies visible during installation. +pub struct PluginContext { + plugin_id: String, + apis: Arc, + dependency_markers: Vec, + core_events: Option, + tasks: Option, + messaging: Option, + iq: Option, +} + +impl PluginContext { + pub fn plugin_id(&self) -> &str { + &self.plugin_id + } + + pub fn plugin(&self) -> Option> { + if !self.dependency_markers.contains(&TypeId::of::

()) { + return None; + } + self.apis.get::

() + } + + pub fn core_events(&self) -> Option<&PluginCoreEvents> { + self.core_events.as_ref() + } + + pub fn tasks(&self) -> Option<&PluginTasks> { + self.tasks.as_ref() + } + + pub fn messaging(&self) -> Option<&PluginMessaging> { + self.messaging.as_ref() + } + + pub fn iq(&self) -> Option<&PluginIq> { + self.iq.as_ref() + } +} + +/// One connection generation plus its optional connection-scoped task capability. +#[derive(Clone)] +pub struct PluginConnectionScope { + scope: ConnectionScope, + tasks: Option, +} + +impl PluginConnectionScope { + pub fn generation(&self) -> u64 { + self.scope.generation() + } + + pub fn state(&self) -> ConnectionScopeState { + self.scope.state() + } + + pub fn is_cancelled(&self) -> bool { + self.scope.is_cancelled() + } + + pub fn cancellation_signal(&self) -> ShutdownSignal { + self.scope.cancellation_signal() + } + + pub fn tasks(&self) -> Option<&PluginConnectionTasks> { + self.tasks.as_ref() + } +} + +/// Task capability whose work is aborted synchronously when its generation retires. +#[derive(Clone)] +pub struct PluginConnectionTasks { + runtime: Arc, + scope: ConnectionScope, +} + +impl PluginConnectionTasks { + pub fn spawn(&self, future: F) -> Result<(), PluginResourceError> + where + F: Future + Spawnable, + { + if self.scope.is_cancelled() { + return Err(PluginResourceError::ShuttingDown); + } + spawn_until_cancelled(&self.runtime, self.scope.cancellation_signal(), future); + Ok(()) + } +} + +fn spawn_until_cancelled(runtime: &Arc, cancellation: ShutdownSignal, future: F) +where + F: Future + Spawnable, +{ + runtime + .spawn(Box::pin(async move { + let cancelled = Box::pin(wait_for_shutdown(&cancellation)); + let work = Box::pin(future); + let _ = futures::future::select(cancelled, work).await; + })) + .detach(); +} + +fn spawn_after_activation(runtime: &Arc, resources: Arc, future: F) +where + F: Future + Spawnable, +{ + let activation = resources.activation.subscribe(); + let cancellation = resources.shutdown.subscribe(); + runtime + .spawn(Box::pin(async move { + let cancelled = Box::pin(wait_for_shutdown(&cancellation)); + let activated = Box::pin(wait_for_shutdown(&activation)); + if matches!( + futures::future::select(cancelled, activated).await, + futures::future::Either::Left(_) + ) { + return; + } + if resources.closed.load(Ordering::Acquire) { + return; + } + let cancelled = Box::pin(wait_for_shutdown(&cancellation)); + let work = Box::pin(future); + let _ = futures::future::select(cancelled, work).await; + })) + .detach(); +} + +trait ErasedApiValue: MaybeSendSync { + fn as_any(&self) -> &dyn Any; +} + +struct TypedApi(Arc); + +impl ErasedApiValue for TypedApi { + fn as_any(&self) -> &dyn Any { + self + } +} + +type ErasedApi = Arc; + +#[derive(Default)] +struct ApiRegistry { + values: Mutex>, +} + +impl ApiRegistry { + fn insert(&self, marker: TypeId, api: ErasedApi) { + self.values + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .insert(marker, api); + } + + fn snapshot(&self) -> HashMap { + self.values + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + } + + fn get(&self) -> Option> { + let values = self + .values + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + downcast_api::(values.get(&TypeId::of::

())?) + } +} + +fn downcast_api(api: &ErasedApi) -> Option> { + api.as_any() + .downcast_ref::>() + .map(|typed| typed.0.clone()) +} + +trait ErasedClientPlugin: MaybeSendSync { + fn marker_type_id(&self) -> TypeId; + fn marker_type_name(&self) -> &'static str; + fn manifest(&self) -> PluginManifest; + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>; + fn on_ready(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>>; + fn on_closed(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>>; + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>>; +} + +struct PluginAdapter

(Arc

); + +impl ErasedClientPlugin for PluginAdapter

{ + fn marker_type_id(&self) -> TypeId { + TypeId::of::

() + } + + fn marker_type_name(&self) -> &'static str { + std::any::type_name::

() + } + + fn manifest(&self) -> PluginManifest { + self.0.manifest() + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result> { + Box::pin(async move { + let api = self.0.install(context).await?; + Ok(Arc::new(TypedApi(api)) as ErasedApi) + }) + } + + fn on_ready(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + self.0.on_ready(scope) + } + + fn on_closed(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + self.0.on_closed(scope) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + self.0.shutdown() + } +} + +pub(crate) struct PluginRegistration { + plugin: Arc, +} + +impl PluginRegistration { + pub(crate) fn new(plugin: P) -> Self { + Self::new_arc(Arc::new(plugin)) + } + + pub(crate) fn new_arc(plugin: Arc

) -> Self { + Self { + plugin: Arc::new(PluginAdapter(plugin)), + } + } +} + +struct PlannedPlugin { + plugin: Arc, + manifest: PluginManifest, + dependency_markers: Vec, +} + +pub(crate) struct PluginPlan { + ordered: Vec, +} + +impl PluginPlan { + pub(crate) fn prepare( + registrations: Vec, + ) -> Result, PluginPlanError> { + if registrations.is_empty() { + return Ok(None); + } + + let mut plugins = Vec::with_capacity(registrations.len()); + let mut ids = HashMap::with_capacity(registrations.len()); + let mut marker_types = HashSet::with_capacity(registrations.len()); + + for registration in registrations { + let plugin = registration.plugin; + let marker = plugin.marker_type_id(); + if !marker_types.insert(marker) { + return Err(PluginPlanError::DuplicateType { + plugin_type: plugin.marker_type_name(), + }); + } + let manifest = std::panic::catch_unwind(AssertUnwindSafe(|| plugin.manifest())) + .map_err(|_| PluginPlanError::ManifestPanicked { + plugin_type: plugin.marker_type_name(), + })?; + validate_manifest(&manifest)?; + let index = plugins.len(); + if ids.insert(manifest.id.clone(), index).is_some() { + return Err(PluginPlanError::DuplicateId { + id: manifest.id.clone(), + }); + } + plugins.push(PlannedPlugin { + plugin, + manifest, + dependency_markers: Vec::new(), + }); + } + + let mut indegree = vec![0usize; plugins.len()]; + let mut dependents = vec![Vec::new(); plugins.len()]; + let mut dependency_markers = vec![Vec::new(); plugins.len()]; + for (plugin_index, planned) in plugins.iter().enumerate() { + let mut seen = HashSet::with_capacity(planned.manifest.dependencies.len()); + for dependency in &planned.manifest.dependencies { + if !seen.insert(dependency) { + return Err(PluginPlanError::DuplicateDependency { + plugin_id: planned.manifest.id.clone(), + dependency: dependency.clone(), + }); + } + let Some(&dependency_index) = ids.get(dependency) else { + return Err(PluginPlanError::MissingDependency { + plugin_id: planned.manifest.id.clone(), + dependency: dependency.clone(), + }); + }; + indegree[plugin_index] += 1; + dependents[dependency_index].push(plugin_index); + dependency_markers[plugin_index] + .push(plugins[dependency_index].plugin.marker_type_id()); + } + } + for (planned, markers) in plugins.iter_mut().zip(dependency_markers) { + planned.dependency_markers = markers; + } + + let mut ready = indegree + .iter() + .enumerate() + .filter_map(|(index, count)| (*count == 0).then_some(index)) + .collect::>(); + let mut order = Vec::with_capacity(plugins.len()); + while let Some(index) = ready.pop_first() { + order.push(index); + for &dependent in &dependents[index] { + indegree[dependent] -= 1; + if indegree[dependent] == 0 { + ready.insert(dependent); + } + } + } + + if order.len() != plugins.len() { + let cycle = indegree + .iter() + .enumerate() + .filter(|(_, count)| **count > 0) + .map(|(index, _)| plugins[index].manifest.id.clone()) + .collect(); + return Err(PluginPlanError::DependencyCycle { plugins: cycle }); + } + + let mut slots = plugins.into_iter().map(Some).collect::>(); + let ordered = order + .into_iter() + .filter_map(|index| slots[index].take()) + .collect(); + Ok(Some(Self { ordered })) + } +} + +fn validate_manifest(manifest: &PluginManifest) -> Result<(), PluginPlanError> { + if !valid_plugin_id(&manifest.id) { + return Err(PluginPlanError::InvalidId { + id: manifest.id.clone(), + }); + } + if manifest.version.is_empty() + || manifest.version.len() > 64 + || !manifest.version.bytes().all(|byte| byte.is_ascii_graphic()) + { + return Err(PluginPlanError::InvalidVersion { + plugin_id: manifest.id.clone(), + version: manifest.version.clone(), + }); + } + Ok(()) +} + +fn valid_plugin_id(id: &str) -> bool { + if id.is_empty() || id.len() > 128 { + return false; + } + let mut previous_separator = true; + for byte in id.bytes() { + let separator = matches!(byte, b'.' | b'-' | b'_'); + if separator { + if previous_separator { + return false; + } + } else if !byte.is_ascii_lowercase() && !byte.is_ascii_digit() { + return false; + } + previous_separator = separator; + } + !previous_separator && id.as_bytes()[0].is_ascii_lowercase() +} + +struct InstalledPlugin { + plugin: Arc, + manifest: PluginManifest, + resources: Arc, +} + +pub(crate) struct PluginHost { + ordered: Vec, + manifests: Vec, + upstream: Option>, + installed: OnceLock>, + apis: OnceLock>, + runtime: OnceLock>, +} + +impl PluginHost { + pub(crate) fn new(plan: PluginPlan, upstream: Option>) -> Arc { + let manifests = plan + .ordered + .iter() + .map(|plugin| plugin.manifest.clone()) + .collect(); + Arc::new(Self { + ordered: plan.ordered, + manifests, + upstream, + installed: OnceLock::new(), + apis: OnceLock::new(), + runtime: OnceLock::new(), + }) + } + + pub(crate) fn plugin(&self) -> Option> { + downcast_api::(self.apis.get()?.get(&TypeId::of::

())?) + } + + pub(crate) fn manifests(&self) -> &[PluginManifest] { + &self.manifests + } + + fn context( + &self, + client: &Weak, + manifest: &PluginManifest, + dependency_markers: &[TypeId], + resources: Arc, + apis: Arc, + runtime: Arc, + ) -> PluginContext { + let capabilities = manifest.capabilities; + PluginContext { + plugin_id: manifest.id.clone(), + apis, + dependency_markers: dependency_markers.to_vec(), + core_events: capabilities + .contains(PluginCapability::CoreEvents) + .then(|| PluginCoreEvents { + client: client.clone(), + resources: Arc::clone(&resources), + }), + tasks: capabilities + .contains(PluginCapability::Tasks) + .then(|| PluginTasks { + runtime: Arc::clone(&runtime), + resources: Arc::clone(&resources), + }), + messaging: capabilities.contains(PluginCapability::Messaging).then(|| { + PluginMessaging { + client: client.clone(), + resources: Arc::clone(&resources), + } + }), + iq: capabilities + .contains(PluginCapability::Iq) + .then(|| PluginIq { + client: client.clone(), + resources: Arc::clone(&resources), + }), + } + } + + fn connection_scope( + &self, + scope: ConnectionScope, + manifest: &PluginManifest, + ) -> PluginConnectionScope { + let tasks = manifest + .capabilities + .contains(PluginCapability::Tasks) + .then(|| self.runtime.get().cloned()) + .flatten() + .map(|runtime| PluginConnectionTasks { + runtime, + scope: scope.clone(), + }); + PluginConnectionScope { scope, tasks } + } + + async fn install_all(&self, client: Weak) -> anyhow::Result<()> { + let Some(strong_client) = client.upgrade() else { + anyhow::bail!("client was dropped during plugin installation"); + }; + let runtime = strong_client.runtime.clone(); + drop(strong_client); + self.runtime + .set(runtime.clone()) + .map_err(|_| anyhow::anyhow!("plugin host was installed more than once"))?; + + if let Some(upstream) = &self.upstream { + plugin_callback(|| upstream.install(client.clone())).await?; + } + + let staging = Arc::new(ApiRegistry::default()); + let mut installed: Vec = Vec::with_capacity(self.ordered.len()); + for planned in &self.ordered { + let resources = PluginResources::new(); + let context = self.context( + &client, + &planned.manifest, + &planned.dependency_markers, + Arc::clone(&resources), + Arc::clone(&staging), + runtime.clone(), + ); + let api = match plugin_install(|| planned.plugin.install(context)).await { + Ok(api) => api, + Err(error) => { + resources.close(); + for plugin in installed.iter().rev() { + plugin.resources.close(); + } + if let Err(shutdown_error) = plugin_callback(|| planned.plugin.shutdown()).await + { + log::warn!( + "Plugin `{}` failed-install rollback failed: {shutdown_error:#}", + planned.manifest.id + ); + } + self.rollback(&mut installed).await; + anyhow::bail!( + "plugin `{}` installation failed: {error:#}", + planned.manifest.id + ); + } + }; + staging.insert(planned.plugin.marker_type_id(), api); + installed.push(InstalledPlugin { + plugin: planned.plugin.clone(), + manifest: planned.manifest.clone(), + resources, + }); + } + + self.apis + .set(staging.snapshot()) + .map_err(|_| anyhow::anyhow!("plugin APIs were published more than once"))?; + self.installed + .set(installed) + .map_err(|_| anyhow::anyhow!("plugins were installed more than once"))?; + Ok(()) + } + + pub(crate) fn activate(&self) { + for plugin in self.installed.get().into_iter().flatten() { + plugin.resources.activate(); + } + } + + async fn rollback(&self, installed: &mut Vec) { + for plugin in installed.iter().rev() { + plugin.resources.close(); + } + while let Some(plugin) = installed.pop() { + if plugin_callback(|| plugin.plugin.shutdown()).await.is_err() { + log::warn!("Plugin `{}` rollback failed", plugin.manifest.id); + } + } + if let Some(upstream) = &self.upstream + && let Err(error) = plugin_callback(|| upstream.shutdown()).await + { + log::warn!("Upstream lifecycle rollback failed: {error:#}"); + } + } +} + +impl ClientLifecycle for PluginHost { + fn install(&self, client: Weak) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { self.install_all(client).await }) + } + + fn on_ready(&self, scope: ConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + if let Some(upstream) = &self.upstream { + plugin_callback(|| upstream.on_ready(scope.clone())).await?; + } + let mut failures = Vec::new(); + for plugin in self.installed.get().into_iter().flatten() { + let plugin_scope = self.connection_scope(scope.clone(), &plugin.manifest); + if let Err(error) = plugin_callback(|| plugin.plugin.on_ready(plugin_scope)).await { + failures.push(format!("{}: {error:#}", plugin.manifest.id)); + } + } + finish_callbacks("ready", failures) + }) + } + + fn on_closed(&self, scope: ConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + let mut failures = Vec::new(); + for plugin in self.installed.get().into_iter().flatten().rev() { + let plugin_scope = self.connection_scope(scope.clone(), &plugin.manifest); + if let Err(error) = plugin_callback(|| plugin.plugin.on_closed(plugin_scope)).await + { + failures.push(format!("{}: {error:#}", plugin.manifest.id)); + } + } + if let Some(upstream) = &self.upstream + && let Err(error) = plugin_callback(|| upstream.on_closed(scope)).await + { + failures.push(format!("upstream: {error:#}")); + } + finish_callbacks("closed", failures) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + let mut failures = Vec::new(); + for plugin in self.installed.get().into_iter().flatten().rev() { + plugin.resources.close(); + } + for plugin in self.installed.get().into_iter().flatten().rev() { + if let Err(error) = plugin_callback(|| plugin.plugin.shutdown()).await { + failures.push(format!("{}: {error:#}", plugin.manifest.id)); + } + } + if let Some(upstream) = &self.upstream + && let Err(error) = plugin_callback(|| upstream.shutdown()).await + { + failures.push(format!("upstream: {error:#}")); + } + finish_callbacks("shutdown", failures) + }) + } +} + +impl Drop for PluginHost { + fn drop(&mut self) { + for plugin in self.installed.get().into_iter().flatten() { + plugin.resources.close(); + } + } +} + +async fn plugin_callback<'a>( + make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result<()>>, +) -> anyhow::Result<()> { + let future = std::panic::catch_unwind(AssertUnwindSafe(make_future)) + .map_err(|_| anyhow::anyhow!("callback panicked before returning a future"))?; + AssertUnwindSafe(future) + .catch_unwind() + .await + .map_err(|_| anyhow::anyhow!("callback future panicked"))? +} + +async fn plugin_install<'a>( + make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result>, +) -> anyhow::Result { + let future = std::panic::catch_unwind(AssertUnwindSafe(make_future)) + .map_err(|_| anyhow::anyhow!("install panicked before returning a future"))?; + AssertUnwindSafe(future) + .catch_unwind() + .await + .map_err(|_| anyhow::anyhow!("install future panicked"))? +} + +fn finish_callbacks(stage: &str, failures: Vec) -> anyhow::Result<()> { + if failures.is_empty() { + Ok(()) + } else { + anyhow::bail!("plugin {stage} callbacks failed: {}", failures.join("; ")) + } +} + +impl Client { + /// Return the API exposed by plugin marker `P`, if that plugin was installed. + pub fn plugin(&self) -> Option> { + self.plugin_host.as_ref()?.plugin::

() + } + + /// Manifests in dependency-resolved installation order. + pub fn plugin_manifests(&self) -> &[PluginManifest] { + self.plugin_host + .as_ref() + .map(|host| host.manifests()) + .unwrap_or_default() + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicBool; + use std::time::Duration; + + use super::*; + use crate::client::{ClientBuilder, ClientBuilderError}; + use crate::runtime_impl::TokioRuntime; + use crate::store::persistence_manager::PersistenceManager; + use crate::test_utils::MockHttpClient; + use crate::transport::mock::MockTransportFactory; + + type Log = Arc>>; + + fn record(log: &Log, value: impl Into) { + log.lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(value.into()); + } + + async fn complete_builder() -> ClientBuilder { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + ClientBuilder::new() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + } + + struct FoundationPlugin { + log: Log, + } + + impl ClientPlugin for FoundationPlugin { + type Api = String; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("foundation", "0.1.0") + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "install:foundation"); + Ok(Arc::new("foundation-api".to_string())) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "shutdown:foundation"); + Ok(()) + }) + } + } + + struct DependentPlugin { + log: Log, + } + + impl ClientPlugin for DependentPlugin { + type Api = String; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("dependent", "0.1.0").with_dependency("foundation") + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + let log = self.log.clone(); + Box::pin(async move { + let foundation = context + .plugin::() + .ok_or_else(|| anyhow::anyhow!("foundation API is unavailable"))?; + anyhow::ensure!(&*foundation == "foundation-api"); + record(&log, "install:dependent"); + Ok(Arc::new("dependent-api".to_string())) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "shutdown:dependent"); + Ok(()) + }) + } + } + + #[tokio::test] + async fn installs_in_dependency_order_and_indexes_by_marker_type() { + let log = Arc::new(Mutex::new(Vec::new())); + let build = complete_builder() + .await + .with_plugin(DependentPlugin { log: log.clone() }) + .with_plugin(FoundationPlugin { log: log.clone() }) + .build() + .await + .expect("valid plugin plan"); + let client = build.client(); + + assert_eq!( + client + .plugin::() + .as_deref() + .map(String::as_str), + Some("foundation-api") + ); + assert_eq!( + client + .plugin::() + .as_deref() + .map(String::as_str), + Some("dependent-api") + ); + assert_eq!( + client + .plugin_manifests() + .iter() + .map(PluginManifest::id) + .collect::>(), + vec!["foundation", "dependent"] + ); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["install:foundation", "install:dependent"] + ); + + client.disconnect().await; + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec![ + "install:foundation", + "install:dependent", + "shutdown:dependent", + "shutdown:foundation" + ] + ); + } + + struct DeclarativePlugin { + id: &'static str, + dependency: Option<&'static str>, + } + + struct TransitiveProbe; + + impl ClientPlugin for TransitiveProbe { + type Api = bool; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("transitive-probe", "0.1.0").with_dependency("dependent") + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async move { + anyhow::ensure!(context.plugin::().is_some()); + Ok(Arc::new(context.plugin::().is_none())) + }) + } + } + + #[tokio::test] + async fn install_context_exposes_only_direct_declared_dependencies() { + let log = Arc::new(Mutex::new(Vec::new())); + let build = complete_builder() + .await + .with_plugin(FoundationPlugin { log: log.clone() }) + .with_plugin(DependentPlugin { log }) + .with_plugin(TransitiveProbe) + .build() + .await + .expect("declared dependency plan"); + let client = build.client(); + assert_eq!(client.plugin::().as_deref(), Some(&true)); + client.disconnect().await; + } + + impl ClientPlugin for DeclarativePlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + let manifest = PluginManifest::new(self.id, "0.1.0"); + match self.dependency { + Some(dependency) => manifest.with_dependency(dependency), + None => manifest, + } + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async { Ok(Arc::new(())) }) + } + } + + #[test] + fn rejects_duplicate_ids_missing_dependencies_and_cycles() { + let duplicate = PluginPlan::prepare(vec![ + PluginRegistration::new(DeclarativePlugin::<1> { + id: "same", + dependency: None, + }), + PluginRegistration::new(DeclarativePlugin::<2> { + id: "same", + dependency: None, + }), + ]); + assert!(matches!( + duplicate, + Err(PluginPlanError::DuplicateId { ref id }) if id == "same" + )); + + let missing = PluginPlan::prepare(vec![PluginRegistration::new(DeclarativePlugin::<3> { + id: "orphan", + dependency: Some("absent"), + })]); + assert!(matches!( + missing, + Err(PluginPlanError::MissingDependency { + ref plugin_id, + ref dependency, + }) if plugin_id == "orphan" && dependency == "absent" + )); + + let cycle = PluginPlan::prepare(vec![ + PluginRegistration::new(DeclarativePlugin::<4> { + id: "cycle-a", + dependency: Some("cycle-b"), + }), + PluginRegistration::new(DeclarativePlugin::<5> { + id: "cycle-b", + dependency: Some("cycle-a"), + }), + ]); + assert!(matches!( + cycle, + Err(PluginPlanError::DependencyCycle { ref plugins }) + if plugins == &["cycle-a", "cycle-b"] + )); + } + + struct FixedManifestPlugin(PluginManifest); + + impl ClientPlugin for FixedManifestPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + self.0.clone() + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async { Ok(Arc::new(())) }) + } + } + + struct PanickingManifestPlugin; + + impl ClientPlugin for PanickingManifestPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + panic!("injected manifest panic") + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async { Ok(Arc::new(())) }) + } + } + + #[test] + fn rejects_invalid_or_ambiguous_manifests_without_installing() { + let duplicate_type = PluginPlan::prepare(vec![ + PluginRegistration::new(FixedManifestPlugin::<1>(PluginManifest::new( + "first", "0.1.0", + ))), + PluginRegistration::new(FixedManifestPlugin::<1>(PluginManifest::new( + "second", "0.1.0", + ))), + ]); + assert!(matches!( + duplicate_type, + Err(PluginPlanError::DuplicateType { .. }) + )); + + let invalid_id = PluginPlan::prepare(vec![PluginRegistration::new( + FixedManifestPlugin::<2>(PluginManifest::new("Invalid", "0.1.0")), + )]); + assert!(matches!(invalid_id, Err(PluginPlanError::InvalidId { .. }))); + + let invalid_version = PluginPlan::prepare(vec![PluginRegistration::new( + FixedManifestPlugin::<3>(PluginManifest::new("invalid-version", "0.1 0")), + )]); + assert!(matches!( + invalid_version, + Err(PluginPlanError::InvalidVersion { .. }) + )); + + let duplicate_dependency = PluginPlan::prepare(vec![ + PluginRegistration::new(FixedManifestPlugin::<4>(PluginManifest::new( + "base", "0.1.0", + ))), + PluginRegistration::new(FixedManifestPlugin::<5>( + PluginManifest::new("duplicate-dependency", "0.1.0") + .with_dependency("base") + .with_dependency("base"), + )), + ]); + assert!(matches!( + duplicate_dependency, + Err(PluginPlanError::DuplicateDependency { .. }) + )); + + let manifest_panic = + PluginPlan::prepare(vec![PluginRegistration::new(PanickingManifestPlugin)]); + assert!(matches!( + manifest_panic, + Err(PluginPlanError::ManifestPanicked { .. }) + )); + } + + struct DropFlag(Arc); + + impl Drop for DropFlag { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } + } + + struct RollbackPlugin { + log: Log, + task_dropped: Arc, + } + + impl ClientPlugin for RollbackPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("rollback", "0.1.0").with_capability(PluginCapability::Tasks) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + let log = self.log.clone(); + let task_dropped = self.task_dropped.clone(); + Box::pin(async move { + record(&log, "install:rollback"); + let guard = DropFlag(task_dropped); + context + .tasks() + .ok_or_else(|| anyhow::anyhow!("tasks capability missing"))? + .spawn(async move { + let _guard = guard; + futures::future::pending::<()>().await; + })?; + Ok(Arc::new(())) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "shutdown:rollback"); + Ok(()) + }) + } + } + + struct FailingPlugin { + log: Log, + } + + impl ClientPlugin for FailingPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("failing", "0.1.0").with_dependency("rollback") + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "install:failing"); + anyhow::bail!("injected failure") + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "shutdown:failing"); + Ok(()) + }) + } + } + + #[tokio::test] + async fn install_failure_rolls_back_resources_and_plugins_in_lifo_order() { + let log = Arc::new(Mutex::new(Vec::new())); + let task_dropped = Arc::new(AtomicBool::new(false)); + let result = complete_builder() + .await + .with_lifecycle(UpstreamLifecycle { log: log.clone() }) + .with_plugin(FailingPlugin { log: log.clone() }) + .with_plugin(RollbackPlugin { + log: log.clone(), + task_dropped: task_dropped.clone(), + }) + .build() + .await; + + assert!(matches!(result, Err(ClientBuilderError::PluginInstall(_)))); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec![ + "install:upstream", + "install:rollback", + "install:failing", + "shutdown:failing", + "shutdown:rollback", + "shutdown:upstream" + ] + ); + tokio::time::timeout(Duration::from_secs(1), async { + while !task_dropped.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("rollback aborted the install-scoped task"); + } + + struct PanickingPlugin { + log: Log, + } + + impl ClientPlugin for PanickingPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("panicking", "0.1.0").with_dependency("rollback") + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + record(&self.log, "install:panicking"); + panic!("injected install panic") + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "shutdown:panicking"); + Ok(()) + }) + } + } + + #[tokio::test] + async fn install_panic_isolated_and_rolled_back() { + let log = Arc::new(Mutex::new(Vec::new())); + let result = complete_builder() + .await + .with_plugin(PanickingPlugin { log: log.clone() }) + .with_plugin(RollbackPlugin { + log: log.clone(), + task_dropped: Arc::new(AtomicBool::new(false)), + }) + .build() + .await; + + assert!(matches!(result, Err(ClientBuilderError::PluginInstall(_)))); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec![ + "install:rollback", + "install:panicking", + "shutdown:panicking", + "shutdown:rollback" + ] + ); + } + + struct ScopedTaskPlugin { + install_started: Arc, + install_dropped: Arc, + connection_started: Arc, + connection_dropped: Arc, + } + + impl ClientPlugin for ScopedTaskPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("scoped-tasks", "0.1.0").with_capability(PluginCapability::Tasks) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + let started = self.install_started.clone(); + let observed = self.install_started.clone(); + let dropped = self.install_dropped.clone(); + Box::pin(async move { + context + .tasks() + .ok_or_else(|| anyhow::anyhow!("tasks capability missing"))? + .spawn(async move { + started.store(true, Ordering::Release); + let _guard = DropFlag(dropped); + futures::future::pending::<()>().await; + })?; + tokio::task::yield_now().await; + anyhow::ensure!(!observed.load(Ordering::Acquire)); + Ok(Arc::new(())) + }) + } + + fn on_ready(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + let started = self.connection_started.clone(); + let dropped = self.connection_dropped.clone(); + Box::pin(async move { + scope + .tasks() + .ok_or_else(|| anyhow::anyhow!("connection tasks capability missing"))? + .spawn(async move { + started.store(true, Ordering::Release); + let _guard = DropFlag(dropped); + futures::future::pending::<()>().await; + })?; + Ok(()) + }) + } + } + + async fn wait_for_flag(flag: &AtomicBool) { + tokio::time::timeout(Duration::from_secs(1), async { + while !flag.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("task state transition"); + } + + #[tokio::test] + async fn install_tasks_start_after_publish_and_outlive_connection_tasks() { + let install_started = Arc::new(AtomicBool::new(false)); + let install_dropped = Arc::new(AtomicBool::new(false)); + let connection_started = Arc::new(AtomicBool::new(false)); + let connection_dropped = Arc::new(AtomicBool::new(false)); + let build = complete_builder() + .await + .with_plugin(ScopedTaskPlugin { + install_started: install_started.clone(), + install_dropped: install_dropped.clone(), + connection_started: connection_started.clone(), + connection_dropped: connection_dropped.clone(), + }) + .build() + .await + .expect("scoped task plugin"); + let client = build.client(); + wait_for_flag(&install_started).await; + + let scope = ConnectionScope::new(88); + client + .plugin_host + .as_ref() + .expect("plugin host") + .on_ready(scope.clone()) + .await + .expect("plugin ready callback"); + wait_for_flag(&connection_started).await; + scope.cancel(); + wait_for_flag(&connection_dropped).await; + assert!(!install_dropped.load(Ordering::Acquire)); + + client.disconnect().await; + wait_for_flag(&install_dropped).await; + } + + #[tokio::test] + async fn connection_scoped_tasks_stop_when_the_generation_is_cancelled() { + let scope = ConnectionScope::new(77); + let task_dropped = Arc::new(AtomicBool::new(false)); + let tasks = PluginConnectionTasks { + runtime: Arc::new(TokioRuntime), + scope: scope.clone(), + }; + let guard = DropFlag(task_dropped.clone()); + tasks + .spawn(async move { + let _guard = guard; + futures::future::pending::<()>().await; + }) + .expect("open connection scope"); + + scope.cancel(); + tokio::time::timeout(Duration::from_secs(1), async { + while !task_dropped.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("connection cancellation stopped the scoped task"); + assert!(matches!( + tasks.spawn(async {}), + Err(PluginResourceError::ShuttingDown) + )); + } + + struct UpstreamLifecycle { + log: Log, + } + + impl ClientLifecycle for UpstreamLifecycle { + fn install(&self, _client: Weak) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "install:upstream"); + Ok(()) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "shutdown:upstream"); + Ok(()) + }) + } + } + + #[tokio::test] + async fn composes_existing_lifecycle_outside_plugin_lifo_order() { + let log = Arc::new(Mutex::new(Vec::new())); + let build = complete_builder() + .await + .with_lifecycle(UpstreamLifecycle { log: log.clone() }) + .with_plugin(FoundationPlugin { log: log.clone() }) + .build() + .await + .expect("composed lifecycle"); + let client = build.client(); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["install:upstream", "install:foundation"] + ); + + client.disconnect().await; + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec![ + "install:upstream", + "install:foundation", + "shutdown:foundation", + "shutdown:upstream" + ] + ); + } + + struct NoopEventHandler; + + impl EventHandler for NoopEventHandler { + fn handle_event(&self, _event: Arc) {} + } + + struct EventSubscriptionPlugin; + + impl ClientPlugin for EventSubscriptionPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("event-subscription", "0.1.0") + .with_capability(PluginCapability::CoreEvents) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async move { + context + .core_events() + .ok_or_else(|| anyhow::anyhow!("core events capability missing"))? + .subscribe( + EventInterest::of(&[wacore::types::events::EventKind::Connected]), + Arc::new(NoopEventHandler), + )?; + Ok(Arc::new(())) + }) + } + } + + #[tokio::test] + async fn shutdown_removes_plugin_event_subscriptions() { + let build = complete_builder() + .await + .with_plugin(EventSubscriptionPlugin) + .build() + .await + .expect("event subscription plugin"); + let client = build.client(); + assert!( + client + .core + .event_bus + .has_handler_for(wacore::types::events::EventKind::Connected) + ); + + client.disconnect().await; + assert!( + !client + .core + .event_bus + .has_handler_for(wacore::types::events::EventKind::Connected) + ); + } + + struct CapabilityProbe; + + impl ClientPlugin for CapabilityProbe { + type Api = [bool; 4]; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("capability-probe", "0.1.0") + .with_capability(PluginCapability::Messaging) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async move { + Ok(Arc::new([ + context.core_events().is_some(), + context.tasks().is_some(), + context.messaging().is_some(), + context.iq().is_some(), + ])) + }) + } + } + + #[tokio::test] + async fn context_exposes_only_declared_capabilities() { + let build = complete_builder() + .await + .with_plugin(CapabilityProbe) + .build() + .await + .expect("capability plugin"); + let client = build.client(); + assert_eq!( + client.plugin::().as_deref(), + Some(&[false, false, true, false]) + ); + client.disconnect().await; + } +} From 7c2e584c85b4f2c8cfc00e174edf4e1f6fbc6535 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 10:27:37 -0300 Subject: [PATCH 12/46] perf(plugins): keep native host opt-in --- Cargo.toml | 3 +++ src/bot.rs | 15 +++++++++++++-- src/client.rs | 1 + src/client/builder.rs | 39 ++++++++++++++++++++++++++++----------- src/client/lifecycle.rs | 2 ++ src/lib.rs | 3 +++ 6 files changed, 50 insertions(+), 13 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 102b0cdc4..b8b73d7f0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -112,6 +112,9 @@ zlib-rs = { version = "0.6.5", default-features = false, features = ["std", "rus [features] debug-snapshots = ["wacore/debug-snapshots"] +# Build-time native plugin host. Kept opt-in so clients that do not use plugins +# retain the pre-host binary footprint. +plugins = [] # Optional observability. Off by default: no `tracing` dep, zero overhead. # Emits tracing spans/events only; the application installs the subscriber # (and any OpenTelemetry bridge). See examples/observability.rs. diff --git a/src/bot.rs b/src/bot.rs index c6e2990e7..1d12132ce 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -1,6 +1,7 @@ use crate::cache_config::CacheConfig; use crate::client::{Client, ClientBuilderError}; use crate::pair_code::PairCodeOptions; +#[cfg(feature = "plugins")] use crate::plugins::{ClientPlugin, PluginRegistration}; use crate::store::commands::DeviceCommand; use crate::store::error::StoreError; @@ -672,6 +673,7 @@ pub struct BotBuilder< resend_rate_limit: Option<(u32, u32)>, task_instrument: Option>, alloc_meter: Option>, + #[cfg(feature = "plugins")] plugins: Vec, _marker: PhantomData<(B, T, H, R)>, } @@ -698,6 +700,7 @@ impl BotBuilder BotBuilder { resend_rate_limit: self.resend_rate_limit, task_instrument: self.task_instrument, alloc_meter: self.alloc_meter, + #[cfg(feature = "plugins")] plugins: self.plugins, _marker: PhantomData, } @@ -849,12 +853,14 @@ impl BotBuilder { } /// Register a native plugin without changing the builder's typestate. + #[cfg(feature = "plugins")] pub fn with_plugin(mut self, plugin: P) -> Self { self.plugins.push(PluginRegistration::new(plugin)); self } /// Register an already-shared native plugin without changing its marker type. + #[cfg(feature = "plugins")] pub fn with_plugin_arc(mut self, plugin: Arc

) -> Self { self.plugins.push(PluginRegistration::new_arc(plugin)); self @@ -1297,16 +1303,18 @@ impl BotBuilder { } info!("Creating client..."); - let mut client_builder = Client::builder() + let client_builder = Client::builder() .with_runtime_arc(runtime) .with_persistence_manager(persistence_manager) .with_transport_factory_arc(transport_factory) .with_http_client_arc(http_client) .with_cache_config(self.cache_config) .with_custom_enc_handlers(self.custom_enc_handlers) - .with_plugin_registrations(self.plugins) .with_skip_history_sync(self.skip_history_sync) .with_background_saver_interval(std::time::Duration::from_secs(30)); + #[cfg(feature = "plugins")] + let client_builder = client_builder.with_plugin_registrations(self.plugins); + let mut client_builder = client_builder; if let Some(version) = self.override_version { client_builder = client_builder.with_version_override(version); @@ -1402,8 +1410,10 @@ mod tests { .client() } + #[cfg(feature = "plugins")] struct BotBuilderPlugin; + #[cfg(feature = "plugins")] impl ClientPlugin for BotBuilderPlugin { type Api = &'static str; @@ -1419,6 +1429,7 @@ mod tests { } } + #[cfg(feature = "plugins")] #[tokio::test] async fn typestate_builder_preserves_registered_plugins() { let bot = Bot::builder() diff --git a/src/client.rs b/src/client.rs index 83e39b3ae..3a819d47d 100644 --- a/src/client.rs +++ b/src/client.rs @@ -690,6 +690,7 @@ pub struct Client { /// Allocated only when an extension host installs lifecycle callbacks. lifecycle: Option>, /// Allocated only when at least one build-time plugin is registered. + #[cfg(feature = "plugins")] pub(crate) plugin_host: Option>, /// Per-session wire I/O and activity counters. Written at the transport /// chokepoints (noise sender task, read loop); the keepalive dead-socket diff --git a/src/client/builder.rs b/src/client/builder.rs index a975621ff..5280e8e5a 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -7,6 +7,7 @@ use thiserror::Error; use super::{Client, ClientLifecycle, LifecycleRegistration}; use crate::cache_config::CacheConfig; use crate::http::HttpClient; +#[cfg(feature = "plugins")] use crate::plugins::{ClientPlugin, PluginHost, PluginPlan, PluginPlanError, PluginRegistration}; use crate::store::error::StoreError; use crate::store::persistence_manager::PersistenceManager; @@ -63,8 +64,10 @@ pub enum ClientBuilderError { UnsupportedDurabilityBackend(String), #[error("client lifecycle installation failed: {0}")] LifecycleInstall(#[source] anyhow::Error), + #[cfg(feature = "plugins")] #[error("plugin host installation failed: {0}")] PluginInstall(#[source] anyhow::Error), + #[cfg(feature = "plugins")] #[error("invalid plugin plan: {0}")] PluginPlan(#[from] PluginPlanError), } @@ -90,6 +93,7 @@ pub struct ClientBuilder { alloc_meter: Option>, background_saver_interval: Option, lifecycle: Option>, + #[cfg(feature = "plugins")] plugins: Vec, } @@ -118,6 +122,7 @@ impl ClientBuilder { alloc_meter: None, background_saver_interval: None, lifecycle: None, + #[cfg(feature = "plugins")] plugins: Vec::new(), } } @@ -283,17 +288,20 @@ impl ClientBuilder { } /// Register a native plugin for transactional installation before services start. + #[cfg(feature = "plugins")] pub fn with_plugin(mut self, plugin: P) -> Self { self.plugins.push(PluginRegistration::new(plugin)); self } /// Register an already-shared native plugin without changing its marker type. + #[cfg(feature = "plugins")] pub fn with_plugin_arc(mut self, plugin: Arc

) -> Self { self.plugins.push(PluginRegistration::new_arc(plugin)); self } + #[cfg(feature = "plugins")] pub(crate) fn with_plugin_registrations( mut self, registrations: Vec, @@ -370,6 +378,7 @@ impl ClientBuilder { transport_factory: Arc, http_client: Arc, ) -> Result { + #[cfg(feature = "plugins")] let plugin_plan = PluginPlan::prepare(self.plugins)?; let runtime: Arc = match self.task_instrument { Some(instrument) => { @@ -378,12 +387,17 @@ impl ClientBuilder { None => runtime, }; - let mut lifecycle_handler = self.lifecycle; - let plugin_host = plugin_plan.map(|plan| { - let host = PluginHost::new(plan, lifecycle_handler.take()); - lifecycle_handler = Some(host.clone()); - host - }); + let lifecycle_handler = self.lifecycle; + #[cfg(feature = "plugins")] + let (lifecycle_handler, plugin_host) = { + let mut lifecycle_handler = lifecycle_handler; + let plugin_host = plugin_plan.map(|plan| { + let host = PluginHost::new(plan, lifecycle_handler.take()); + lifecycle_handler = Some(host.clone()); + host + }); + (lifecycle_handler, plugin_host) + }; let lifecycle = lifecycle_handler .map(|handler| Arc::new(LifecycleRegistration::new(handler, Arc::clone(&runtime)))); let assembly = Client::assemble( @@ -395,6 +409,7 @@ impl ClientBuilder { self.cache_config, ClientExtensions { lifecycle, + #[cfg(feature = "plugins")] plugin_host, }, ); @@ -421,11 +436,11 @@ impl ClientBuilder { if let Some(lifecycle) = &client.lifecycle && let Err(error) = lifecycle.install(Arc::downgrade(&client)).await { - return Err(if client.plugin_host.is_some() { - ClientBuilderError::PluginInstall(error) - } else { - ClientBuilderError::LifecycleInstall(error) - }); + #[cfg(feature = "plugins")] + if client.plugin_host.is_some() { + return Err(ClientBuilderError::PluginInstall(error)); + } + return Err(ClientBuilderError::LifecycleInstall(error)); } let build = assembly.start(); @@ -437,6 +452,7 @@ impl ClientBuilder { ); let _ = build.client.saver_handle.set(saver_handle); } + #[cfg(feature = "plugins")] if let Some(plugin_host) = &client.plugin_host { plugin_host.activate(); } @@ -491,6 +507,7 @@ pub(super) struct ClientAssembly { #[derive(Default)] pub(super) struct ClientExtensions { pub(super) lifecycle: Option>, + #[cfg(feature = "plugins")] pub(super) plugin_host: Option>, } diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 31aa2556f..1aab79374 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -156,6 +156,7 @@ impl Client { ) -> ClientAssembly { let ClientExtensions { lifecycle, + #[cfg(feature = "plugins")] plugin_host, } = extensions; let mut unique_id_bytes = [0u8; 2]; @@ -185,6 +186,7 @@ impl Client { shutdown_notifier: wacore::runtime::ShutdownNotifier::new(), connection_shutdown: std::sync::Mutex::new(wacore::runtime::ShutdownNotifier::new()), lifecycle, + #[cfg(feature = "plugins")] plugin_host, stats: Arc::new(wacore::stats::SessionStats::new()), diff --git a/src/lib.rs b/src/lib.rs index 544d464d1..a94bb0a81 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -111,7 +111,9 @@ pub(crate) mod msg_secret_buffer; pub mod pair; pub mod pair_code; pub mod passkey; +#[cfg(feature = "plugins")] pub mod plugins; +#[cfg(feature = "plugins")] pub use plugins::{ ClientPlugin, PluginCapabilities, PluginCapability, PluginConnectionScope, PluginConnectionTasks, PluginContext, PluginCoreEvents, PluginIq, PluginIqError, @@ -190,6 +192,7 @@ pub mod prelude { Client, ClientBuilder, ClientBuilderError, ClientError, ClientLifecycle, ConnectionScope, ConnectionScopeState, RawNodeLease, }; + #[cfg(feature = "plugins")] pub use crate::plugins::{ ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, PluginManifest, }; From e41901a6a87b273babcc5e75c175cdc4095092ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:02:42 -0300 Subject: [PATCH 13/46] fix(plugins): harden host lifecycle ownership --- src/bot.rs | 41 +-- src/client/app_state.rs | 36 ++ src/client/builder.rs | 111 ++++++- src/client/extension_lifecycle.rs | 43 ++- src/client/lifecycle.rs | 2 +- src/plugins/mod.rs | 529 ++++++++++++++++++++++++++---- 6 files changed, 637 insertions(+), 125 deletions(-) diff --git a/src/bot.rs b/src/bot.rs index 1d12132ce..ac36db814 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -546,46 +546,7 @@ impl Bot { } = self; if let Some(receiver) = sync_task_receiver { - // This channel carries only HistorySync tasks: app-state sync runs via - // its own direct path (fetch_app_state_with_retry), nothing enqueues - // AppStateSync here. Chunks are independent (order-free upserts; the - // event carries chunk_order), so ingest concurrently, bounded low — each - // in-flight chunk decompresses a blob and the connect path is peak- - // memory-conscious (WA Web caps at histSyncChunk=3). Taking the permit in - // the recv loop backpressures history intake on a burst; since no - // app-state task flows here, that can't head-of-line block one. - const HISTORY_SYNC_CONCURRENCY: usize = 2; - let worker_client = Arc::downgrade(&client); - let history_permits = Arc::new(async_lock::Semaphore::new(HISTORY_SYNC_CONCURRENCY)); - client - .runtime - .spawn(Box::pin(async move { - while let Ok(task) = receiver.recv().await { - let Some(worker_client) = worker_client.upgrade() else { - break; - }; - - if matches!(task, crate::sync_task::MajorSyncTask::HistorySync { .. }) { - let permit = history_permits.acquire_arc().await; - let task_client = worker_client.clone(); - worker_client - .runtime - .spawn(Box::pin(async move { - let _permit = permit; - task_client.process_sync_task(task).await; - })) - .detach(); - } else { - // Defensive: nothing enqueues AppStateSync today, but if - // that changes it must run serially (ordered patches). - worker_client.process_sync_task(task).await; - } - } - info!( - "Sync worker intake loop finished (detached history-sync tasks may still be running)." - ); - })) - .detach(); + client.start_sync_task_worker(receiver); } if !event_handlers.is_empty() { diff --git a/src/client/app_state.rs b/src/client/app_state.rs index d7137a32a..f768f53a8 100644 --- a/src/client/app_state.rs +++ b/src/client/app_state.rs @@ -276,6 +276,42 @@ impl Client { pre_downloaded } + pub(crate) fn start_sync_task_worker( + self: &Arc, + receiver: async_channel::Receiver, + ) { + const HISTORY_SYNC_CONCURRENCY: usize = 2; + + let worker_client = Arc::downgrade(self); + let history_permits = Arc::new(async_lock::Semaphore::new(HISTORY_SYNC_CONCURRENCY)); + self.runtime + .spawn(Box::pin(async move { + while let Ok(task) = receiver.recv().await { + let Some(worker_client) = worker_client.upgrade() else { + break; + }; + + if matches!(task, crate::sync_task::MajorSyncTask::HistorySync { .. }) { + let permit = history_permits.acquire_arc().await; + let task_client = worker_client.clone(); + worker_client + .runtime + .spawn(Box::pin(async move { + let _permit = permit; + task_client.process_sync_task(task).await; + })) + .detach(); + } else { + worker_client.process_sync_task(task).await; + } + } + info!( + "Sync worker intake loop finished (detached history-sync tasks may still be running)." + ); + })) + .detach(); + } + /// Public entry point for processing [`MajorSyncTask`] from the sync channel. #[cfg_attr( feature = "tracing", diff --git a/src/client/builder.rs b/src/client/builder.rs index 5280e8e5a..69c7528fb 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -19,8 +19,8 @@ use wacore::runtime::Runtime; /// Result of constructing a [`Client`]. /// -/// The sync-task receiver has a single consumer and is therefore transferred -/// together with the client rather than hidden behind a cloneable handle. +/// Consume with [`ClientBuild::into_client`] for the standard worker or +/// [`ClientBuild::into_parts`] when the host owns that worker itself. pub struct ClientBuild { client: Arc, sync_task_receiver: async_channel::Receiver, @@ -37,12 +37,15 @@ impl ClientBuild { } } - /// Return the constructed client. - pub fn client(&self) -> Arc { - Arc::clone(&self.client) + /// Transfer the client and start its default major-sync worker. + pub fn into_client(self) -> Arc { + let (client, sync_task_receiver) = self.into_parts(); + client.start_sync_task_worker(sync_task_receiver); + client } - /// Transfer ownership of the client and its sync-task receiver. + /// Transfer ownership of the client and its sole sync-task receiver. + /// The caller must drain the receiver for history sync to keep working. pub fn into_parts(self) -> (Arc, async_channel::Receiver) { (self.client, self.sync_task_receiver) } @@ -60,6 +63,8 @@ pub enum ClientBuilderError { MissingTransportFactory, #[error("missing HTTP client")] MissingHttpClient, + #[error("background saver interval must be greater than zero")] + InvalidBackgroundSaverInterval, #[error("the configured backend does not support the inbound durability hook: {0}")] UnsupportedDurabilityBackend(String), #[error("client lifecycle installation failed: {0}")] @@ -341,6 +346,10 @@ impl ClientBuilder { .cloned() .ok_or(ClientBuilderError::MissingHttpClient)?; + if self.background_saver_interval == Some(Duration::ZERO) { + return Err(ClientBuilderError::InvalidBackgroundSaverInterval); + } + if self.inbound_durability_hook.is_some() { probe_durability_backend(&persistence_manager.backend()).await?; } @@ -398,8 +407,17 @@ impl ClientBuilder { }); (lifecycle_handler, plugin_host) }; - let lifecycle = lifecycle_handler - .map(|handler| Arc::new(LifecycleRegistration::new(handler, Arc::clone(&runtime)))); + let lifecycle = lifecycle_handler.map(|handler| { + #[cfg(feature = "plugins")] + if let Some(plugin_host) = &plugin_host { + return Arc::new(LifecycleRegistration::new_with_timeout( + handler, + Arc::clone(&runtime), + plugin_host.lifecycle_callback_timeout(), + )); + } + Arc::new(LifecycleRegistration::new(handler, Arc::clone(&runtime))) + }); let assembly = Client::assemble( Arc::clone(&runtime), Arc::clone(&persistence_manager), @@ -593,6 +611,19 @@ mod tests { } } + async fn complete_builder() -> ClientBuilder { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + ClientBuilder::new() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + } + #[tokio::test] async fn validates_required_dependencies_before_assembly() { assert!(matches!( @@ -658,7 +689,7 @@ mod tests { let build = assembly.start(); assert_eq!(spawns.load(Ordering::SeqCst), 2); - build.client().signal_shutdown_sync(); + build.into_client().signal_shutdown_sync(); } #[tokio::test] @@ -685,7 +716,7 @@ mod tests { .build() .await .expect("complete builder"); - let client = build.client(); + let client = build.into_client(); assert!(client.skip_history_sync_enabled()); assert_eq!(client.wanted_pre_key_count(), 123); @@ -696,10 +727,68 @@ mod tests { .is_some_and(|installed| Arc::ptr_eq(installed, &meter)) ); assert!(client.saver_handle.get().is_some()); - assert_eq!(spawns.load(Ordering::SeqCst), 3); + assert_eq!(spawns.load(Ordering::SeqCst), 4); + client.signal_shutdown_sync(); + } + + #[tokio::test] + async fn consuming_build_as_client_keeps_major_sync_worker_alive() { + let client = complete_builder() + .await + .build() + .await + .expect("complete builder") + .into_client(); + + assert!(!client.major_sync_task_sender.is_closed()); client.signal_shutdown_sync(); } + #[tokio::test] + async fn rejects_zero_background_saver_interval() { + let result = complete_builder() + .await + .with_background_saver_interval(Duration::ZERO) + .build() + .await; + + assert!(matches!( + result, + Err(ClientBuilderError::InvalidBackgroundSaverInterval) + )); + } + + struct PanickingInstallLifecycle { + when_polled: bool, + } + + impl ClientLifecycle for PanickingInstallLifecycle { + fn install( + &self, + _client: std::sync::Weak, + ) -> wacore::runtime::BoxFuture<'_, anyhow::Result<()>> { + if !self.when_polled { + panic!("injected synchronous install panic"); + } + Box::pin(async { panic!("injected asynchronous install panic") }) + } + } + + #[tokio::test] + async fn lifecycle_install_panics_are_typed_build_errors() { + for when_polled in [false, true] { + let result = complete_builder() + .await + .with_lifecycle(PanickingInstallLifecycle { when_polled }) + .build() + .await; + assert!(matches!( + result, + Err(ClientBuilderError::LifecycleInstall(_)) + )); + } + } + #[tokio::test] async fn lifecycle_install_failure_publishes_nothing_and_starts_no_tasks() { let persistence_manager = Arc::new( diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 35e4e31ac..a5286437c 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -131,6 +131,8 @@ impl fmt::Debug for ConnectionScope { /// stalled extension cannot block reconnect. A future plugin host owns /// per-plugin ordering and isolation behind this client-level seam. `install` /// receives a weak client reference so retaining it cannot create a cycle. +/// `signal_shutdown` is the non-blocking boundary for resources that must stop +/// even when an FFI host cannot await `shutdown`. pub trait ClientLifecycle: wacore::sync_marker::MaybeSendSync { fn install<'a>(&'a self, _client: Weak) -> BoxFuture<'a, anyhow::Result<()>> { Box::pin(async { Ok(()) }) @@ -144,6 +146,10 @@ pub trait ClientLifecycle: wacore::sync_marker::MaybeSendSync { Box::pin(async { Ok(()) }) } + /// Stop synchronously owned resources before asynchronous shutdown begins. + /// Implementations must return promptly and make repeated calls harmless. + fn signal_shutdown(&self) {} + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { Box::pin(async { Ok(()) }) } @@ -209,7 +215,7 @@ impl LifecycleRegistration { Self::new_with_timeout(handler, runtime, CALLBACK_TIMEOUT) } - fn new_with_timeout( + pub(super) fn new_with_timeout( handler: Arc, runtime: Arc, callback_timeout: Duration, @@ -227,7 +233,14 @@ impl LifecycleRegistration { } pub(super) async fn install(&self, client: Weak) -> anyhow::Result<()> { - self.handler.install(client).await + let install = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.handler.install(client) + })) + .map_err(|_| anyhow::anyhow!("lifecycle install panicked before returning a future"))?; + std::panic::AssertUnwindSafe(install) + .catch_unwind() + .await + .map_err(|_| anyhow::anyhow!("lifecycle install future panicked"))? } pub(super) fn begin_scope_if_current( @@ -352,8 +365,7 @@ impl LifecycleRegistration { } pub(super) async fn shutdown(self: &Arc) { - self.terminal.store(true, Ordering::Release); - self.cancel_active_scope(); + self.signal_shutdown_sync(); { let mut queue = self.callback_queue(); queue.shutdown_requested = true; @@ -371,6 +383,19 @@ impl LifecycleRegistration { wacore::runtime::wait_for_shutdown(&completed).await; } + pub(super) fn signal_shutdown_sync(&self) { + let first_signal = !self.terminal.swap(true, Ordering::AcqRel); + self.cancel_active_scope(); + if first_signal + && std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.handler.signal_shutdown(); + })) + .is_err() + { + log::warn!("Client lifecycle synchronous shutdown signal panicked"); + } + } + fn enqueue_callback(self: &Arc, callback: LifecycleCallback) { let should_spawn = { let mut queue = self.callback_queue(); @@ -806,7 +831,7 @@ mod tests { .build() .await .expect("client build"); - let client = build.client(); + let client = build.into_client(); const GENERATION: u64 = 9; client .connection_generation @@ -884,7 +909,7 @@ mod tests { .build() .await .expect("client build") - .client(); + .into_client(); const GENERATION: u64 = 13; client .connection_generation @@ -959,7 +984,7 @@ mod tests { .build() .await .expect("client build") - .client(); + .into_client(); const GENERATION: u64 = 17; client .connection_generation @@ -1152,7 +1177,7 @@ mod tests { .build() .await .expect("client build") - .client(); + .into_client(); const GENERATION: u64 = 37; client .connection_generation @@ -1306,7 +1331,7 @@ mod tests { .build() .await .expect("client build") - .client(); + .into_client(); client .subscribe_handler(Arc::new(LogoutOrderHandler { lifecycle: lifecycle.clone(), diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 1aab79374..0dbc3c640 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -31,7 +31,7 @@ impl Client { self.is_running.store(false, Ordering::Relaxed); self.shutdown_notifier.notify(); if let Some(lifecycle) = &self.lifecycle { - lifecycle.cancel_active_scope(); + lifecycle.signal_shutdown_sync(); } self.notify_connection_shutdown(); } diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 83fdf9eae..da0c12ecb 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -6,20 +6,22 @@ use std::future::Future; use std::panic::AssertUnwindSafe; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, OnceLock, Weak}; +use std::time::Duration; use futures::FutureExt; use thiserror::Error; use wacore::iq::spec::IqSpec; use wacore::runtime::{ - BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, Spawnable, wait_for_shutdown, + BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, Spawnable, timeout as runtime_timeout, + wait_for_shutdown, }; use wacore::sync_marker::MaybeSendSync; -use wacore::types::events::{EventHandler, EventInterest, Subscription}; +use wacore::types::events::{EventHandler, EventInterest, EventKind, Subscription}; use wacore_binary::Jid; use waproto::whatsapp::Message; use crate::Client; -use crate::client::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; +use crate::client::{ClientLifecycle, ConnectionScope, ConnectionScopeState, RawNodeLease}; use crate::request::IqError; use crate::send::{SendError, SendResult}; @@ -27,6 +29,7 @@ const CAP_CORE_EVENTS: u8 = 1 << 0; const CAP_TASKS: u8 = 1 << 1; const CAP_MESSAGING: u8 = 1 << 2; const CAP_IQ: u8 = 1 << 3; +const PLUGIN_CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); /// A capability a plugin asks the host to expose during installation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -207,7 +210,12 @@ struct PluginResources { closed: AtomicBool, activation: ShutdownNotifier, shutdown: ShutdownNotifier, - subscriptions: Mutex>, + subscriptions: Mutex>, +} + +struct PluginCoreEventSubscription { + _subscription: Subscription, + _raw_node_lease: Option, } impl PluginResources { @@ -239,7 +247,11 @@ impl PluginResources { } } - fn retain_subscription(&self, subscription: Subscription) -> Result<(), PluginResourceError> { + fn retain_subscription( + &self, + subscription: Subscription, + raw_node_lease: Option, + ) -> Result<(), PluginResourceError> { let mut subscriptions = self .subscriptions .lock() @@ -248,7 +260,10 @@ impl PluginResources { drop(subscription); return Err(PluginResourceError::ShuttingDown); } - subscriptions.push(subscription); + subscriptions.push(PluginCoreEventSubscription { + _subscription: subscription, + _raw_node_lease: raw_node_lease, + }); Ok(()) } @@ -307,8 +322,12 @@ impl PluginCoreEvents { .client .upgrade() .ok_or(PluginResourceError::ClientUnavailable)?; + let raw_node_lease = interest + .wants(EventKind::RawNode) + .then(|| client.acquire_raw_node_forwarding()); let subscription = client.subscribe(interest, handler); - self.resources.retain_subscription(subscription) + self.resources + .retain_subscription(subscription, raw_node_lease) } } @@ -760,6 +779,79 @@ struct InstalledPlugin { resources: Arc, } +struct PluginInstallRollback { + runtime: Arc, + installed: Vec, + current: Option, + upstream: Option>, + armed: bool, +} + +impl PluginInstallRollback { + fn new(runtime: Arc, capacity: usize) -> Self { + Self { + runtime, + installed: Vec::with_capacity(capacity), + current: None, + upstream: None, + armed: true, + } + } + + fn close_resources(&self) { + if let Some(current) = &self.current { + current.resources.close(); + } + for plugin in self.installed.iter().rev() { + plugin.resources.close(); + } + } + + async fn rollback(&mut self) { + self.close_resources(); + let current = self.current.take(); + let installed = std::mem::take(&mut self.installed); + let upstream = self.upstream.take(); + self.armed = false; + shutdown_staged_plugins(self.runtime.clone(), current, installed, upstream).await; + } + + fn take_installed(&mut self) -> Vec { + std::mem::take(&mut self.installed) + } + + fn restore_installed(&mut self, installed: Vec) { + self.installed = installed; + } + + fn disarm(&mut self) { + self.armed = false; + self.upstream = None; + } +} + +impl Drop for PluginInstallRollback { + fn drop(&mut self) { + if !self.armed { + return; + } + self.close_resources(); + let current = self.current.take(); + let installed = std::mem::take(&mut self.installed); + let upstream = self.upstream.take(); + if current.is_none() && installed.is_empty() && upstream.is_none() { + return; + } + let runtime = self.runtime.clone(); + let cleanup_runtime = runtime.clone(); + runtime + .spawn(Box::pin(async move { + shutdown_staged_plugins(cleanup_runtime, current, installed, upstream).await; + })) + .detach(); + } +} + pub(crate) struct PluginHost { ordered: Vec, manifests: Vec, @@ -767,10 +859,19 @@ pub(crate) struct PluginHost { installed: OnceLock>, apis: OnceLock>, runtime: OnceLock>, + callback_timeout: Duration, } impl PluginHost { pub(crate) fn new(plan: PluginPlan, upstream: Option>) -> Arc { + Self::new_with_callback_timeout(plan, upstream, PLUGIN_CALLBACK_TIMEOUT) + } + + fn new_with_callback_timeout( + plan: PluginPlan, + upstream: Option>, + callback_timeout: Duration, + ) -> Arc { let manifests = plan .ordered .iter() @@ -783,6 +884,7 @@ impl PluginHost { installed: OnceLock::new(), apis: OnceLock::new(), runtime: OnceLock::new(), + callback_timeout, }) } @@ -794,6 +896,13 @@ impl PluginHost { &self.manifests } + pub(crate) fn lifecycle_callback_timeout(&self) -> Duration { + let callback_count = self.ordered.len() + usize::from(self.upstream.is_some()); + self.callback_timeout + .saturating_mul(callback_count as u32) + .saturating_add(Duration::from_secs(1)) + } + fn context( &self, client: &Weak, @@ -862,12 +971,16 @@ impl PluginHost { .set(runtime.clone()) .map_err(|_| anyhow::anyhow!("plugin host was installed more than once"))?; + let mut rollback = PluginInstallRollback::new(runtime.clone(), self.ordered.len()); if let Some(upstream) = &self.upstream { - plugin_callback(|| upstream.install(client.clone())).await?; + if let Err(error) = plugin_callback(|| upstream.install(client.clone())).await { + rollback.disarm(); + return Err(error); + } + rollback.upstream = Some(upstream.clone()); } let staging = Arc::new(ApiRegistry::default()); - let mut installed: Vec = Vec::with_capacity(self.ordered.len()); for planned in &self.ordered { let resources = PluginResources::new(); let context = self.context( @@ -878,21 +991,15 @@ impl PluginHost { Arc::clone(&staging), runtime.clone(), ); + rollback.current = Some(InstalledPlugin { + plugin: planned.plugin.clone(), + manifest: planned.manifest.clone(), + resources, + }); let api = match plugin_install(|| planned.plugin.install(context)).await { Ok(api) => api, Err(error) => { - resources.close(); - for plugin in installed.iter().rev() { - plugin.resources.close(); - } - if let Err(shutdown_error) = plugin_callback(|| planned.plugin.shutdown()).await - { - log::warn!( - "Plugin `{}` failed-install rollback failed: {shutdown_error:#}", - planned.manifest.id - ); - } - self.rollback(&mut installed).await; + rollback.rollback().await; anyhow::bail!( "plugin `{}` installation failed: {error:#}", planned.manifest.id @@ -900,19 +1007,22 @@ impl PluginHost { } }; staging.insert(planned.plugin.marker_type_id(), api); - installed.push(InstalledPlugin { - plugin: planned.plugin.clone(), - manifest: planned.manifest.clone(), - resources, - }); + let Some(installed) = rollback.current.take() else { + rollback.rollback().await; + anyhow::bail!("plugin installation rollback state was lost"); + }; + rollback.installed.push(installed); } self.apis .set(staging.snapshot()) .map_err(|_| anyhow::anyhow!("plugin APIs were published more than once"))?; - self.installed - .set(installed) - .map_err(|_| anyhow::anyhow!("plugins were installed more than once"))?; + let installed = rollback.take_installed(); + if let Err(installed) = self.installed.set(installed) { + rollback.restore_installed(installed); + anyhow::bail!("plugins were installed more than once"); + } + rollback.disarm(); Ok(()) } @@ -922,20 +1032,15 @@ impl PluginHost { } } - async fn rollback(&self, installed: &mut Vec) { - for plugin in installed.iter().rev() { - plugin.resources.close(); - } - while let Some(plugin) = installed.pop() { - if plugin_callback(|| plugin.plugin.shutdown()).await.is_err() { - log::warn!("Plugin `{}` rollback failed", plugin.manifest.id); - } - } - if let Some(upstream) = &self.upstream - && let Err(error) = plugin_callback(|| upstream.shutdown()).await - { - log::warn!("Upstream lifecycle rollback failed: {error:#}"); - } + async fn run_callback<'a>( + &'a self, + make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result<()>>, + ) -> anyhow::Result<()> { + let runtime = self + .runtime + .get() + .ok_or_else(|| anyhow::anyhow!("plugin runtime is unavailable"))?; + bounded_plugin_callback(&**runtime, self.callback_timeout, make_future).await } } @@ -946,13 +1051,18 @@ impl ClientLifecycle for PluginHost { fn on_ready(&self, scope: ConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { Box::pin(async move { - if let Some(upstream) = &self.upstream { - plugin_callback(|| upstream.on_ready(scope.clone())).await?; - } let mut failures = Vec::new(); + if let Some(upstream) = &self.upstream + && let Err(error) = self.run_callback(|| upstream.on_ready(scope.clone())).await + { + failures.push(format!("upstream: {error:#}")); + } for plugin in self.installed.get().into_iter().flatten() { let plugin_scope = self.connection_scope(scope.clone(), &plugin.manifest); - if let Err(error) = plugin_callback(|| plugin.plugin.on_ready(plugin_scope)).await { + if let Err(error) = self + .run_callback(|| plugin.plugin.on_ready(plugin_scope)) + .await + { failures.push(format!("{}: {error:#}", plugin.manifest.id)); } } @@ -965,13 +1075,15 @@ impl ClientLifecycle for PluginHost { let mut failures = Vec::new(); for plugin in self.installed.get().into_iter().flatten().rev() { let plugin_scope = self.connection_scope(scope.clone(), &plugin.manifest); - if let Err(error) = plugin_callback(|| plugin.plugin.on_closed(plugin_scope)).await + if let Err(error) = self + .run_callback(|| plugin.plugin.on_closed(plugin_scope)) + .await { failures.push(format!("{}: {error:#}", plugin.manifest.id)); } } if let Some(upstream) = &self.upstream - && let Err(error) = plugin_callback(|| upstream.on_closed(scope)).await + && let Err(error) = self.run_callback(|| upstream.on_closed(scope)).await { failures.push(format!("upstream: {error:#}")); } @@ -979,19 +1091,28 @@ impl ClientLifecycle for PluginHost { }) } + fn signal_shutdown(&self) { + for plugin in self.installed.get().into_iter().flatten().rev() { + plugin.resources.close(); + } + if let Some(upstream) = &self.upstream + && std::panic::catch_unwind(AssertUnwindSafe(|| upstream.signal_shutdown())).is_err() + { + log::warn!("Upstream lifecycle synchronous shutdown signal panicked"); + } + } + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { Box::pin(async move { let mut failures = Vec::new(); + self.signal_shutdown(); for plugin in self.installed.get().into_iter().flatten().rev() { - plugin.resources.close(); - } - for plugin in self.installed.get().into_iter().flatten().rev() { - if let Err(error) = plugin_callback(|| plugin.plugin.shutdown()).await { + if let Err(error) = self.run_callback(|| plugin.plugin.shutdown()).await { failures.push(format!("{}: {error:#}", plugin.manifest.id)); } } if let Some(upstream) = &self.upstream - && let Err(error) = plugin_callback(|| upstream.shutdown()).await + && let Err(error) = self.run_callback(|| upstream.shutdown()).await { failures.push(format!("upstream: {error:#}")); } @@ -1002,10 +1123,57 @@ impl ClientLifecycle for PluginHost { impl Drop for PluginHost { fn drop(&mut self) { - for plugin in self.installed.get().into_iter().flatten() { - plugin.resources.close(); + self.signal_shutdown(); + } +} + +async fn shutdown_staged_plugins( + runtime: Arc, + current: Option, + mut installed: Vec, + upstream: Option>, +) { + if let Some(plugin) = current + && let Err(error) = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { + plugin.plugin.shutdown() + }) + .await + { + log::warn!( + "Plugin `{}` failed-install rollback failed: {error:#}", + plugin.manifest.id + ); + } + while let Some(plugin) = installed.pop() { + if let Err(error) = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { + plugin.plugin.shutdown() + }) + .await + { + log::warn!("Plugin `{}` rollback failed: {error:#}", plugin.manifest.id); } } + if let Some(upstream) = upstream + && let Err(error) = + bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || upstream.shutdown()) + .await + { + log::warn!("Upstream lifecycle rollback failed: {error:#}"); + } +} + +async fn bounded_plugin_callback<'a>( + runtime: &dyn Runtime, + timeout: Duration, + make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result<()>>, +) -> anyhow::Result<()> { + match runtime_timeout(runtime, timeout, plugin_callback(make_future)).await { + Ok(result) => result, + Err(_) => anyhow::bail!( + "callback timed out after {:.3} seconds", + timeout.as_secs_f64() + ), + } } async fn plugin_callback<'a>( @@ -1159,7 +1327,7 @@ mod tests { .build() .await .expect("valid plugin plan"); - let client = build.client(); + let client = build.into_client(); assert_eq!( client @@ -1233,7 +1401,7 @@ mod tests { .build() .await .expect("declared dependency plan"); - let client = build.client(); + let client = build.into_client(); assert_eq!(client.plugin::().as_deref(), Some(&true)); client.disconnect().await; } @@ -1501,6 +1669,93 @@ mod tests { .expect("rollback aborted the install-scoped task"); } + struct BlockingInstallPlugin { + log: Log, + started: async_channel::Sender<()>, + release: async_channel::Receiver<()>, + } + + impl ClientPlugin for BlockingInstallPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("blocking-install", "0.1.0").with_dependency("rollback") + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + let log = self.log.clone(); + let started = self.started.clone(); + let release = self.release.clone(); + Box::pin(async move { + record(&log, "install:blocking"); + let _ = started.try_send(()); + release.recv().await?; + Ok(Arc::new(())) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "shutdown:blocking"); + Ok(()) + }) + } + } + + #[tokio::test] + async fn cancelled_build_closes_resources_and_schedules_lifo_rollback() { + let log = Arc::new(Mutex::new(Vec::new())); + let task_dropped = Arc::new(AtomicBool::new(false)); + let (started_tx, started_rx) = async_channel::bounded(1); + let (_release_tx, release_rx) = async_channel::bounded(1); + let builder = complete_builder() + .await + .with_plugin(BlockingInstallPlugin { + log: log.clone(), + started: started_tx, + release: release_rx, + }) + .with_plugin(RollbackPlugin { + log: log.clone(), + task_dropped: task_dropped.clone(), + }); + + let build = tokio::spawn(async move { builder.build().await }); + started_rx.recv().await.expect("blocking install started"); + build.abort(); + let _ = build.await; + + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let complete = log + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .last() + .is_some_and(|entry| entry == "shutdown:rollback"); + if complete && task_dropped.load(Ordering::Acquire) { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled build rollback completed"); + + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec![ + "install:rollback", + "install:blocking", + "shutdown:blocking", + "shutdown:rollback" + ] + ); + } + struct PanickingPlugin { log: Log, } @@ -1631,7 +1886,7 @@ mod tests { .build() .await .expect("scoped task plugin"); - let client = build.client(); + let client = build.into_client(); wait_for_flag(&install_started).await; let scope = ConnectionScope::new(88); @@ -1703,6 +1958,121 @@ mod tests { } } + struct FailingReadyLifecycle; + + impl ClientLifecycle for FailingReadyLifecycle { + fn on_ready(&self, _scope: ConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async { anyhow::bail!("injected upstream ready failure") }) + } + } + + struct ReadyPlugin { + id: &'static str, + dependency: Option<&'static str>, + log: Log, + stalls: bool, + } + + impl ClientPlugin for ReadyPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + let manifest = PluginManifest::new(self.id, "0.1.0"); + match self.dependency { + Some(dependency) => manifest.with_dependency(dependency), + None => manifest, + } + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async { Ok(Arc::new(())) }) + } + + fn on_ready(&self, _scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + let id = self.id; + let log = self.log.clone(); + let stalls = self.stalls; + Box::pin(async move { + record(&log, format!("ready:{id}")); + if stalls { + futures::future::pending::<()>().await; + } + Ok(()) + }) + } + } + + #[tokio::test] + async fn upstream_ready_failure_does_not_suppress_plugins() { + let log = Arc::new(Mutex::new(Vec::new())); + let client = complete_builder() + .await + .with_lifecycle(FailingReadyLifecycle) + .with_plugin(ReadyPlugin::<1> { + id: "ready-probe", + dependency: None, + log: log.clone(), + stalls: false, + }) + .build() + .await + .expect("ready probe client") + .into_client(); + + let result = client + .plugin_host + .as_ref() + .expect("plugin host") + .on_ready(ConnectionScope::new(91)) + .await; + assert!(result.is_err()); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["ready:ready-probe"] + ); + client.disconnect().await; + } + + #[tokio::test] + async fn timed_out_plugin_callback_does_not_suppress_following_plugins() { + let log = Arc::new(Mutex::new(Vec::new())); + let plan = PluginPlan::prepare(vec![ + PluginRegistration::new(ReadyPlugin::<2> { + id: "stalling-ready", + dependency: None, + log: log.clone(), + stalls: true, + }), + PluginRegistration::new(ReadyPlugin::<3> { + id: "following-ready", + dependency: Some("stalling-ready"), + log: log.clone(), + stalls: false, + }), + ]) + .expect("valid callback plan") + .expect("non-empty callback plan"); + let host = PluginHost::new_with_callback_timeout(plan, None, Duration::from_millis(10)); + let client = complete_builder() + .await + .with_lifecycle_arc(host.clone()) + .build() + .await + .expect("callback timeout client") + .into_client(); + + let result = host.on_ready(ConnectionScope::new(92)).await; + assert!(result.is_err()); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["ready:stalling-ready", "ready:following-ready"] + ); + client.disconnect().await; + } + #[tokio::test] async fn composes_existing_lifecycle_outside_plugin_lifo_order() { let log = Arc::new(Mutex::new(Vec::new())); @@ -1713,7 +2083,7 @@ mod tests { .build() .await .expect("composed lifecycle"); - let client = build.client(); + let client = build.into_client(); assert_eq!( *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), vec!["install:upstream", "install:foundation"] @@ -1753,7 +2123,7 @@ mod tests { .core_events() .ok_or_else(|| anyhow::anyhow!("core events capability missing"))? .subscribe( - EventInterest::of(&[wacore::types::events::EventKind::Connected]), + EventInterest::of(&[EventKind::Connected, EventKind::RawNode]), Arc::new(NoopEventHandler), )?; Ok(Arc::new(())) @@ -1762,20 +2132,21 @@ mod tests { } #[tokio::test] - async fn shutdown_removes_plugin_event_subscriptions() { + async fn shutdown_removes_plugin_event_subscriptions_and_raw_lease() { let build = complete_builder() .await .with_plugin(EventSubscriptionPlugin) .build() .await .expect("event subscription plugin"); - let client = build.client(); + let client = build.into_client(); assert!( client .core .event_bus .has_handler_for(wacore::types::events::EventKind::Connected) ); + assert!(client.raw_node_forwarding_enabled()); client.disconnect().await; assert!( @@ -1784,6 +2155,36 @@ mod tests { .event_bus .has_handler_for(wacore::types::events::EventKind::Connected) ); + assert!(!client.raw_node_forwarding_enabled()); + } + + #[tokio::test] + async fn synchronous_shutdown_closes_plugin_resources_with_live_client_refs() { + let task_dropped = Arc::new(AtomicBool::new(false)); + let client = complete_builder() + .await + .with_plugin(RollbackPlugin { + log: Arc::new(Mutex::new(Vec::new())), + task_dropped: task_dropped.clone(), + }) + .with_plugin(EventSubscriptionPlugin) + .build() + .await + .expect("plugin resource client") + .into_client(); + let retained_client = client.clone(); + + client.signal_shutdown_sync(); + wait_for_flag(&task_dropped).await; + + assert!(!retained_client.raw_node_forwarding_enabled()); + assert!( + !retained_client + .core + .event_bus + .has_handler_for(EventKind::Connected) + ); + retained_client.disconnect().await; } struct CapabilityProbe; @@ -1816,7 +2217,7 @@ mod tests { .build() .await .expect("capability plugin"); - let client = build.client(); + let client = build.into_client(); assert_eq!( client.plugin::().as_deref(), Some(&[false, false, true, false]) From 9e7aba15c2bc729cc3b01ed70e61446c3d7bbb22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:37:38 -0300 Subject: [PATCH 14/46] fix(plugins): make terminal cleanup cancellation-safe --- src/client/extension_lifecycle.rs | 82 ++++++++++++- src/client/lifecycle.rs | 7 ++ src/plugins/mod.rs | 188 +++++++++++++++++++++++++++--- 3 files changed, 252 insertions(+), 25 deletions(-) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index a5286437c..15b3415d9 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -365,12 +365,7 @@ impl LifecycleRegistration { } pub(super) async fn shutdown(self: &Arc) { - self.signal_shutdown_sync(); - { - let mut queue = self.callback_queue(); - queue.shutdown_requested = true; - } - self.enqueue_shutdown_if_ready(); + self.request_shutdown(); if self.shutdown_complete.load(Ordering::Acquire) || callback_context_active(self) { return; @@ -383,6 +378,15 @@ impl LifecycleRegistration { wacore::runtime::wait_for_shutdown(&completed).await; } + pub(super) fn request_shutdown(self: &Arc) { + self.signal_shutdown_sync(); + { + let mut queue = self.callback_queue(); + queue.shutdown_requested = true; + } + self.enqueue_shutdown_if_ready(); + } + pub(super) fn signal_shutdown_sync(&self) { let first_signal = !self.terminal.swap(true, Ordering::AcqRel); self.cancel_active_scope(); @@ -722,6 +726,26 @@ mod tests { } } + #[derive(Default)] + struct EarlyShutdownLifecycle { + signalled: AtomicBool, + shutdowns: AtomicUsize, + } + + impl ClientLifecycle for EarlyShutdownLifecycle { + fn signal_shutdown(&self) { + self.signalled.store(true, Ordering::Release); + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + assert!(self.signalled.load(Ordering::Acquire)); + self.shutdowns.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + } + } + #[derive(Default)] struct SynchronousPanicLifecycle { ready_calls: AtomicUsize, @@ -885,6 +909,52 @@ mod tests { client.signal_shutdown_sync(); } + #[tokio::test] + async fn disconnect_requests_lifecycle_shutdown_before_cancellable_io() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let lifecycle = Arc::new(EarlyShutdownLifecycle::default()); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await + .expect("client build") + .into_client(); + let (started_tx, started_rx) = async_channel::bounded(1); + let (_release_tx, release_rx) = async_channel::bounded(1); + *client.transport.lock().await = Some(Arc::new(BlockingDisconnect { + started: started_tx, + release: release_rx, + })); + + let disconnect_client = Arc::clone(&client); + let disconnect = tokio::spawn(async move { + disconnect_client.disconnect().await; + }); + tokio::time::timeout(std::time::Duration::from_secs(2), started_rx.recv()) + .await + .expect("disconnect reached cancellable transport I/O") + .expect("transport remained alive"); + assert!(lifecycle.signalled.load(Ordering::Acquire)); + + disconnect.abort(); + let _ = disconnect.await; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while lifecycle.shutdowns.load(Ordering::SeqCst) != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("detached lifecycle shutdown completed"); + } + #[tokio::test] async fn cancellation_does_not_wait_for_a_running_callback() { let persistence_manager = Arc::new( diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 0dbc3c640..83374add2 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -101,6 +101,12 @@ impl Client { } } + fn request_lifecycle_shutdown(&self) { + if let Some(lifecycle) = &self.lifecycle { + lifecycle.request_shutdown(); + } + } + /// Create a new `Client` with default cache configuration. /// /// This is the standard constructor. Use [`Client::new_with_cache_config`] @@ -684,6 +690,7 @@ impl Client { self.expected_disconnect.store(true, Ordering::Relaxed); self.is_running.store(false, Ordering::Relaxed); self.shutdown_notifier.notify(); + self.request_lifecycle_shutdown(); // Drain buffered offline receipts into the flush window before // closing it, so a disconnect mid-offline-sync still acks the diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index da0c12ecb..b53e71257 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -807,13 +807,41 @@ impl PluginInstallRollback { } } - async fn rollback(&mut self) { + fn schedule_rollback(&mut self) -> Option { + if !self.armed { + return None; + } self.close_resources(); let current = self.current.take(); let installed = std::mem::take(&mut self.installed); let upstream = self.upstream.take(); self.armed = false; - shutdown_staged_plugins(self.runtime.clone(), current, installed, upstream).await; + if current.is_none() && installed.is_empty() && upstream.is_none() { + return None; + } + if let Some(upstream) = &upstream + && std::panic::catch_unwind(AssertUnwindSafe(|| upstream.signal_shutdown())).is_err() + { + log::warn!("Upstream lifecycle rollback shutdown signal panicked"); + } + + let completed = ShutdownNotifier::new(); + let completion = completed.subscribe(); + let runtime = self.runtime.clone(); + let cleanup_runtime = runtime.clone(); + runtime + .spawn(Box::pin(async move { + shutdown_staged_plugins(cleanup_runtime, current, installed, upstream).await; + completed.notify(); + })) + .detach(); + Some(completion) + } + + async fn rollback(&mut self) { + if let Some(completion) = self.schedule_rollback() { + wait_for_shutdown(&completion).await; + } } fn take_installed(&mut self) -> Vec { @@ -832,23 +860,7 @@ impl PluginInstallRollback { impl Drop for PluginInstallRollback { fn drop(&mut self) { - if !self.armed { - return; - } - self.close_resources(); - let current = self.current.take(); - let installed = std::mem::take(&mut self.installed); - let upstream = self.upstream.take(); - if current.is_none() && installed.is_empty() && upstream.is_none() { - return; - } - let runtime = self.runtime.clone(); - let cleanup_runtime = runtime.clone(); - runtime - .spawn(Box::pin(async move { - shutdown_staged_plugins(cleanup_runtime, current, installed, upstream).await; - })) - .detach(); + let _ = self.schedule_rollback(); } } @@ -1669,6 +1681,144 @@ mod tests { .expect("rollback aborted the install-scoped task"); } + struct BlockingFailingPlugin { + log: Log, + started: async_channel::Sender<()>, + release: async_channel::Receiver<()>, + } + + impl ClientPlugin for BlockingFailingPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("blocking-failure", "0.1.0").with_dependency("rollback") + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "install:blocking-failure"); + anyhow::bail!("injected failure") + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + let started = self.started.clone(); + let release = self.release.clone(); + Box::pin(async move { + record(&log, "shutdown:blocking-failure-started"); + let _ = started.try_send(()); + let _ = release.recv().await; + record(&log, "shutdown:blocking-failure-finished"); + Ok(()) + }) + } + } + + struct SignalAwareUpstream { + log: Log, + signalled: Arc, + shutdown_saw_signal: Arc, + } + + impl ClientLifecycle for SignalAwareUpstream { + fn install(&self, _client: Weak) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + Box::pin(async move { + record(&log, "install:upstream"); + Ok(()) + }) + } + + fn signal_shutdown(&self) { + if !self.signalled.swap(true, Ordering::AcqRel) { + record(&self.log, "signal:upstream"); + } + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = self.log.clone(); + let signalled = self.signalled.clone(); + let shutdown_saw_signal = self.shutdown_saw_signal.clone(); + Box::pin(async move { + shutdown_saw_signal.store(signalled.load(Ordering::Acquire), Ordering::Release); + record(&log, "shutdown:upstream"); + Ok(()) + }) + } + } + + #[tokio::test] + async fn cancelled_explicit_rollback_finishes_detached_and_signals_upstream() { + let log = Arc::new(Mutex::new(Vec::new())); + let task_dropped = Arc::new(AtomicBool::new(false)); + let signalled = Arc::new(AtomicBool::new(false)); + let shutdown_saw_signal = Arc::new(AtomicBool::new(false)); + let (started_tx, started_rx) = async_channel::bounded(1); + let (release_tx, release_rx) = async_channel::bounded(1); + let builder = complete_builder() + .await + .with_lifecycle(SignalAwareUpstream { + log: log.clone(), + signalled: signalled.clone(), + shutdown_saw_signal: shutdown_saw_signal.clone(), + }) + .with_plugin(BlockingFailingPlugin { + log: log.clone(), + started: started_tx, + release: release_rx, + }) + .with_plugin(RollbackPlugin { + log: log.clone(), + task_dropped: task_dropped.clone(), + }); + + let build = tokio::spawn(async move { builder.build().await }); + started_rx + .recv() + .await + .expect("failed plugin rollback started"); + assert!(signalled.load(Ordering::Acquire)); + build.abort(); + let _ = build.await; + release_tx.send(()).await.expect("release rollback hook"); + + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let complete = log + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .last() + .is_some_and(|entry| entry == "shutdown:upstream"); + if complete && task_dropped.load(Ordering::Acquire) { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("detached rollback completed after build cancellation"); + + assert!(shutdown_saw_signal.load(Ordering::Acquire)); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec![ + "install:upstream", + "install:rollback", + "install:blocking-failure", + "signal:upstream", + "shutdown:blocking-failure-started", + "shutdown:blocking-failure-finished", + "shutdown:rollback", + "shutdown:upstream" + ] + ); + } + struct BlockingInstallPlugin { log: Log, started: async_channel::Sender<()>, From 776f907472cd295d204c3785f74cd831845e84ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:59:10 -0300 Subject: [PATCH 15/46] fix(core): harden reentrant teardown races --- src/client.rs | 1 + src/client/extension_lifecycle.rs | 31 ++++++++++++++++++++++++ src/client/lifecycle.rs | 6 +++++ src/client/node_io.rs | 6 +++++ wacore/src/types/events.rs | 40 ++++++++++++++++++++++++++++++- 5 files changed, 83 insertions(+), 1 deletion(-) diff --git a/src/client.rs b/src/client.rs index 3a819d47d..d488e1ee6 100644 --- a/src/client.rs +++ b/src/client.rs @@ -659,6 +659,7 @@ pub struct Client { pub(crate) media_conn: Arc>>, pub(crate) is_logged_in: Arc, + pub(crate) login_transition: std::sync::Mutex<()>, pub(crate) is_connecting: Arc, pub(crate) is_running: Arc, /// Whether the noise socket is established (connected to WhatsApp servers). diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 15b3415d9..12347af5c 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -1314,6 +1314,37 @@ mod tests { assert_eq!(lifecycle.shutdowns.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn rejected_success_restores_logged_out_state() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let lifecycle = Arc::new(RecordingLifecycle::default()); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle) + .build() + .await + .expect("client build") + .into_client(); + client + .lifecycle + .as_ref() + .expect("lifecycle registration") + .signal_shutdown_sync(); + + let success = wacore_binary::builder::NodeBuilder::new("success").build(); + client.handle_success(&success.as_node_ref()).await; + + assert!(!client.is_logged_in()); + assert_eq!(client.connection_generation.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn replaced_scope_stays_closeable_by_its_generation() { let lifecycle = Arc::new(RecordingLifecycle::default()); diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 83374add2..cd12e6747 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -184,6 +184,7 @@ impl Client { persistence_manager: persistence_manager.clone(), media_conn: Arc::new(RwLock::new(None)), is_logged_in: Arc::new(AtomicBool::new(false)), + login_transition: std::sync::Mutex::new(()), is_connecting: Arc::new(AtomicBool::new(false)), is_running: Arc::new(AtomicBool::new(false)), is_connected: Arc::new(AtomicBool::new(false)), @@ -846,6 +847,10 @@ impl Client { } async fn cleanup_connection_state_inner(&self) { + let login_transition = self + .login_transition + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); // Bump the generation FIRST: it is the "this connection is over" // signal every per-connection loop already polls. Chat-lane workers // stop draining their queues (their remaining stanzas were never @@ -869,6 +874,7 @@ impl Client { // outgoing stanzas, which are transport-scoped. self.clear_sent_node_waiters(); self.is_logged_in.store(false, Ordering::Relaxed); + drop(login_transition); self.is_ready.store(false, Ordering::Relaxed); // Publish the disconnected state BEFORE draining VoIP calls (it used to be cleared only after // the socket teardown below): a concurrent accept()/call() setup that finishes its async work diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 3036354cd..82b3fbd4d 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -679,6 +679,10 @@ impl Client { tracing::instrument(name = "wa.conn.success", level = "debug", skip_all) )] pub(crate) async fn handle_success(self: &Arc, node: &wacore_binary::NodeRef<'_>) { + let login_transition = self + .login_transition + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); // Skip processing if an expected disconnect is pending (e.g., 515 received). // This prevents race conditions where a spawned success handler runs after // cleanup_connection_state has already reset is_logged_in. @@ -703,10 +707,12 @@ impl Client { && !self.expected_disconnect.load(Ordering::Acquire) }); if !opened { + self.is_logged_in.store(false, Ordering::SeqCst); debug!("Ignoring stanza retired during lifecycle publication"); return; } } + drop(login_transition); info!( "Successfully authenticated with WhatsApp servers! (gen={})", diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index b00f8c8c4..9ebd18178 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -448,8 +448,10 @@ impl CoreEventBusInner { let snapshot = Arc::new(HandlerSnapshot { handlers }); // Retire the entry before clearing bits; an early read may only be a // harmless false positive. - *guard = Arc::clone(&snapshot); + let retired = std::mem::replace(&mut *guard, Arc::clone(&snapshot)); self.store_aggregate(&snapshot); + drop(guard); + drop(retired); true } @@ -2326,6 +2328,42 @@ mod tests { assert_eq!(calls.load(Ordering::SeqCst), 1); } + #[test] + fn handler_drop_can_unsubscribe_from_the_same_bus() { + use std::sync::Mutex; + use std::sync::mpsc; + use std::time::Duration; + + struct OwnsSubscription(Mutex>); + impl EventHandler for OwnsSubscription { + fn handle_event(&self, _: Arc) {} + } + impl Drop for OwnsSubscription { + fn drop(&mut self) { + self.0 + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + } + } + + let bus = CoreEventBus::new(); + let owned = bus.subscribe_handler(ChannelEventHandler::new().0); + let owner = Arc::new(OwnsSubscription(Mutex::new(Some(owned)))); + let outer = bus.subscribe_handler(owner.clone()); + drop(owner); + + let (done_tx, done_rx) = mpsc::channel(); + std::thread::spawn(move || { + drop(outer); + let _ = done_tx.send(()); + }); + done_rx + .recv_timeout(Duration::from_secs(2)) + .expect("handler drop re-entered event-bus removal"); + assert!(!bus.has_handlers()); + } + #[test] fn in_flight_dispatch_can_finish_after_unsubscribe() { use std::sync::Barrier; From 35d3315b25c287ca2b2dfac32436409475072ca1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:15:54 -0300 Subject: [PATCH 16/46] fix(plugins): avoid reentrant teardown deadlock --- src/plugins/mod.rs | 79 +++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 75 insertions(+), 4 deletions(-) diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index b53e71257..21f7f6d32 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -272,10 +272,14 @@ impl PluginResources { return; } self.shutdown.notify(); - self.subscriptions - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clear(); + let subscriptions = { + let mut subscriptions = self + .subscriptions + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + std::mem::take(&mut *subscriptions) + }; + drop(subscriptions); } } @@ -2281,6 +2285,50 @@ mod tests { } } + struct ReentrantSubscriptionHandler { + events: PluginCoreEvents, + } + + impl EventHandler for ReentrantSubscriptionHandler { + fn handle_event(&self, _event: Arc) {} + } + + impl Drop for ReentrantSubscriptionHandler { + fn drop(&mut self) { + let _ = self.events.subscribe( + EventInterest::of(&[EventKind::Connected]), + Arc::new(NoopEventHandler), + ); + } + } + + struct ReentrantSubscriptionPlugin; + + impl ClientPlugin for ReentrantSubscriptionPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("reentrant-subscription", "0.1.0") + .with_capability(PluginCapability::CoreEvents) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async move { + let events = context + .core_events() + .cloned() + .ok_or_else(|| anyhow::anyhow!("core events capability missing"))?; + events.subscribe( + EventInterest::of(&[EventKind::Connected]), + Arc::new(ReentrantSubscriptionHandler { + events: events.clone(), + }), + )?; + Ok(Arc::new(())) + }) + } + } + #[tokio::test] async fn shutdown_removes_plugin_event_subscriptions_and_raw_lease() { let build = complete_builder() @@ -2308,6 +2356,29 @@ mod tests { assert!(!client.raw_node_forwarding_enabled()); } + #[tokio::test] + async fn resource_close_drops_reentrant_handlers_outside_the_subscription_lock() { + let client = complete_builder() + .await + .with_plugin(ReentrantSubscriptionPlugin) + .build() + .await + .expect("reentrant subscription plugin") + .into_client(); + let shutdown_client = client.clone(); + let (completed_tx, completed_rx) = std::sync::mpsc::sync_channel(1); + let shutdown = std::thread::spawn(move || { + shutdown_client.signal_shutdown_sync(); + let _ = completed_tx.send(()); + }); + + completed_rx + .recv_timeout(Duration::from_secs(2)) + .expect("reentrant handler teardown must not deadlock"); + shutdown.join().expect("shutdown thread"); + client.disconnect().await; + } + #[tokio::test] async fn synchronous_shutdown_closes_plugin_resources_with_live_client_refs() { let task_dropped = Arc::new(AtomicBool::new(false)); From 2137fcb5965128b2351ca1bd1fedd1b926f0659b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:45:18 -0300 Subject: [PATCH 17/46] fix(plugins): close late lifecycle ownership races --- src/client/extension_lifecycle.rs | 195 ++++++++++++++++++++++++++++-- src/client/lifecycle.rs | 9 ++ src/plugins/mod.rs | 143 ++++++++++++++++++---- 3 files changed, 312 insertions(+), 35 deletions(-) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 12347af5c..8f01f70f7 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -19,6 +19,7 @@ const CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); std::thread_local! { static ACTIVE_CALLBACK: Cell<*const LifecycleRegistration> = const { Cell::new(std::ptr::null()) }; + static ACTIVE_READY_PUBLICATION: Cell<*const LifecycleRegistration> = const { Cell::new(std::ptr::null()) }; } /// Observable state of one authenticated connection generation. @@ -158,6 +159,7 @@ pub trait ClientLifecycle: wacore::sync_marker::MaybeSendSync { pub(super) struct LifecycleRegistration { handler: Arc, runtime: Arc, + ready_publication: std::sync::Mutex<()>, scopes: std::sync::Mutex, callback_queue: std::sync::Mutex, shutdown_complete: AtomicBool, @@ -193,6 +195,23 @@ struct CallbackContextGuard { previous: *const LifecycleRegistration, } +struct ReadyPublicationGuard { + previous: *const LifecycleRegistration, +} + +impl ReadyPublicationGuard { + fn enter(registration: &LifecycleRegistration) -> Self { + let previous = ACTIVE_READY_PUBLICATION.replace(registration); + Self { previous } + } +} + +impl Drop for ReadyPublicationGuard { + fn drop(&mut self) { + ACTIVE_READY_PUBLICATION.set(self.previous); + } +} + impl CallbackContextGuard { fn enter(registration: &LifecycleRegistration) -> Self { let previous = ACTIVE_CALLBACK.replace(registration); @@ -210,6 +229,10 @@ fn callback_context_active(registration: &LifecycleRegistration) -> bool { ACTIVE_CALLBACK.with(|active| std::ptr::eq(active.get(), registration)) } +fn ready_publication_active(registration: &LifecycleRegistration) -> bool { + ACTIVE_READY_PUBLICATION.with(|active| std::ptr::eq(active.get(), registration)) +} + impl LifecycleRegistration { pub(super) fn new(handler: Arc, runtime: Arc) -> Self { Self::new_with_timeout(handler, runtime, CALLBACK_TIMEOUT) @@ -223,6 +246,7 @@ impl LifecycleRegistration { Self { handler, runtime, + ready_publication: std::sync::Mutex::new(()), scopes: std::sync::Mutex::new(ScopeRegistry::default()), callback_queue: std::sync::Mutex::new(CallbackQueue::default()), shutdown_complete: AtomicBool::new(false), @@ -300,19 +324,29 @@ impl LifecycleRegistration { done_rx.recv().await.unwrap_or(false) && !scope.is_cancelled() } - pub(super) fn cancel_scope(&self, generation: u64) { - if let Some(scope) = self.scope_for(generation) { - scope.cancel(); + pub(super) fn publish_ready(&self, generation: u64, publish: impl FnOnce()) -> bool { + let _publication = self.ready_publication(); + if self.terminal.load(Ordering::Acquire) { + return false; } + let Some(scope) = self.scope_for(generation) else { + return false; + }; + if scope.state() != ConnectionScopeState::Ready { + return false; + } + + let _publication_context = ReadyPublicationGuard::enter(self); + publish(); + true } - pub(super) fn cancel_active_scope(&self) { - let scopes = self.scopes(); - if let Some(scope) = &scopes.active { - scope.cancel(); - } - for scope in &scopes.retired { - scope.cancel(); + pub(super) fn cancel_scope(&self, generation: u64) { + if ready_publication_active(self) { + self.cancel_scope_inner(generation); + } else { + let _publication = self.ready_publication(); + self.cancel_scope_inner(generation); } } @@ -388,8 +422,12 @@ impl LifecycleRegistration { } pub(super) fn signal_shutdown_sync(&self) { - let first_signal = !self.terminal.swap(true, Ordering::AcqRel); - self.cancel_active_scope(); + let first_signal = if ready_publication_active(self) { + self.mark_terminal_and_cancel_scopes() + } else { + let _publication = self.ready_publication(); + self.mark_terminal_and_cancel_scopes() + }; if first_signal && std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { self.handler.signal_shutdown(); @@ -530,6 +568,30 @@ impl LifecycleRegistration { .unwrap_or_else(|poisoned| poisoned.into_inner()) } + fn ready_publication(&self) -> std::sync::MutexGuard<'_, ()> { + self.ready_publication + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn cancel_scope_inner(&self, generation: u64) { + if let Some(scope) = self.scope_for(generation) { + scope.cancel(); + } + } + + fn mark_terminal_and_cancel_scopes(&self) -> bool { + let first_signal = !self.terminal.swap(true, Ordering::AcqRel); + let scopes = self.scopes(); + if let Some(scope) = &scopes.active { + scope.cancel(); + } + for scope in &scopes.retired { + scope.cancel(); + } + first_signal + } + fn scope_for(&self, generation: u64) -> Option { let scopes = self.scopes(); scopes @@ -838,6 +900,115 @@ mod tests { assert_eq!(scope.state(), ConnectionScopeState::Closed); } + #[tokio::test] + async fn terminal_signal_rejects_ready_publication() { + let registration = Arc::new(LifecycleRegistration::new( + Arc::new(RecordingLifecycle::default()), + Arc::new(TokioRuntime), + )); + const GENERATION: u64 = 43; + assert!(registration.begin_scope_if_current(GENERATION, || true)); + assert!(registration.ready(GENERATION).await); + registration.signal_shutdown_sync(); + + let published = AtomicBool::new(false); + assert!(!registration.publish_ready(GENERATION, || { + published.store(true, Ordering::Release); + })); + assert!(!published.load(Ordering::Acquire)); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn terminal_cancellation_waits_for_ready_publication() { + let registration = Arc::new(LifecycleRegistration::new( + Arc::new(RecordingLifecycle::default()), + Arc::new(TokioRuntime), + )); + const GENERATION: u64 = 47; + assert!(registration.begin_scope_if_current(GENERATION, || true)); + assert!(registration.ready(GENERATION).await); + let scope = registration + .scope_for(GENERATION) + .expect("ready connection scope"); + + let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); + let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); + let (published_tx, published_rx) = std::sync::mpsc::sync_channel(1); + let publish_registration = registration.clone(); + let publish = std::thread::spawn(move || { + let published = publish_registration.publish_ready(GENERATION, || { + let _ = started_tx.send(()); + let _ = release_rx.recv(); + }); + let _ = published_tx.send(published); + }); + started_rx + .recv_timeout(Duration::from_secs(2)) + .expect("ready publication started"); + + let (attempted_tx, attempted_rx) = std::sync::mpsc::sync_channel(1); + let (cancelled_tx, cancelled_rx) = std::sync::mpsc::sync_channel(1); + let cancel_registration = registration.clone(); + let cancel = std::thread::spawn(move || { + let _ = attempted_tx.send(()); + cancel_registration.signal_shutdown_sync(); + let _ = cancelled_tx.send(()); + }); + attempted_rx + .recv_timeout(Duration::from_secs(2)) + .expect("terminal cancellation attempted"); + assert!( + cancelled_rx + .recv_timeout(Duration::from_millis(100)) + .is_err() + ); + + release_tx.send(()).expect("release ready publication"); + assert!( + published_rx + .recv_timeout(Duration::from_secs(2)) + .expect("ready publication completed") + ); + cancelled_rx + .recv_timeout(Duration::from_secs(2)) + .expect("terminal cancellation completed"); + publish.join().expect("publication thread"); + cancel.join().expect("cancellation thread"); + assert_eq!(scope.state(), ConnectionScopeState::Cancelled); + } + + #[tokio::test] + async fn ready_publication_allows_reentrant_terminal_signal() { + let registration = Arc::new(LifecycleRegistration::new( + Arc::new(RecordingLifecycle::default()), + Arc::new(TokioRuntime), + )); + const GENERATION: u64 = 53; + assert!(registration.begin_scope_if_current(GENERATION, || true)); + assert!(registration.ready(GENERATION).await); + let scope = registration + .scope_for(GENERATION) + .expect("ready connection scope"); + + let (completed_tx, completed_rx) = std::sync::mpsc::sync_channel(1); + let publish_registration = registration.clone(); + let signal_registration = registration.clone(); + let publish = std::thread::spawn(move || { + let published = publish_registration.publish_ready(GENERATION, || { + signal_registration.signal_shutdown_sync(); + }); + let _ = completed_tx.send(published); + }); + + assert!( + completed_rx + .recv_timeout(Duration::from_secs(2)) + .expect("reentrant terminal signal completed") + ); + publish.join().expect("publication thread"); + assert_eq!(scope.state(), ConnectionScopeState::Cancelled); + } + #[tokio::test] async fn cleanup_cancels_before_io_and_closes_after_authoritative_teardown() { let persistence_manager = Arc::new( diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index cd12e6747..52fbec60b 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -86,7 +86,16 @@ impl Client { debug!("Skipping Connected dispatch after generation changed"); return; } + if !lifecycle.publish_ready(generation, || self.publish_connected()) { + debug!("Skipping Connected dispatch after lifecycle cancellation"); + } + return; } + + self.publish_connected(); + } + + fn publish_connected(&self) { self.is_ready.store(true, Ordering::Relaxed); wacore::telemetry::set_connected(true); self.core.event_bus.dispatch(Event::Connected( diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 21f7f6d32..05d265fae 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -252,19 +252,28 @@ impl PluginResources { subscription: Subscription, raw_node_lease: Option, ) -> Result<(), PluginResourceError> { - let mut subscriptions = self - .subscriptions - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if self.closed.load(Ordering::Acquire) { - drop(subscription); - return Err(PluginResourceError::ShuttingDown); - } - subscriptions.push(PluginCoreEventSubscription { + let registration = PluginCoreEventSubscription { _subscription: subscription, _raw_node_lease: raw_node_lease, - }); - Ok(()) + }; + let rejected = { + let mut subscriptions = self + .subscriptions + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.closed.load(Ordering::Acquire) { + Some(registration) + } else { + subscriptions.push(registration); + None + } + }; + if let Some(rejected) = rejected { + drop(rejected); + Err(PluginResourceError::ShuttingDown) + } else { + Ok(()) + } } fn close(&self) { @@ -394,8 +403,7 @@ impl PluginIq { /// Capabilities and already-installed dependencies visible during installation. pub struct PluginContext { plugin_id: String, - apis: Arc, - dependency_markers: Vec, + dependencies: HashMap, core_events: Option, tasks: Option, messaging: Option, @@ -407,11 +415,11 @@ impl PluginContext { &self.plugin_id } + /// Return a declared dependency without making retained contexts own it. + /// Clone the returned API during installation if it must outlive this call. pub fn plugin(&self) -> Option> { - if !self.dependency_markers.contains(&TypeId::of::

()) { - return None; - } - self.apis.get::

() + let api = self.dependencies.get(&TypeId::of::

())?.upgrade()?; + downcast_api::(&api) } pub fn core_events(&self) -> Option<&PluginCoreEvents> { @@ -532,6 +540,7 @@ impl ErasedApiValue for TypedApi { } type ErasedApi = Arc; +type WeakErasedApi = Weak; #[derive(Default)] struct ApiRegistry { @@ -553,12 +562,15 @@ impl ApiRegistry { .clone() } - fn get(&self) -> Option> { + fn dependency_view(&self, markers: &[TypeId]) -> HashMap { let values = self .values .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); - downcast_api::(values.get(&TypeId::of::

())?) + markers + .iter() + .filter_map(|marker| values.get(marker).map(|api| (*marker, Arc::downgrade(api)))) + .collect() } } @@ -931,8 +943,7 @@ impl PluginHost { let capabilities = manifest.capabilities; PluginContext { plugin_id: manifest.id.clone(), - apis, - dependency_markers: dependency_markers.to_vec(), + dependencies: apis.dependency_view(dependency_markers), core_events: capabilities .contains(PluginCapability::CoreEvents) .then(|| PluginCoreEvents { @@ -1580,6 +1591,58 @@ mod tests { } } + struct ContextRetainingApi { + _context: PluginContext, + _drop_flag: DropFlag, + } + + struct ContextRetainingPlugin { + api_dropped: Arc, + } + + impl ClientPlugin for ContextRetainingPlugin { + type Api = ContextRetainingApi; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("context-retaining", "0.1.0") + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + let api_dropped = self.api_dropped.clone(); + Box::pin(async move { + Ok(Arc::new(ContextRetainingApi { + _context: context, + _drop_flag: DropFlag(api_dropped), + })) + }) + } + } + + #[tokio::test] + async fn retained_context_does_not_cycle_with_the_api_registry() { + let api_dropped = Arc::new(AtomicBool::new(false)); + let build = complete_builder() + .await + .with_plugin(ContextRetainingPlugin { + api_dropped: api_dropped.clone(), + }) + .build() + .await + .expect("context-retaining plugin"); + let (client, sync_tasks) = build.into_parts(); + drop(sync_tasks); + let api = client + .plugin::() + .expect("retained-context API"); + let weak_api = Arc::downgrade(&api); + drop(api); + + client.disconnect().await; + drop(client); + wait_for_flag(&api_dropped).await; + assert!(weak_api.upgrade().is_none()); + } + struct RollbackPlugin { log: Log, task_dropped: Arc, @@ -2305,7 +2368,7 @@ mod tests { struct ReentrantSubscriptionPlugin; impl ClientPlugin for ReentrantSubscriptionPlugin { - type Api = (); + type Api = PluginCoreEvents; fn manifest(&self) -> PluginManifest { PluginManifest::new("reentrant-subscription", "0.1.0") @@ -2324,7 +2387,7 @@ mod tests { events: events.clone(), }), )?; - Ok(Arc::new(())) + Ok(Arc::new(events)) }) } } @@ -2379,6 +2442,40 @@ mod tests { client.disconnect().await; } + #[tokio::test] + async fn rejected_subscription_drops_reentrant_handler_outside_the_subscription_lock() { + let client = complete_builder() + .await + .with_plugin(ReentrantSubscriptionPlugin) + .build() + .await + .expect("reentrant subscription plugin") + .into_client(); + let events = client + .plugin::() + .expect("plugin event API"); + client.signal_shutdown_sync(); + + let (completed_tx, completed_rx) = std::sync::mpsc::sync_channel(1); + let subscribe_events = events.clone(); + let subscribe = std::thread::spawn(move || { + let result = subscribe_events.subscribe( + EventInterest::of(&[EventKind::Connected]), + Arc::new(ReentrantSubscriptionHandler { + events: (*subscribe_events).clone(), + }), + ); + let _ = completed_tx.send(result); + }); + + let result = completed_rx + .recv_timeout(Duration::from_secs(2)) + .expect("rejected reentrant subscription must not deadlock"); + assert!(matches!(result, Err(PluginResourceError::ShuttingDown))); + subscribe.join().expect("subscription thread"); + client.disconnect().await; + } + #[tokio::test] async fn synchronous_shutdown_closes_plugin_resources_with_live_client_refs() { let task_dropped = Arc::new(AtomicBool::new(false)); From 9cef55db33a37328ef0c6c5965f2207a21089f34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:56:49 -0300 Subject: [PATCH 18/46] feat(plugins): add bounded custom event routing --- Cargo.lock | 1 + Cargo.toml | 3 +- src/lib.rs | 15 +- src/plugins/events.rs | 925 ++++++++++++++++++++++++++++++++++++++++++ src/plugins/mod.rs | 187 ++++++++- 5 files changed, 1118 insertions(+), 13 deletions(-) create mode 100644 src/plugins/events.rs diff --git a/Cargo.lock b/Cargo.lock index 66dfc9507..0e4aaf7eb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4313,6 +4313,7 @@ dependencies = [ "async-lock", "async-trait", "base64", + "bon", "buffa", "bytes", "cbc 0.2.1", diff --git a/Cargo.toml b/Cargo.toml index b8b73d7f0..bd3b6fc6a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -114,7 +114,7 @@ zlib-rs = { version = "0.6.5", default-features = false, features = ["std", "rus debug-snapshots = ["wacore/debug-snapshots"] # Build-time native plugin host. Kept opt-in so clients that do not use plugins # retain the pre-host binary footprint. -plugins = [] +plugins = ["dep:bon"] # Optional observability. Off by default: no `tracing` dep, zero overhead. # Emits tracing spans/events only; the application installs the subscriber # (and any OpenTelemetry bridge). See examples/observability.rs. @@ -173,6 +173,7 @@ async-trait = { workspace = true } base64 = { workspace = true } buffa = { workspace = true } bytes = { workspace = true } +bon = { workspace = true, optional = true } chrono = { workspace = true, features = ["clock"] } event-listener = { workspace = true } futures = { workspace = true, features = ["std"] } diff --git a/src/lib.rs b/src/lib.rs index a94bb0a81..cd13e2286 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -116,9 +116,13 @@ pub mod plugins; #[cfg(feature = "plugins")] pub use plugins::{ ClientPlugin, PluginCapabilities, PluginCapability, PluginConnectionScope, - PluginConnectionTasks, PluginContext, PluginCoreEvents, PluginIq, PluginIqError, - PluginManifest, PluginMessaging, PluginMessagingError, PluginPlanError, PluginResourceError, - PluginTasks, + PluginConnectionTasks, PluginContext, PluginCoreEvents, PluginEventEndpointConfig, + PluginEventEndpointStats, PluginEventEnvelope, PluginEventOverflow, PluginEventPayloadEncoding, + PluginEventPublishError, PluginEventPublishReport, PluginEventReceiveError, + PluginEventRouteError, PluginEventRouter, PluginEventSelector, PluginEventSubscribeError, + PluginEventSubscription, PluginEventTopic, PluginEventTryReceiveError, PluginEvents, PluginIq, + PluginIqError, PluginManifest, PluginMessaging, PluginMessagingError, PluginPlanError, + PluginResourceError, PluginTasks, }; pub mod request; pub(crate) mod signal_flush; @@ -194,7 +198,10 @@ pub mod prelude { }; #[cfg(feature = "plugins")] pub use crate::plugins::{ - ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, PluginManifest, + ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, + PluginEventEndpointConfig, PluginEventOverflow, PluginEventPayloadEncoding, + PluginEventRouter, PluginEventSelector, PluginEventSubscription, PluginEventTopic, + PluginEvents, PluginManifest, }; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] diff --git a/src/plugins/events.rs b/src/plugins/events.rs new file mode 100644 index 000000000..6f30d5498 --- /dev/null +++ b/src/plugins/events.rs @@ -0,0 +1,925 @@ +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::sync::{Arc, Mutex, RwLock}; + +use async_channel::{Receiver, Sender, TryRecvError, TrySendError}; +use bytes::Bytes; +use portable_atomic::{AtomicBool, AtomicU64, Ordering}; +use thiserror::Error; + +use super::{PluginResourceError, PluginResources, valid_plugin_id}; + +const MAX_ENDPOINT_CAPACITY: usize = 65_536; +const MAX_ENDPOINT_SELECTORS: usize = 1_024; + +/// Encoding of a custom plugin event payload. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum PluginEventPayloadEncoding { + Json, + Binary, +} + +impl PluginEventPayloadEncoding { + pub const fn identifier(self) -> &'static str { + match self { + Self::Json => "json", + Self::Binary => "binary", + } + } +} + +/// Validated second-level topic within one plugin namespace. +#[derive(Clone, PartialEq, Eq, Hash)] +pub struct PluginEventTopic(Arc); + +impl PluginEventTopic { + pub fn new(topic: impl Into) -> Result { + let topic = topic.into(); + if !valid_topic(&topic) { + return Err(PluginEventRouteError::InvalidTopic { topic }); + } + Ok(Self(Arc::from(topic))) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl fmt::Debug for PluginEventTopic { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_tuple("PluginEventTopic") + .field(&self.0) + .finish() + } +} + +impl fmt::Display for PluginEventTopic { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.0) + } +} + +/// Exact `(plugin_id, topic)` route selected by one endpoint. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct PluginEventSelector { + route: RouteKey, +} + +impl PluginEventSelector { + pub fn new( + plugin_id: impl Into, + topic: PluginEventTopic, + ) -> Result { + let plugin_id = plugin_id.into(); + if !valid_plugin_id(&plugin_id) { + return Err(PluginEventRouteError::InvalidPluginId { plugin_id }); + } + Ok(Self { + route: RouteKey { + plugin_id: Arc::from(plugin_id), + topic, + }, + }) + } + + pub fn plugin_id(&self) -> &str { + &self.route.plugin_id + } + + pub fn topic(&self) -> &PluginEventTopic { + &self.route.topic + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct RouteKey { + plugin_id: Arc, + topic: PluginEventTopic, +} + +/// Routed event shared by every matching endpoint without copying its payload. +#[derive(Debug, Clone, bon::Builder)] +#[non_exhaustive] +pub struct PluginEventEnvelope { + pub plugin_id: Arc, + pub topic: PluginEventTopic, + pub schema_version: u32, + pub payload_encoding: PluginEventPayloadEncoding, + pub payload: Bytes, + pub connection_generation: u64, + /// Monotonic sequence for this route while it has at least one subscriber. + /// + /// Dropped events consume a sequence number, allowing one endpoint to detect loss. The + /// sequence resets after the last subscriber to the route is removed. + pub sequence: u64, +} + +/// Behavior when one endpoint cannot keep up with publishers. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum PluginEventOverflow { + DropNewest, + DropOldest, +} + +/// Required queue policy for one independent consumer endpoint. +/// +/// Capacity counts envelopes rather than bytes. Native plugins are trusted, and payloads are +/// shared across matching endpoints. A foreign adapter must enforce its wire payload limit before +/// publishing into this router. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PluginEventEndpointConfig { + capacity: usize, + overflow: PluginEventOverflow, +} + +impl PluginEventEndpointConfig { + pub const fn new(capacity: usize, overflow: PluginEventOverflow) -> Self { + Self { capacity, overflow } + } + + pub const fn capacity(self) -> usize { + self.capacity + } + + pub const fn overflow(self) -> PluginEventOverflow { + self.overflow + } +} + +/// Syntactic route validation failure. +#[derive(Debug, Error, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum PluginEventRouteError { + #[error("invalid plugin id `{plugin_id}`")] + InvalidPluginId { plugin_id: String }, + #[error("invalid plugin event topic `{topic}`")] + InvalidTopic { topic: String }, +} + +/// Endpoint registration failure. +#[derive(Debug, Error, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum PluginEventSubscribeError { + #[error(transparent)] + Resource(#[from] PluginResourceError), + #[error("at least one plugin event selector is required")] + EmptySelectors, + #[error("plugin event endpoint selector count exceeds the maximum of {max}")] + TooManySelectors { max: usize }, + #[error("plugin event endpoint capacity {capacity} is outside 1..={max}")] + InvalidCapacity { capacity: usize, max: usize }, + #[error("plugin `{plugin_id}` is not registered as a custom-event publisher")] + UnknownPublisher { plugin_id: String }, + #[error("plugin event endpoint identifiers are exhausted")] + EndpointIdsExhausted, + #[error("the plugin event router is closed")] + Closed, +} + +/// Custom event publication failure. +#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum PluginEventPublishError { + #[error(transparent)] + Resource(#[from] PluginResourceError), + #[error("plugin event schema version must be greater than zero")] + InvalidSchemaVersion, + #[error("the plugin event router is closed")] + Closed, + #[error("the plugin event sequence is exhausted")] + SequenceExhausted, +} + +/// Result of one non-blocking fan-out attempt. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +#[non_exhaustive] +pub struct PluginEventPublishReport { + pub matched: u64, + pub enqueued: u64, + pub dropped: u64, + pub closed: u64, +} + +/// Cumulative state for one endpoint queue. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub struct PluginEventEndpointStats { + pub enqueued: u64, + pub delivered: u64, + pub dropped: u64, + pub queue_depth: usize, + pub capacity: usize, +} + +#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +#[error("the plugin event endpoint is closed")] +pub struct PluginEventReceiveError; + +#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum PluginEventTryReceiveError { + #[error("the plugin event endpoint queue is empty")] + Empty, + #[error("the plugin event endpoint is closed")] + Closed, +} + +enum EnqueueOutcome { + Enqueued, + Dropped, + EnqueuedAfterDrop, + Closed, +} + +struct EventEndpoint { + id: u64, + sender: Sender>, + overflow: PluginEventOverflow, + capacity: usize, + enqueued: AtomicU64, + delivered: AtomicU64, + dropped: AtomicU64, +} + +impl EventEndpoint { + fn enqueue(&self, event: Arc) -> EnqueueOutcome { + match self.overflow { + PluginEventOverflow::DropNewest => match self.sender.try_send(event) { + Ok(()) => { + self.enqueued.fetch_add(1, Ordering::Relaxed); + EnqueueOutcome::Enqueued + } + Err(TrySendError::Full(_)) => { + self.dropped.fetch_add(1, Ordering::Relaxed); + EnqueueOutcome::Dropped + } + Err(TrySendError::Closed(_)) => EnqueueOutcome::Closed, + }, + PluginEventOverflow::DropOldest => match self.sender.force_send(event) { + Ok(evicted) => { + self.enqueued.fetch_add(1, Ordering::Relaxed); + if evicted.is_some() { + self.dropped.fetch_add(1, Ordering::Relaxed); + EnqueueOutcome::EnqueuedAfterDrop + } else { + EnqueueOutcome::Enqueued + } + } + Err(_) => EnqueueOutcome::Closed, + }, + } + } + + fn close(&self) { + self.sender.close(); + } + + fn stats(&self) -> PluginEventEndpointStats { + PluginEventEndpointStats { + enqueued: self.enqueued.load(Ordering::Relaxed), + delivered: self.delivered.load(Ordering::Relaxed), + dropped: self.dropped.load(Ordering::Relaxed), + queue_depth: self.sender.len(), + capacity: self.capacity, + } + } +} + +struct RouteClock { + sequence: Mutex, +} + +struct RouteEntry { + clock: Arc, + endpoints: Arc<[Arc]>, +} + +#[derive(Default)] +struct RouterState { + routes: HashMap, + endpoints: HashMap>, +} + +struct PluginEventRouterInner { + plugin_ids: HashSet>, + state: RwLock, + next_endpoint_id: AtomicU64, + closed: AtomicBool, +} + +impl PluginEventRouterInner { + fn unsubscribe(&self, endpoint_id: u64, selectors: &[PluginEventSelector]) { + let endpoint = { + let mut state = self + .state + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let endpoint = state.endpoints.remove(&endpoint_id); + for selector in selectors { + let remove_route = if let Some(route) = state.routes.get_mut(&selector.route) { + let remaining = route + .endpoints + .iter() + .filter(|endpoint| endpoint.id != endpoint_id) + .cloned() + .collect::>(); + route.endpoints = remaining.into(); + route.endpoints.is_empty() + } else { + false + }; + if remove_route { + state.routes.remove(&selector.route); + } + } + endpoint + }; + if let Some(endpoint) = endpoint { + endpoint.close(); + } + } + + fn close(&self) { + if self.closed.swap(true, Ordering::AcqRel) { + return; + } + let endpoints = { + let mut state = self + .state + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.routes.clear(); + std::mem::take(&mut state.endpoints) + }; + for endpoint in endpoints.into_values() { + endpoint.close(); + } + } +} + +/// Read-only subscription boundary for native consumers and future foreign adapters. +/// +/// Routes are exact `(plugin_id, topic)` matches. Closing the router prevents new publications and +/// subscriptions, while already queued envelopes remain available before receivers observe closure. +#[derive(Clone)] +pub struct PluginEventRouter { + inner: Arc, +} + +impl PluginEventRouter { + pub(super) fn new(plugin_ids: impl IntoIterator) -> Self { + Self { + inner: Arc::new(PluginEventRouterInner { + plugin_ids: plugin_ids.into_iter().map(Arc::from).collect(), + state: RwLock::new(RouterState::default()), + next_endpoint_id: AtomicU64::new(1), + closed: AtomicBool::new(false), + }), + } + } + + pub fn has_subscribers(&self, selector: &PluginEventSelector) -> bool { + if self.inner.closed.load(Ordering::Acquire) { + return false; + } + self.inner + .state + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .routes + .get(&selector.route) + .is_some_and(|route| !route.endpoints.is_empty()) + } + + pub fn subscribe( + &self, + selectors: impl IntoIterator, + config: PluginEventEndpointConfig, + ) -> Result { + if config.capacity == 0 || config.capacity > MAX_ENDPOINT_CAPACITY { + return Err(PluginEventSubscribeError::InvalidCapacity { + capacity: config.capacity, + max: MAX_ENDPOINT_CAPACITY, + }); + } + if self.inner.closed.load(Ordering::Acquire) { + return Err(PluginEventSubscribeError::Closed); + } + + let mut seen = HashSet::new(); + let mut unique_selectors = Vec::new(); + for selector in selectors { + if !seen.insert(selector.route.clone()) { + continue; + } + if unique_selectors.len() == MAX_ENDPOINT_SELECTORS { + return Err(PluginEventSubscribeError::TooManySelectors { + max: MAX_ENDPOINT_SELECTORS, + }); + } + unique_selectors.push(selector); + } + let selectors = unique_selectors; + if selectors.is_empty() { + return Err(PluginEventSubscribeError::EmptySelectors); + } + for selector in &selectors { + if !self.inner.plugin_ids.contains(selector.plugin_id()) { + return Err(PluginEventSubscribeError::UnknownPublisher { + plugin_id: selector.plugin_id().to_string(), + }); + } + } + + let endpoint_id = self + .inner + .next_endpoint_id + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1)) + .map_err(|_| PluginEventSubscribeError::EndpointIdsExhausted)?; + let (sender, receiver) = async_channel::bounded(config.capacity); + let endpoint = Arc::new(EventEndpoint { + id: endpoint_id, + sender, + overflow: config.overflow, + capacity: config.capacity, + enqueued: AtomicU64::new(0), + delivered: AtomicU64::new(0), + dropped: AtomicU64::new(0), + }); + + { + let mut state = self + .inner + .state + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.inner.closed.load(Ordering::Acquire) { + return Err(PluginEventSubscribeError::Closed); + } + state.endpoints.insert(endpoint_id, endpoint.clone()); + for selector in &selectors { + let route = state + .routes + .entry(selector.route.clone()) + .or_insert_with(|| RouteEntry { + clock: Arc::new(RouteClock { + sequence: Mutex::new(0), + }), + endpoints: Arc::from([]), + }); + let endpoints = route + .endpoints + .iter() + .cloned() + .chain(std::iter::once(endpoint.clone())) + .collect::>(); + route.endpoints = endpoints.into(); + } + } + + Ok(PluginEventSubscription { + router: self.clone(), + endpoint, + receiver, + selectors, + }) + } + + fn publish( + &self, + plugin_id: &Arc, + topic: &PluginEventTopic, + schema_version: u32, + payload_encoding: PluginEventPayloadEncoding, + payload: Bytes, + connection_generation: u64, + ) -> Result { + if schema_version == 0 { + return Err(PluginEventPublishError::InvalidSchemaVersion); + } + if self.inner.closed.load(Ordering::Acquire) { + return Err(PluginEventPublishError::Closed); + } + + let route_key = RouteKey { + plugin_id: plugin_id.clone(), + topic: topic.clone(), + }; + let Some((clock, endpoints)) = self + .inner + .state + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .routes + .get(&route_key) + .map(|route| (route.clock.clone(), route.endpoints.clone())) + else { + return Ok(PluginEventPublishReport::default()); + }; + + let mut sequence = clock + .sequence + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let next_sequence = sequence + .checked_add(1) + .ok_or(PluginEventPublishError::SequenceExhausted)?; + *sequence = next_sequence; + let event = Arc::new( + PluginEventEnvelope::builder() + .plugin_id(plugin_id.clone()) + .topic(topic.clone()) + .schema_version(schema_version) + .payload_encoding(payload_encoding) + .payload(payload) + .connection_generation(connection_generation) + .sequence(next_sequence) + .build(), + ); + + let mut report = PluginEventPublishReport { + matched: u64::try_from(endpoints.len()).unwrap_or(u64::MAX), + ..PluginEventPublishReport::default() + }; + for endpoint in endpoints.iter() { + match endpoint.enqueue(event.clone()) { + EnqueueOutcome::Enqueued => report.enqueued += 1, + EnqueueOutcome::Dropped => report.dropped += 1, + EnqueueOutcome::EnqueuedAfterDrop => { + report.enqueued += 1; + report.dropped += 1; + } + EnqueueOutcome::Closed => report.closed += 1, + } + } + Ok(report) + } + + pub(super) fn close(&self) { + self.inner.close(); + } +} + +/// One bounded endpoint. Dropping it unregisters every selected route atomically. +#[must_use = "dropping the subscription unregisters its plugin event routes"] +pub struct PluginEventSubscription { + router: PluginEventRouter, + endpoint: Arc, + receiver: Receiver>, + selectors: Vec, +} + +impl PluginEventSubscription { + pub fn id(&self) -> u64 { + self.endpoint.id + } + + pub fn selectors(&self) -> &[PluginEventSelector] { + &self.selectors + } + + pub fn stats(&self) -> PluginEventEndpointStats { + self.endpoint.stats() + } + + pub async fn recv(&self) -> Result, PluginEventReceiveError> { + let event = self + .receiver + .recv() + .await + .map_err(|_| PluginEventReceiveError)?; + self.endpoint.delivered.fetch_add(1, Ordering::Relaxed); + Ok(event) + } + + pub fn try_recv(&self) -> Result, PluginEventTryReceiveError> { + let event = self.receiver.try_recv().map_err(|error| match error { + TryRecvError::Empty => PluginEventTryReceiveError::Empty, + TryRecvError::Closed => PluginEventTryReceiveError::Closed, + })?; + self.endpoint.delivered.fetch_add(1, Ordering::Relaxed); + Ok(event) + } +} + +impl Drop for PluginEventSubscription { + fn drop(&mut self) { + self.router + .inner + .unsubscribe(self.endpoint.id, &self.selectors); + } +} + +/// Context-bound custom event capability. A plugin can publish only under its own ID. +/// +/// Consumers subscribe through [`PluginEventRouter`], keeping publication authority separate from +/// native or future foreign endpoints. +#[derive(Clone)] +pub struct PluginEvents { + plugin_id: Arc, + router: PluginEventRouter, + resources: Arc, + connection_generation: Arc, +} + +impl PluginEvents { + pub fn selector(&self, topic: &PluginEventTopic) -> PluginEventSelector { + PluginEventSelector { + route: RouteKey { + plugin_id: self.plugin_id.clone(), + topic: topic.clone(), + }, + } + } + + pub fn has_subscribers(&self, topic: &PluginEventTopic) -> bool { + self.router.has_subscribers(&self.selector(topic)) + } + + pub fn publish( + &self, + topic: &PluginEventTopic, + schema_version: u32, + payload_encoding: PluginEventPayloadEncoding, + payload: impl Into, + ) -> Result { + self.resources.ensure_active()?; + self.router.publish( + &self.plugin_id, + topic, + schema_version, + payload_encoding, + payload.into(), + self.connection_generation.load(Ordering::Acquire), + ) + } +} + +pub(super) fn publisher( + plugin_id: &str, + router: PluginEventRouter, + resources: Arc, + connection_generation: Arc, +) -> PluginEvents { + PluginEvents { + plugin_id: Arc::from(plugin_id), + router, + resources, + connection_generation, + } +} + +fn valid_topic(topic: &str) -> bool { + valid_plugin_id(topic) +} + +#[cfg(test)] +mod tests { + use std::thread; + + use super::*; + + fn topic(value: &str) -> PluginEventTopic { + PluginEventTopic::new(value).expect("valid topic") + } + + fn selector(plugin_id: &str, topic: &PluginEventTopic) -> PluginEventSelector { + PluginEventSelector::new(plugin_id, topic.clone()).expect("valid selector") + } + + fn publish( + router: &PluginEventRouter, + plugin_id: &str, + topic: &PluginEventTopic, + value: u32, + ) -> PluginEventPublishReport { + router + .publish( + &Arc::from(plugin_id), + topic, + 1, + PluginEventPayloadEncoding::Binary, + Bytes::copy_from_slice(&value.to_be_bytes()), + 7, + ) + .expect("event publication") + } + + #[test] + fn routes_only_exact_plugin_and_topic_matches() { + let router = PluginEventRouter::new(["metrics".to_string(), "audit".to_string()]); + let tick = topic("tick"); + let other = topic("other"); + let subscription = router + .subscribe( + [selector("metrics", &tick)], + PluginEventEndpointConfig::new(4, PluginEventOverflow::DropNewest), + ) + .expect("subscription"); + + assert_eq!(publish(&router, "metrics", &other, 1).matched, 0); + assert_eq!(publish(&router, "audit", &tick, 2).matched, 0); + assert_eq!(publish(&router, "metrics", &tick, 3).enqueued, 1); + let event = subscription.try_recv().expect("routed event"); + assert_eq!(&*event.plugin_id, "metrics"); + assert_eq!(event.topic, tick); + assert_eq!(event.connection_generation, 7); + assert_eq!(event.sequence, 1); + assert_eq!(event.payload, Bytes::copy_from_slice(&3u32.to_be_bytes())); + } + + #[test] + fn drop_newest_preserves_the_queued_prefix_and_counts_loss() { + let router = PluginEventRouter::new(["metrics".to_string()]); + let tick = topic("tick"); + let subscription = router + .subscribe( + [selector("metrics", &tick)], + PluginEventEndpointConfig::new(2, PluginEventOverflow::DropNewest), + ) + .expect("subscription"); + + for value in 1..=4 { + publish(&router, "metrics", &tick, value); + } + + assert_eq!(subscription.try_recv().expect("first").sequence, 1); + assert_eq!(subscription.try_recv().expect("second").sequence, 2); + assert!(matches!( + subscription.try_recv(), + Err(PluginEventTryReceiveError::Empty) + )); + assert_eq!( + subscription.stats(), + PluginEventEndpointStats { + enqueued: 2, + delivered: 2, + dropped: 2, + queue_depth: 0, + capacity: 2, + } + ); + } + + #[test] + fn drop_oldest_preserves_the_latest_events_and_counts_evictions() { + let router = PluginEventRouter::new(["metrics".to_string()]); + let tick = topic("tick"); + let subscription = router + .subscribe( + [selector("metrics", &tick)], + PluginEventEndpointConfig::new(2, PluginEventOverflow::DropOldest), + ) + .expect("subscription"); + + for value in 1..=4 { + publish(&router, "metrics", &tick, value); + } + + assert_eq!(subscription.try_recv().expect("third").sequence, 3); + assert_eq!(subscription.try_recv().expect("fourth").sequence, 4); + assert_eq!(subscription.stats().enqueued, 4); + assert_eq!(subscription.stats().dropped, 2); + } + + #[test] + fn backpressure_is_isolated_per_endpoint() { + let router = PluginEventRouter::new(["metrics".to_string()]); + let tick = topic("tick"); + let slow = router + .subscribe( + [selector("metrics", &tick)], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ) + .expect("slow endpoint"); + let fast = router + .subscribe( + [selector("metrics", &tick)], + PluginEventEndpointConfig::new(4, PluginEventOverflow::DropNewest), + ) + .expect("fast endpoint"); + + publish(&router, "metrics", &tick, 1); + publish(&router, "metrics", &tick, 2); + + assert_eq!(slow.stats().dropped, 1); + assert_eq!(fast.stats().dropped, 0); + assert_eq!(fast.try_recv().expect("fast first").sequence, 1); + assert_eq!(fast.try_recv().expect("fast second").sequence, 2); + } + + #[tokio::test] + async fn drop_unregisters_and_router_close_wakes_receivers() { + let router = PluginEventRouter::new(["metrics".to_string()]); + let tick = topic("tick"); + let selector = selector("metrics", &tick); + let subscription = router + .subscribe( + [selector.clone()], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ) + .expect("subscription"); + assert!(router.has_subscribers(&selector)); + drop(subscription); + assert!(!router.has_subscribers(&selector)); + assert_eq!(publish(&router, "metrics", &tick, 1).matched, 0); + + let subscription = router + .subscribe( + [selector], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ) + .expect("second subscription"); + assert_eq!(publish(&router, "metrics", &tick, 2).enqueued, 1); + router.close(); + assert_eq!(subscription.recv().await.expect("queued event").sequence, 1); + assert!(matches!( + subscription.recv().await, + Err(PluginEventReceiveError) + )); + } + + #[test] + fn rejects_invalid_or_unknown_endpoint_configuration() { + assert!(PluginEventTopic::new("Invalid").is_err()); + let tick = topic("tick"); + assert!(PluginEventSelector::new("Invalid", tick.clone()).is_err()); + let router = PluginEventRouter::new(["metrics".to_string()]); + assert!(matches!( + router.subscribe( + [selector("unknown", &tick)], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ), + Err(PluginEventSubscribeError::UnknownPublisher { .. }) + )); + assert!(matches!( + router.subscribe( + [selector("metrics", &tick)], + PluginEventEndpointConfig::new(0, PluginEventOverflow::DropNewest), + ), + Err(PluginEventSubscribeError::InvalidCapacity { .. }) + )); + assert!(matches!( + router.subscribe( + [], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ), + Err(PluginEventSubscribeError::EmptySelectors) + )); + let too_many = (0..=MAX_ENDPOINT_SELECTORS) + .map(|index| selector("metrics", &topic(&format!("topic-{index}")))) + .collect::>(); + assert!(matches!( + router.subscribe( + too_many, + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ), + Err(PluginEventSubscribeError::TooManySelectors { .. }) + )); + } + + #[test] + fn concurrent_publish_keeps_route_sequences_in_queue_order() { + const THREADS: usize = 8; + const EVENTS_PER_THREAD: usize = 100; + let router = PluginEventRouter::new(["metrics".to_string()]); + let tick = topic("tick"); + let subscription = router + .subscribe( + [selector("metrics", &tick)], + PluginEventEndpointConfig::new( + THREADS * EVENTS_PER_THREAD, + PluginEventOverflow::DropNewest, + ), + ) + .expect("subscription"); + + let threads = (0..THREADS) + .map(|_| { + let router = router.clone(); + let tick = tick.clone(); + thread::spawn(move || { + for value in 0..EVENTS_PER_THREAD { + publish(&router, "metrics", &tick, value as u32); + } + }) + }) + .collect::>(); + for thread in threads { + thread.join().expect("publisher thread"); + } + + for sequence in 1..=(THREADS * EVENTS_PER_THREAD) as u64 { + assert_eq!( + subscription.try_recv().expect("ordered event").sequence, + sequence + ); + } + assert_eq!(subscription.stats().dropped, 0); + } +} diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 05d265fae..e260108dc 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -1,5 +1,15 @@ //! Build-time client plugins and their capability-scoped host. +mod events; + +pub use events::{ + PluginEventEndpointConfig, PluginEventEndpointStats, PluginEventEnvelope, PluginEventOverflow, + PluginEventPayloadEncoding, PluginEventPublishError, PluginEventPublishReport, + PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, PluginEventSelector, + PluginEventSubscribeError, PluginEventSubscription, PluginEventTopic, + PluginEventTryReceiveError, PluginEvents, +}; + use std::any::{Any, TypeId}; use std::collections::{BTreeSet, HashMap, HashSet}; use std::future::Future; @@ -9,6 +19,7 @@ use std::sync::{Arc, Mutex, OnceLock, Weak}; use std::time::Duration; use futures::FutureExt; +use portable_atomic::AtomicU64; use thiserror::Error; use wacore::iq::spec::IqSpec; use wacore::runtime::{ @@ -29,6 +40,7 @@ const CAP_CORE_EVENTS: u8 = 1 << 0; const CAP_TASKS: u8 = 1 << 1; const CAP_MESSAGING: u8 = 1 << 2; const CAP_IQ: u8 = 1 << 3; +const CAP_PLUGIN_EVENTS: u8 = 1 << 4; const PLUGIN_CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); /// A capability a plugin asks the host to expose during installation. @@ -39,6 +51,7 @@ pub enum PluginCapability { Tasks, Messaging, Iq, + PluginEvents, } impl PluginCapability { @@ -48,6 +61,7 @@ impl PluginCapability { Self::Tasks => "tasks.spawn", Self::Messaging => "messaging.send", Self::Iq => "iq.execute", + Self::PluginEvents => "events.plugin.publish", } } @@ -57,6 +71,7 @@ impl PluginCapability { Self::Tasks => CAP_TASKS, Self::Messaging => CAP_MESSAGING, Self::Iq => CAP_IQ, + Self::PluginEvents => CAP_PLUGIN_EVENTS, } } } @@ -408,6 +423,7 @@ pub struct PluginContext { tasks: Option, messaging: Option, iq: Option, + plugin_events: Option, } impl PluginContext { @@ -437,6 +453,10 @@ impl PluginContext { pub fn iq(&self) -> Option<&PluginIq> { self.iq.as_ref() } + + pub fn plugin_events(&self) -> Option<&PluginEvents> { + self.plugin_events.as_ref() + } } /// One connection generation plus its optional connection-scoped task capability. @@ -887,6 +907,7 @@ pub(crate) struct PluginHost { installed: OnceLock>, apis: OnceLock>, runtime: OnceLock>, + event_router: Option, callback_timeout: Duration, } @@ -904,7 +925,18 @@ impl PluginHost { .ordered .iter() .map(|plugin| plugin.manifest.clone()) - .collect(); + .collect::>(); + let event_publishers = manifests + .iter() + .filter(|manifest| { + manifest + .capabilities + .contains(PluginCapability::PluginEvents) + }) + .map(|manifest| manifest.id.clone()) + .collect::>(); + let event_router = + (!event_publishers.is_empty()).then(|| PluginEventRouter::new(event_publishers)); Arc::new(Self { ordered: plan.ordered, manifests, @@ -912,6 +944,7 @@ impl PluginHost { installed: OnceLock::new(), apis: OnceLock::new(), runtime: OnceLock::new(), + event_router, callback_timeout, }) } @@ -934,16 +967,17 @@ impl PluginHost { fn context( &self, client: &Weak, - manifest: &PluginManifest, - dependency_markers: &[TypeId], + planned: &PlannedPlugin, resources: Arc, apis: Arc, runtime: Arc, + connection_generation: Arc, ) -> PluginContext { + let manifest = &planned.manifest; let capabilities = manifest.capabilities; PluginContext { plugin_id: manifest.id.clone(), - dependencies: apis.dependency_view(dependency_markers), + dependencies: apis.dependency_view(&planned.dependency_markers), core_events: capabilities .contains(PluginCapability::CoreEvents) .then(|| PluginCoreEvents { @@ -968,6 +1002,18 @@ impl PluginHost { client: client.clone(), resources: Arc::clone(&resources), }), + plugin_events: self + .event_router + .as_ref() + .filter(|_| capabilities.contains(PluginCapability::PluginEvents)) + .map(|router| { + events::publisher( + &manifest.id, + router.clone(), + Arc::clone(&resources), + connection_generation, + ) + }), } } @@ -993,6 +1039,7 @@ impl PluginHost { anyhow::bail!("client was dropped during plugin installation"); }; let runtime = strong_client.runtime.clone(); + let connection_generation = strong_client.connection_generation.clone(); drop(strong_client); self.runtime .set(runtime.clone()) @@ -1012,11 +1059,11 @@ impl PluginHost { let resources = PluginResources::new(); let context = self.context( &client, - &planned.manifest, - &planned.dependency_markers, + planned, Arc::clone(&resources), Arc::clone(&staging), runtime.clone(), + connection_generation.clone(), ); rollback.current = Some(InstalledPlugin { plugin: planned.plugin.clone(), @@ -1119,6 +1166,9 @@ impl ClientLifecycle for PluginHost { } fn signal_shutdown(&self) { + if let Some(router) = &self.event_router { + router.close(); + } for plugin in self.installed.get().into_iter().flatten().rev() { plugin.resources.close(); } @@ -1246,6 +1296,15 @@ impl Client { .map(|host| host.manifests()) .unwrap_or_default() } + + /// Subscribe to custom events emitted by installed plugins. + /// + /// Returns `None` when no manifest requested custom-event publication. + pub fn plugin_event_router(&self) -> Option { + self.plugin_host + .as_ref() + .and_then(|host| host.event_router.clone()) + } } #[cfg(test)] @@ -1253,6 +1312,8 @@ mod tests { use std::sync::atomic::AtomicBool; use std::time::Duration; + use bytes::Bytes; + use super::*; use crate::client::{ClientBuilder, ClientBuilderError}; use crate::runtime_impl::TokioRuntime; @@ -2508,7 +2569,7 @@ mod tests { struct CapabilityProbe; impl ClientPlugin for CapabilityProbe { - type Api = [bool; 4]; + type Api = [bool; 5]; fn manifest(&self) -> PluginManifest { PluginManifest::new("capability-probe", "0.1.0") @@ -2522,6 +2583,7 @@ mod tests { context.tasks().is_some(), context.messaging().is_some(), context.iq().is_some(), + context.plugin_events().is_some(), ])) }) } @@ -2538,8 +2600,117 @@ mod tests { let client = build.into_client(); assert_eq!( client.plugin::().as_deref(), - Some(&[false, false, true, false]) + Some(&[false, false, true, false, false]) + ); + assert!(client.plugin_event_router().is_none()); + client.disconnect().await; + } + + struct PluginEventPublisher; + + impl ClientPlugin for PluginEventPublisher { + type Api = PluginEvents; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("event-publisher", "0.1.0") + .with_capability(PluginCapability::PluginEvents) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async move { + context + .plugin_events() + .cloned() + .map(Arc::new) + .ok_or_else(|| anyhow::anyhow!("plugin events capability missing")) + }) + } + } + + #[tokio::test] + async fn typed_plugin_api_publishes_only_to_exact_bounded_routes() { + let client = complete_builder() + .await + .with_plugin(PluginEventPublisher) + .with_plugin(CapabilityProbe) + .build() + .await + .expect("plugin event publisher") + .into_client(); + let publisher = client + .plugin::() + .expect("typed publisher API"); + let router = client.plugin_event_router().expect("plugin event router"); + let tick = PluginEventTopic::new("tick").expect("valid topic"); + let silent_selector = + PluginEventSelector::new("capability-probe", tick.clone()).expect("valid selector"); + assert!(matches!( + router.subscribe( + [silent_selector], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ), + Err(PluginEventSubscribeError::UnknownPublisher { .. }) + )); + let selector = publisher.selector(&tick); + let subscription = router + .subscribe( + [selector.clone()], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ) + .expect("bounded event endpoint"); + + assert!(publisher.has_subscribers(&tick)); + let generation = client.connection_generation.load(Ordering::Acquire); + assert_eq!( + publisher + .publish( + &tick, + 2, + PluginEventPayloadEncoding::Json, + r#"{"messages":1}"#, + ) + .expect("publish tick"), + PluginEventPublishReport { + matched: 1, + enqueued: 1, + dropped: 0, + closed: 0, + } ); + let event = subscription.recv().await.expect("routed tick"); + assert_eq!(&*event.plugin_id, "event-publisher"); + assert_eq!(event.topic, tick); + assert_eq!(event.schema_version, 2); + assert_eq!(event.payload_encoding, PluginEventPayloadEncoding::Json); + assert_eq!(event.payload, Bytes::from_static(br#"{"messages":1}"#)); + assert_eq!(event.connection_generation, generation); + assert_eq!(event.sequence, 1); + + let next_generation = client.connection_generation.fetch_add(1, Ordering::SeqCst) + 1; + publisher + .publish(&tick, 2, PluginEventPayloadEncoding::Json, Bytes::new()) + .expect("publish after generation change"); + let event = subscription.recv().await.expect("next generation tick"); + assert_eq!(event.connection_generation, next_generation); + assert_eq!(event.sequence, 2); + client.disconnect().await; + assert!(matches!( + publisher.publish(&tick, 2, PluginEventPayloadEncoding::Json, Bytes::new(),), + Err(PluginEventPublishError::Resource( + PluginResourceError::ShuttingDown + )) + )); + assert!(matches!( + subscription.recv().await, + Err(PluginEventReceiveError) + )); + assert!(matches!( + router.subscribe( + [selector], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ), + Err(PluginEventSubscribeError::Closed) + )); } } From 57ec05dc170afd4c4b550da15dc3c3187a3ed40d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:39:33 -0300 Subject: [PATCH 19/46] feat(plugins): prove native metrics vertical slice --- Cargo.lock | 13 + Cargo.toml | 1 + plugins/metrics/Cargo.toml | 24 ++ plugins/metrics/src/lib.rs | 483 +++++++++++++++++++++++++++++++++++++ src/lib.rs | 8 +- src/plugins/mod.rs | 77 +++++- 6 files changed, 598 insertions(+), 8 deletions(-) create mode 100644 plugins/metrics/Cargo.toml create mode 100644 plugins/metrics/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 0e4aaf7eb..79ad66dc2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4378,6 +4378,19 @@ dependencies = [ "whatsapp-rust-sqlite-storage", ] +[[package]] +name = "whatsapp-rust-plugin-metrics" +version = "0.1.0" +dependencies = [ + "anyhow", + "bon", + "portable-atomic", + "serde", + "serde_json", + "tokio", + "whatsapp-rust", +] + [[package]] name = "whatsapp-rust-sqlite-storage" version = "0.6.0" diff --git a/Cargo.toml b/Cargo.toml index bd3b6fc6a..0d4ca88d9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ members = [ ".", "examples/voip-cli", "http_clients/ureq-client", + "plugins/metrics", "storages/chat-store", "storages/sqlite-storage", "tests/bench-integration", diff --git a/plugins/metrics/Cargo.toml b/plugins/metrics/Cargo.toml new file mode 100644 index 000000000..a8d01ad8e --- /dev/null +++ b/plugins/metrics/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "whatsapp-rust-plugin-metrics" +version = "0.1.0" +edition = "2024" +publish = false +description = "Reference external metrics plugin for whatsapp-rust" + +[dependencies] +anyhow = { workspace = true } +bon = { workspace = true } +portable-atomic = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +whatsapp-rust = { path = "../..", default-features = false, features = ["plugins"] } + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } +whatsapp-rust = { path = "../..", default-features = false, features = [ + "plugins", + "tokio-runtime", +] } + +[lints] +workspace = true diff --git a/plugins/metrics/src/lib.rs b/plugins/metrics/src/lib.rs new file mode 100644 index 000000000..caa507a9d --- /dev/null +++ b/plugins/metrics/src/lib.rs @@ -0,0 +1,483 @@ +//! Reference out-of-core plugin exercising typed APIs, scoped tasks, and custom events. +//! +//! ```ignore +//! let client = Client::builder() +//! // platform dependencies... +//! .with_plugin(MetricsPlugin::default()) +//! .build() +//! .await? +//! .into_client(); +//! let metrics = client +//! .plugin::() +//! .expect("metrics plugin was registered"); +//! let events = client +//! .plugin_event_router() +//! .expect("metrics publishes custom events") +//! .subscribe( +//! [metrics.tick_selector()], +//! PluginEventEndpointConfig::new(64, PluginEventOverflow::DropOldest), +//! )?; +//! ``` + +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; + +use anyhow::{Context, Result, ensure}; +use portable_atomic::{AtomicBool, AtomicU64, Ordering}; +use serde::{Deserialize, Serialize}; +use whatsapp_rust::wacore::types::events::{Event, EventHandler, EventInterest, EventKind}; +use whatsapp_rust::{ + ClientPlugin, PluginCapability, PluginConnectionScope, PluginEventPayloadEncoding, + PluginEventSelector, PluginEventTopic, PluginEvents, PluginFuture, PluginManifest, PluginTasks, +}; + +pub const METRICS_PLUGIN_ID: &str = "wa.metrics"; +pub const METRICS_TICK_TOPIC: &str = "tick"; +pub const METRICS_TICK_SCHEMA_VERSION: u32 = 1; + +/// Stable snapshot exposed by [`MetricsApi`] and encoded in every tick event. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, bon::Builder)] +#[non_exhaustive] +pub struct MetricsSnapshot { + pub core_events: u64, + pub messages: u64, + pub receipts: u64, + pub connected: u64, + pub disconnected: u64, + pub install_ticks: u64, + pub connection_ticks: u64, + pub ready_scopes: u64, + pub closed_scopes: u64, + pub events_enqueued: u64, + pub events_dropped: u64, + pub publish_failures: u64, + pub active_generation: Option, + pub last_closed_generation: Option, + pub shutdown: bool, +} + +#[derive(Default)] +struct MetricsState { + core_events: AtomicU64, + messages: AtomicU64, + receipts: AtomicU64, + connected: AtomicU64, + disconnected: AtomicU64, + install_ticks: AtomicU64, + connection_ticks: AtomicU64, + ready_scopes: AtomicU64, + closed_scopes: AtomicU64, + events_enqueued: AtomicU64, + events_dropped: AtomicU64, + publish_failures: AtomicU64, + active_generation: Mutex>, + last_closed_generation: Mutex>, + shutdown: AtomicBool, +} + +impl MetricsState { + fn snapshot(&self) -> MetricsSnapshot { + let active_generation = *self + .active_generation + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let last_closed_generation = *self + .last_closed_generation + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + MetricsSnapshot::builder() + .core_events(self.core_events.load(Ordering::Relaxed)) + .messages(self.messages.load(Ordering::Relaxed)) + .receipts(self.receipts.load(Ordering::Relaxed)) + .connected(self.connected.load(Ordering::Relaxed)) + .disconnected(self.disconnected.load(Ordering::Relaxed)) + .install_ticks(self.install_ticks.load(Ordering::Relaxed)) + .connection_ticks(self.connection_ticks.load(Ordering::Relaxed)) + .ready_scopes(self.ready_scopes.load(Ordering::Relaxed)) + .closed_scopes(self.closed_scopes.load(Ordering::Relaxed)) + .events_enqueued(self.events_enqueued.load(Ordering::Relaxed)) + .events_dropped(self.events_dropped.load(Ordering::Relaxed)) + .publish_failures(self.publish_failures.load(Ordering::Relaxed)) + .maybe_active_generation(active_generation) + .maybe_last_closed_generation(last_closed_generation) + .shutdown(self.shutdown.load(Ordering::Acquire)) + .build() + } + + fn open_scope(&self, generation: u64) { + self.ready_scopes.fetch_add(1, Ordering::Relaxed); + *self + .active_generation + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(generation); + } + + fn close_scope(&self, generation: u64) { + let mut active = self + .active_generation + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if *active == Some(generation) { + *active = None; + } + } + + fn record_closed(&self, generation: u64) { + self.closed_scopes.fetch_add(1, Ordering::Relaxed); + *self + .last_closed_generation + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(generation); + self.close_scope(generation); + } +} + +struct MetricsEventHandler(Arc); + +impl EventHandler for MetricsEventHandler { + fn handle_event(&self, event: Arc) { + self.0.core_events.fetch_add(1, Ordering::Relaxed); + match event.kind() { + EventKind::Messages => { + self.0.messages.fetch_add(1, Ordering::Relaxed); + } + EventKind::Receipt => { + self.0.receipts.fetch_add(1, Ordering::Relaxed); + } + EventKind::Connected => { + self.0.connected.fetch_add(1, Ordering::Relaxed); + } + EventKind::Disconnected => { + self.0.disconnected.fetch_add(1, Ordering::Relaxed); + } + _ => {} + } + } +} + +/// Type-safe API returned by `client.plugin::()`. +pub struct MetricsApi { + state: Arc, + tick_selector: PluginEventSelector, +} + +impl MetricsApi { + pub fn snapshot(&self) -> MetricsSnapshot { + self.state.snapshot() + } + + pub fn tick_selector(&self) -> PluginEventSelector { + self.tick_selector.clone() + } +} + +/// One metrics-plugin installation. Construct a fresh value for each client. +pub struct MetricsPlugin { + interval: Duration, + state: OnceLock>, +} + +impl MetricsPlugin { + pub const fn new(interval: Duration) -> Self { + Self { + interval, + state: OnceLock::new(), + } + } + + fn state(&self) -> Result> { + self.state + .get() + .cloned() + .context("metrics plugin is not installed") + } +} + +impl Default for MetricsPlugin { + fn default() -> Self { + Self::new(Duration::from_secs(10)) + } +} + +impl ClientPlugin for MetricsPlugin { + type Api = MetricsApi; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new(METRICS_PLUGIN_ID, env!("CARGO_PKG_VERSION")) + .with_capability(PluginCapability::CoreEvents) + .with_capability(PluginCapability::Tasks) + .with_capability(PluginCapability::PluginEvents) + } + + fn install( + &self, + context: whatsapp_rust::PluginContext, + ) -> PluginFuture<'_, Result>> { + Box::pin(async move { + ensure!(self.interval != Duration::ZERO, "metrics interval is zero"); + let core_events = context + .core_events() + .cloned() + .context("core-events capability is missing")?; + let tasks = context + .tasks() + .cloned() + .context("tasks capability is missing")?; + let plugin_events = context + .plugin_events() + .cloned() + .context("plugin-events capability is missing")?; + let tick = PluginEventTopic::new(METRICS_TICK_TOPIC)?; + let state = Arc::new(MetricsState::default()); + self.state + .set(state.clone()) + .map_err(|_| anyhow::anyhow!("metrics plugin was installed more than once"))?; + + core_events.subscribe( + EventInterest::of(&[ + EventKind::Messages, + EventKind::Receipt, + EventKind::Connected, + EventKind::Disconnected, + ]), + Arc::new(MetricsEventHandler(state.clone())), + )?; + + let api = Arc::new(MetricsApi { + state: state.clone(), + tick_selector: plugin_events.selector(&tick), + }); + spawn_install_ticker(tasks, plugin_events, tick, state, self.interval)?; + Ok(api) + }) + } + + fn on_ready(&self, scope: PluginConnectionScope) -> PluginFuture<'_, Result<()>> { + Box::pin(async move { + let state = self.state()?; + let generation = scope.generation(); + let tasks = scope + .tasks() + .cloned() + .context("connection tasks capability is missing")?; + state.open_scope(generation); + let worker_tasks = tasks.clone(); + let worker_state = state.clone(); + let interval = self.interval; + let guard = ConnectionGuard { state, generation }; + tasks.spawn(async move { + let _guard = guard; + while worker_tasks.sleep(interval).await.is_ok() { + worker_state + .connection_ticks + .fetch_add(1, Ordering::Relaxed); + } + })?; + Ok(()) + }) + } + + fn on_closed(&self, scope: PluginConnectionScope) -> PluginFuture<'_, Result<()>> { + Box::pin(async move { + self.state()?.record_closed(scope.generation()); + Ok(()) + }) + } + + fn shutdown(&self) -> PluginFuture<'_, Result<()>> { + Box::pin(async move { + if let Some(state) = self.state.get() { + state.shutdown.store(true, Ordering::Release); + } + Ok(()) + }) + } +} + +fn spawn_install_ticker( + tasks: PluginTasks, + plugin_events: PluginEvents, + topic: PluginEventTopic, + state: Arc, + interval: Duration, +) -> Result<()> { + let worker_tasks = tasks.clone(); + tasks.spawn(async move { + while worker_tasks.sleep(interval).await.is_ok() { + state.install_ticks.fetch_add(1, Ordering::Relaxed); + if !plugin_events.has_subscribers(&topic) { + continue; + } + let Ok(payload) = serde_json::to_vec(&state.snapshot()) else { + state.publish_failures.fetch_add(1, Ordering::Relaxed); + continue; + }; + match plugin_events.publish( + &topic, + METRICS_TICK_SCHEMA_VERSION, + PluginEventPayloadEncoding::Json, + payload, + ) { + Ok(report) => { + state + .events_enqueued + .fetch_add(report.enqueued, Ordering::Relaxed); + state + .events_dropped + .fetch_add(report.dropped, Ordering::Relaxed); + } + Err(_) => { + state.publish_failures.fetch_add(1, Ordering::Relaxed); + } + } + } + })?; + Ok(()) +} + +struct ConnectionGuard { + state: Arc, + generation: u64, +} + +impl Drop for ConnectionGuard { + fn drop(&mut self) { + self.state.close_scope(self.generation); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::time::Duration; + + use whatsapp_rust::async_channel::Receiver; + use whatsapp_rust::bytes::Bytes; + use whatsapp_rust::http::{HttpClient, HttpRequest, HttpResponse}; + use whatsapp_rust::store::persistence_manager::PersistenceManager; + use whatsapp_rust::transport::{Transport, TransportEvent, TransportFactory}; + use whatsapp_rust::wacore::store::InMemoryBackend; + use whatsapp_rust::{ + Client, PluginEventEndpointConfig, PluginEventOverflow, PluginEventSubscribeError, + TokioRuntime, + }; + + use super::*; + + #[test] + fn retired_scope_cannot_clear_a_newer_generation() { + let state = Arc::new(MetricsState::default()); + state.open_scope(4); + let old_guard = ConnectionGuard { + state: state.clone(), + generation: 4, + }; + state.open_scope(5); + drop(old_guard); + assert_eq!(state.snapshot().active_generation, Some(5)); + + state.record_closed(4); + let snapshot = state.snapshot(); + assert_eq!(snapshot.active_generation, Some(5)); + assert_eq!(snapshot.last_closed_generation, Some(4)); + state.record_closed(5); + assert_eq!(state.snapshot().active_generation, None); + } + + struct TestTransport; + + #[whatsapp_rust::async_trait] + impl Transport for TestTransport { + async fn send(&self, _data: Bytes) -> Result<()> { + Ok(()) + } + + async fn disconnect(&self) {} + } + + struct TestTransportFactory; + + #[whatsapp_rust::async_trait] + impl TransportFactory for TestTransportFactory { + async fn create_transport(&self) -> Result<(Arc, Receiver)> { + let (_sender, receiver) = whatsapp_rust::async_channel::bounded(1); + Ok((Arc::new(TestTransport), receiver)) + } + } + + struct TestHttpClient; + + #[whatsapp_rust::async_trait] + impl HttpClient for TestHttpClient { + async fn execute(&self, _request: HttpRequest) -> Result { + Ok(HttpResponse { + status_code: 200, + body: Vec::new(), + }) + } + } + + async fn test_client(interval: Duration) -> Arc { + let backend = Arc::new(InMemoryBackend::new()); + let persistence = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager"), + ); + Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence) + .with_transport_factory(TestTransportFactory) + .with_http_client(TestHttpClient) + .with_plugin(MetricsPlugin::new(interval)) + .build() + .await + .expect("metrics plugin client") + .into_client() + } + + #[tokio::test] + async fn external_plugin_exposes_typed_api_and_bounded_events() { + let client = test_client(Duration::from_millis(2)).await; + let api = client.plugin::().expect("typed metrics API"); + let router = client.plugin_event_router().expect("plugin event router"); + let selector = api.tick_selector(); + let events = router + .subscribe( + [selector.clone()], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ) + .expect("metrics event endpoint"); + + let event = tokio::time::timeout(Duration::from_secs(1), events.recv()) + .await + .expect("metrics tick timeout") + .expect("metrics tick"); + assert_eq!(&*event.plugin_id, METRICS_PLUGIN_ID); + assert_eq!(event.topic.as_str(), METRICS_TICK_TOPIC); + assert_eq!(event.schema_version, METRICS_TICK_SCHEMA_VERSION); + assert_eq!(event.payload_encoding, PluginEventPayloadEncoding::Json); + let payload: MetricsSnapshot = + serde_json::from_slice(&event.payload).expect("typed metrics payload"); + assert!(payload.install_ticks > 0); + + tokio::time::timeout(Duration::from_secs(1), async { + while events.stats().dropped == 0 { + tokio::time::sleep(Duration::from_millis(2)).await; + } + }) + .await + .expect("bounded endpoint reports pressure"); + assert!(api.snapshot().install_ticks >= payload.install_ticks); + + client.disconnect().await; + assert!(api.snapshot().shutdown); + assert!(matches!( + router.subscribe( + [selector], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ), + Err(PluginEventSubscribeError::Closed) + )); + } +} diff --git a/src/lib.rs b/src/lib.rs index cd13e2286..5a7451e11 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -120,9 +120,9 @@ pub use plugins::{ PluginEventEndpointStats, PluginEventEnvelope, PluginEventOverflow, PluginEventPayloadEncoding, PluginEventPublishError, PluginEventPublishReport, PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, PluginEventSelector, PluginEventSubscribeError, - PluginEventSubscription, PluginEventTopic, PluginEventTryReceiveError, PluginEvents, PluginIq, - PluginIqError, PluginManifest, PluginMessaging, PluginMessagingError, PluginPlanError, - PluginResourceError, PluginTasks, + PluginEventSubscription, PluginEventTopic, PluginEventTryReceiveError, PluginEvents, + PluginFuture, PluginIq, PluginIqError, PluginManifest, PluginMessaging, PluginMessagingError, + PluginPlanError, PluginResourceError, PluginTasks, }; pub mod request; pub(crate) mod signal_flush; @@ -201,7 +201,7 @@ pub mod prelude { ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, PluginEventEndpointConfig, PluginEventOverflow, PluginEventPayloadEncoding, PluginEventRouter, PluginEventSelector, PluginEventSubscription, PluginEventTopic, - PluginEvents, PluginManifest, + PluginEvents, PluginFuture, PluginManifest, }; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index e260108dc..c351229f8 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -139,25 +139,29 @@ impl PluginManifest { } } +/// Target-correct future returned by native plugin entry points. +pub type PluginFuture<'a, T> = BoxFuture<'a, T>; + /// A trusted native plugin installed exactly once while the client is still inert. /// Capabilities shape the handles it receives; they are not an in-process sandbox. +/// A plugin value belongs to one client installation, even when registered through an `Arc`. pub trait ClientPlugin: MaybeSendSync + 'static { type Api: MaybeSendSync + 'static; fn manifest(&self) -> PluginManifest; - fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>>; + fn install(&self, context: PluginContext) -> PluginFuture<'_, anyhow::Result>>; - fn on_ready(&self, _scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + fn on_ready(&self, _scope: PluginConnectionScope) -> PluginFuture<'_, anyhow::Result<()>> { Box::pin(async { Ok(()) }) } - fn on_closed(&self, _scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + fn on_closed(&self, _scope: PluginConnectionScope) -> PluginFuture<'_, anyhow::Result<()>> { Box::pin(async { Ok(()) }) } /// Release plugin-owned state. This may run after `install` began but returned an error. - fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + fn shutdown(&self) -> PluginFuture<'_, anyhow::Result<()>> { Box::pin(async { Ok(()) }) } } @@ -330,6 +334,17 @@ impl PluginTasks { pub fn shutdown_signal(&self) -> ShutdownSignal { self.resources.shutdown.subscribe() } + + /// Sleep through the configured runtime, returning promptly when this plugin shuts down. + pub async fn sleep(&self, duration: Duration) -> Result<(), PluginResourceError> { + self.resources.ensure_active()?; + let shutdown = self.resources.shutdown.subscribe(); + let cancelled = Box::pin(wait_for_shutdown(&shutdown)); + match futures::future::select(cancelled, self.runtime.sleep(duration)).await { + futures::future::Either::Left(_) => Err(PluginResourceError::ShuttingDown), + futures::future::Either::Right(_) => self.resources.ensure_active(), + } + } } /// Selective subscription access to the sealed core event bus. @@ -506,6 +521,22 @@ impl PluginConnectionTasks { spawn_until_cancelled(&self.runtime, self.scope.cancellation_signal(), future); Ok(()) } + + /// Sleep through the configured runtime, returning promptly when this generation retires. + pub async fn sleep(&self, duration: Duration) -> Result<(), PluginResourceError> { + if self.scope.is_cancelled() { + return Err(PluginResourceError::ShuttingDown); + } + let cancellation = self.scope.cancellation_signal(); + let cancelled = Box::pin(wait_for_shutdown(&cancellation)); + match futures::future::select(cancelled, self.runtime.sleep(duration)).await { + futures::future::Either::Left(_) => Err(PluginResourceError::ShuttingDown), + futures::future::Either::Right(_) if self.scope.is_cancelled() => { + Err(PluginResourceError::ShuttingDown) + } + futures::future::Either::Right(_) => Ok(()), + } + } } fn spawn_until_cancelled(runtime: &Arc, cancellation: ShutdownSignal, future: F) @@ -2214,6 +2245,44 @@ mod tests { )); } + #[tokio::test] + async fn task_sleeps_return_when_their_owner_is_cancelled() { + let resources = PluginResources::new(); + resources.activate(); + let install_tasks = PluginTasks { + runtime: Arc::new(TokioRuntime), + resources: resources.clone(), + }; + let install_sleeper = + tokio::spawn(async move { install_tasks.sleep(Duration::from_secs(60)).await }); + tokio::task::yield_now().await; + resources.close(); + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), install_sleeper) + .await + .expect("install sleep cancellation") + .expect("install sleeper task"), + Err(PluginResourceError::ShuttingDown) + ); + + let scope = ConnectionScope::new(91); + let connection_tasks = PluginConnectionTasks { + runtime: Arc::new(TokioRuntime), + scope: scope.clone(), + }; + let connection_sleeper = + tokio::spawn(async move { connection_tasks.sleep(Duration::from_secs(60)).await }); + tokio::task::yield_now().await; + scope.cancel(); + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), connection_sleeper) + .await + .expect("connection sleep cancellation") + .expect("connection sleeper task"), + Err(PluginResourceError::ShuttingDown) + ); + } + struct UpstreamLifecycle { log: Log, } From 06050d0c86210bb44b3a176acc60f68fa6b2990b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:02:25 -0300 Subject: [PATCH 20/46] fix(lifecycle): guarantee terminal cleanup signaling --- src/client/device_registry.rs | 5 +- src/client/extension_lifecycle.rs | 77 +++++++++++++++++++++++++++++++ src/client/lifecycle.rs | 25 +++++++--- 3 files changed, 97 insertions(+), 10 deletions(-) diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index e0441f7d7..ef5c29942 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -964,9 +964,8 @@ impl Client { /// Background loop placeholder for device registry cleanup. /// Note: Cleanup functionality was removed as part of trait simplification. /// Device registry entries are managed through normal update/get operations. - pub(super) async fn device_registry_cleanup_loop(&self) { - // Simply wait for shutdown signal - self.shutdown_notifier.listen().await; + pub(super) async fn device_registry_cleanup_loop(shutdown: wacore::runtime::ShutdownSignal) { + wacore::runtime::wait_for_shutdown(&shutdown).await; debug!( target: "Client/DeviceRegistry", "Shutdown signaled, exiting cleanup loop" diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 8f01f70f7..e34e039bc 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -702,6 +702,19 @@ mod tests { } } + struct PanickingDisconnect; + + #[async_trait] + impl crate::transport::Transport for PanickingDisconnect { + async fn send(&self, _data: Bytes) -> anyhow::Result<()> { + Ok(()) + } + + async fn disconnect(&self) { + panic!("injected disconnect panic"); + } + } + struct BlockingReadyLifecycle { ready_started: async_channel::Sender<()>, release_ready: async_channel::Receiver<()>, @@ -1126,6 +1139,38 @@ mod tests { .expect("detached lifecycle shutdown completed"); } + #[tokio::test] + async fn dropping_last_client_owner_signals_standalone_lifecycle() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let lifecycle = Arc::new(EarlyShutdownLifecycle::default()); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await + .expect("client build") + .into_client(); + let weak = Arc::downgrade(&client); + + drop(client); + + tokio::time::timeout(Duration::from_secs(2), async { + while weak.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .expect("background services released the client"); + assert!(lifecycle.signalled.load(Ordering::Acquire)); + } + #[tokio::test] async fn cancellation_does_not_wait_for_a_running_callback() { let persistence_manager = Arc::new( @@ -1464,6 +1509,38 @@ mod tests { client.signal_shutdown_sync(); } + #[tokio::test] + async fn detached_cleanup_propagates_panics_to_its_waiter() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle(RecordingLifecycle::default()) + .build() + .await + .expect("client build") + .into_client(); + *client.transport.lock().await = Some(Arc::new(PanickingDisconnect)); + + let cleanup_client = Arc::clone(&client); + let cleanup = tokio::spawn(async move { + cleanup_client.cleanup_connection_state().await; + }); + let panic = tokio::time::timeout(Duration::from_secs(2), cleanup) + .await + .expect("cleanup waiter did not hang") + .expect_err("cleanup panic should reach its waiter"); + + assert!(panic.is_panic()); + client.signal_shutdown_sync(); + } + #[tokio::test] async fn stale_generation_is_rejected_before_scope_publication() { let lifecycle = Arc::new(RecordingLifecycle::default()); diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 52fbec60b..50c5f359d 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -6,6 +6,12 @@ use super::*; /// accounts in more groups; an evicted entry just recomputes on next send. const GROUP_DEVICES_MEMO_CAPACITY: u64 = 64; +impl Drop for Client { + fn drop(&mut self) { + self.signal_shutdown_sync(); + } +} + impl Client { /// WA Web `resetDelay: 30000` — only after a connection has stayed up this /// long is the reconnect backoff counter reset to its base. @@ -387,10 +393,10 @@ impl Client { })) .detach(); - let cleanup_arc = self.clone(); + let cleanup_shutdown = self.shutdown_signal(); self.runtime .spawn(Box::pin(async move { - cleanup_arc.device_registry_cleanup_loop().await; + Client::device_registry_cleanup_loop(cleanup_shutdown).await; })) .detach(); } @@ -843,16 +849,21 @@ impl Client { } // Scope closure must survive a caller dropping its cleanup waiter. - let completed = wacore::runtime::ShutdownNotifier::new(); - let completion = completed.subscribe(); + let (completed, completion) = futures::channel::oneshot::channel(); let client = Arc::clone(self); self.runtime .spawn(Box::pin(async move { - client.cleanup_connection_state_inner().await; - completed.notify(); + let result = std::panic::AssertUnwindSafe(client.cleanup_connection_state_inner()) + .catch_unwind() + .await; + let _ = completed.send(result); })) .detach(); - wacore::runtime::wait_for_shutdown(&completion).await; + match completion.await { + Ok(Ok(())) => {} + Ok(Err(panic)) => std::panic::resume_unwind(panic), + Err(_) => error!("Detached connection cleanup stopped before completion"), + } } async fn cleanup_connection_state_inner(&self) { From 5b4a3b4fb30fe740c94db4b935cb76f99c766ba5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:21:36 -0300 Subject: [PATCH 21/46] fix(plugins): drain scoped tasks before teardown --- src/plugins/mod.rs | 373 +++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 346 insertions(+), 27 deletions(-) diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index c351229f8..18cefbb73 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -204,6 +204,8 @@ pub enum PluginResourceError { NotActive, #[error("the plugin scope is shutting down")] ShuttingDown, + #[error("the plugin task capacity is exhausted")] + TaskCapacityExceeded, } #[derive(Debug, Error)] @@ -229,9 +231,99 @@ struct PluginResources { closed: AtomicBool, activation: ShutdownNotifier, shutdown: ShutdownNotifier, + install_tasks: Arc, + connection_tasks: Mutex, subscriptions: Mutex>, } +#[derive(Default)] +struct ConnectionTaskRegistry { + closed: bool, + trackers: HashMap>, +} + +#[derive(Default)] +struct TaskTrackerState { + active: usize, + closed: bool, +} + +struct TaskTracker { + state: Mutex, + idle: ShutdownNotifier, +} + +impl TaskTracker { + fn new() -> Arc { + Arc::new(Self { + state: Mutex::new(TaskTrackerState::default()), + idle: ShutdownNotifier::new(), + }) + } + + fn closed() -> Arc { + let tracker = Self::new(); + tracker.close(); + tracker + } + + fn register(self: &Arc) -> Result { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if state.closed { + return Err(PluginResourceError::ShuttingDown); + } + state.active = state + .active + .checked_add(1) + .ok_or(PluginResourceError::TaskCapacityExceeded)?; + Ok(TaskLease { + tracker: Arc::clone(self), + }) + } + + fn close(&self) { + let idle = { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.closed = true; + state.active == 0 + }; + if idle { + self.idle.notify(); + } + } + + fn completion_signal(&self) -> ShutdownSignal { + self.idle.subscribe() + } +} + +struct TaskLease { + tracker: Arc, +} + +impl Drop for TaskLease { + fn drop(&mut self) { + let idle = { + let mut state = self + .tracker + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.active = state.active.saturating_sub(1); + state.closed && state.active == 0 + }; + if idle { + self.tracker.idle.notify(); + } + } +} + struct PluginCoreEventSubscription { _subscription: Subscription, _raw_node_lease: Option, @@ -244,6 +336,8 @@ impl PluginResources { closed: AtomicBool::new(false), activation: ShutdownNotifier::new(), shutdown: ShutdownNotifier::new(), + install_tasks: TaskTracker::new(), + connection_tasks: Mutex::new(ConnectionTaskRegistry::default()), subscriptions: Mutex::new(Vec::new()), }) } @@ -295,10 +389,78 @@ impl PluginResources { } } + fn connection_task_tracker(&self, generation: u64) -> Arc { + let mut registry = self + .connection_tasks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if registry.closed { + return TaskTracker::closed(); + } + Arc::clone( + registry + .trackers + .entry(generation) + .or_insert_with(TaskTracker::new), + ) + } + + fn close_connection_tasks(&self, generation: u64) -> Arc { + let tracker = Arc::clone( + self.connection_tasks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .trackers + .entry(generation) + .or_insert_with(TaskTracker::closed), + ); + tracker.close(); + tracker + } + + fn forget_connection_tasks(&self, generation: u64, tracker: &Arc) { + let mut registry = self + .connection_tasks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if registry + .trackers + .get(&generation) + .is_some_and(|current| Arc::ptr_eq(current, tracker)) + { + registry.trackers.remove(&generation); + } + } + + fn task_completion_signals(&self) -> Vec { + let mut signals = vec![self.install_tasks.completion_signal()]; + signals.extend( + self.connection_tasks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .trackers + .values() + .map(|tracker| tracker.completion_signal()), + ); + signals + } + fn close(&self) { if self.closed.swap(true, Ordering::AcqRel) { return; } + self.install_tasks.close(); + let connection_trackers = { + let mut registry = self + .connection_tasks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + registry.closed = true; + registry.trackers.values().cloned().collect::>() + }; + for tracker in connection_trackers { + tracker.close(); + } self.shutdown.notify(); let subscriptions = { let mut subscriptions = self @@ -327,7 +489,8 @@ impl PluginTasks { if self.resources.closed.load(Ordering::Acquire) { return Err(PluginResourceError::ShuttingDown); } - spawn_after_activation(&self.runtime, Arc::clone(&self.resources), future); + let lease = self.resources.install_tasks.register()?; + spawn_after_activation(&self.runtime, Arc::clone(&self.resources), lease, future); Ok(()) } @@ -508,6 +671,7 @@ impl PluginConnectionScope { pub struct PluginConnectionTasks { runtime: Arc, scope: ConnectionScope, + tracker: Arc, } impl PluginConnectionTasks { @@ -518,7 +682,13 @@ impl PluginConnectionTasks { if self.scope.is_cancelled() { return Err(PluginResourceError::ShuttingDown); } - spawn_until_cancelled(&self.runtime, self.scope.cancellation_signal(), future); + let lease = self.tracker.register()?; + spawn_until_cancelled( + &self.runtime, + self.scope.cancellation_signal(), + lease, + future, + ); Ok(()) } @@ -539,12 +709,17 @@ impl PluginConnectionTasks { } } -fn spawn_until_cancelled(runtime: &Arc, cancellation: ShutdownSignal, future: F) -where +fn spawn_until_cancelled( + runtime: &Arc, + cancellation: ShutdownSignal, + lease: TaskLease, + future: F, +) where F: Future + Spawnable, { runtime .spawn(Box::pin(async move { + let _lease = lease; let cancelled = Box::pin(wait_for_shutdown(&cancellation)); let work = Box::pin(future); let _ = futures::future::select(cancelled, work).await; @@ -552,14 +727,19 @@ where .detach(); } -fn spawn_after_activation(runtime: &Arc, resources: Arc, future: F) -where +fn spawn_after_activation( + runtime: &Arc, + resources: Arc, + lease: TaskLease, + future: F, +) where F: Future + Spawnable, { let activation = resources.activation.subscribe(); let cancellation = resources.shutdown.subscribe(); runtime .spawn(Box::pin(async move { + let _lease = lease; let cancelled = Box::pin(wait_for_shutdown(&cancellation)); let activated = Box::pin(wait_for_shutdown(&activation)); if matches!( @@ -990,8 +1170,18 @@ impl PluginHost { pub(crate) fn lifecycle_callback_timeout(&self) -> Duration { let callback_count = self.ordered.len() + usize::from(self.upstream.is_some()); + let task_barrier_count = self + .ordered + .iter() + .filter(|plugin| { + plugin + .manifest + .capabilities + .contains(PluginCapability::Tasks) + }) + .count(); self.callback_timeout - .saturating_mul(callback_count as u32) + .saturating_mul((callback_count + task_barrier_count) as u32) .saturating_add(Duration::from_secs(1)) } @@ -1052,19 +1242,32 @@ impl PluginHost { &self, scope: ConnectionScope, manifest: &PluginManifest, + task_tracker: Option>, ) -> PluginConnectionScope { - let tasks = manifest - .capabilities - .contains(PluginCapability::Tasks) - .then(|| self.runtime.get().cloned()) - .flatten() - .map(|runtime| PluginConnectionTasks { - runtime, - scope: scope.clone(), - }); + let tasks = if manifest.capabilities.contains(PluginCapability::Tasks) { + self.runtime + .get() + .cloned() + .zip(task_tracker) + .map(|(runtime, tracker)| PluginConnectionTasks { + runtime, + scope: scope.clone(), + tracker, + }) + } else { + None + }; PluginConnectionScope { scope, tasks } } + async fn wait_for_tasks(&self, completion_signals: Vec) -> anyhow::Result<()> { + let runtime = self + .runtime + .get() + .ok_or_else(|| anyhow::anyhow!("plugin runtime is unavailable"))?; + wait_for_plugin_tasks(&**runtime, self.callback_timeout, completion_signals).await + } + async fn install_all(&self, client: Weak) -> anyhow::Result<()> { let Some(strong_client) = client.upgrade() else { anyhow::bail!("client was dropped during plugin installation"); @@ -1163,7 +1366,13 @@ impl ClientLifecycle for PluginHost { failures.push(format!("upstream: {error:#}")); } for plugin in self.installed.get().into_iter().flatten() { - let plugin_scope = self.connection_scope(scope.clone(), &plugin.manifest); + let task_tracker = plugin + .manifest + .capabilities + .contains(PluginCapability::Tasks) + .then(|| plugin.resources.connection_task_tracker(scope.generation())); + let plugin_scope = + self.connection_scope(scope.clone(), &plugin.manifest, task_tracker); if let Err(error) = self .run_callback(|| plugin.plugin.on_ready(plugin_scope)) .await @@ -1179,7 +1388,26 @@ impl ClientLifecycle for PluginHost { Box::pin(async move { let mut failures = Vec::new(); for plugin in self.installed.get().into_iter().flatten().rev() { - let plugin_scope = self.connection_scope(scope.clone(), &plugin.manifest); + let task_tracker = plugin + .manifest + .capabilities + .contains(PluginCapability::Tasks) + .then(|| plugin.resources.close_connection_tasks(scope.generation())); + if let Some(task_tracker) = &task_tracker { + match self + .wait_for_tasks(vec![task_tracker.completion_signal()]) + .await + { + Ok(()) => plugin + .resources + .forget_connection_tasks(scope.generation(), task_tracker), + Err(error) => { + failures.push(format!("{} tasks: {error:#}", plugin.manifest.id)); + } + } + } + let plugin_scope = + self.connection_scope(scope.clone(), &plugin.manifest, task_tracker); if let Err(error) = self .run_callback(|| plugin.plugin.on_closed(plugin_scope)) .await @@ -1215,6 +1443,12 @@ impl ClientLifecycle for PluginHost { let mut failures = Vec::new(); self.signal_shutdown(); for plugin in self.installed.get().into_iter().flatten().rev() { + if let Err(error) = self + .wait_for_tasks(plugin.resources.task_completion_signals()) + .await + { + failures.push(format!("{} tasks: {error:#}", plugin.manifest.id)); + } if let Err(error) = self.run_callback(|| plugin.plugin.shutdown()).await { failures.push(format!("{}: {error:#}", plugin.manifest.id)); } @@ -1241,18 +1475,43 @@ async fn shutdown_staged_plugins( mut installed: Vec, upstream: Option>, ) { - if let Some(plugin) = current - && let Err(error) = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { + if let Some(plugin) = current { + if let Err(error) = wait_for_plugin_tasks( + &*runtime, + PLUGIN_CALLBACK_TIMEOUT, + plugin.resources.task_completion_signals(), + ) + .await + { + log::warn!( + "Plugin `{}` failed-install task cleanup failed: {error:#}", + plugin.manifest.id + ); + } + if let Err(error) = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { plugin.plugin.shutdown() }) .await - { - log::warn!( - "Plugin `{}` failed-install rollback failed: {error:#}", - plugin.manifest.id - ); + { + log::warn!( + "Plugin `{}` failed-install rollback failed: {error:#}", + plugin.manifest.id + ); + } } while let Some(plugin) = installed.pop() { + if let Err(error) = wait_for_plugin_tasks( + &*runtime, + PLUGIN_CALLBACK_TIMEOUT, + plugin.resources.task_completion_signals(), + ) + .await + { + log::warn!( + "Plugin `{}` rollback task cleanup failed: {error:#}", + plugin.manifest.id + ); + } if let Err(error) = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { plugin.plugin.shutdown() }) @@ -1270,6 +1529,26 @@ async fn shutdown_staged_plugins( } } +async fn wait_for_plugin_tasks( + runtime: &dyn Runtime, + timeout: Duration, + completion_signals: Vec, +) -> anyhow::Result<()> { + let wait_for_all = async move { + for signal in completion_signals { + wait_for_shutdown(&signal).await; + } + }; + runtime_timeout(runtime, timeout, wait_for_all) + .await + .map_err(|_| { + anyhow::anyhow!( + "plugin tasks did not stop within {:.3} seconds", + timeout.as_secs_f64() + ) + }) +} + async fn bounded_plugin_callback<'a>( runtime: &dyn Runtime, timeout: Duration, @@ -1766,7 +2045,12 @@ mod tests { fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { let log = self.log.clone(); + let task_dropped = self.task_dropped.clone(); Box::pin(async move { + anyhow::ensure!( + task_dropped.load(Ordering::Acquire), + "rollback task still running during shutdown" + ); record(&log, "shutdown:rollback"); Ok(()) }) @@ -2123,6 +2407,8 @@ mod tests { install_dropped: Arc, connection_started: Arc, connection_dropped: Arc, + closed_after_task: Arc, + shutdown_after_task: Arc, } impl ClientPlugin for ScopedTaskPlugin { @@ -2166,6 +2452,24 @@ mod tests { Ok(()) }) } + + fn on_closed(&self, _scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + let task_dropped = self.connection_dropped.load(Ordering::Acquire); + let closed_after_task = self.closed_after_task.clone(); + Box::pin(async move { + closed_after_task.store(task_dropped, Ordering::Release); + Ok(()) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let task_dropped = self.install_dropped.load(Ordering::Acquire); + let shutdown_after_task = self.shutdown_after_task.clone(); + Box::pin(async move { + shutdown_after_task.store(task_dropped, Ordering::Release); + Ok(()) + }) + } } async fn wait_for_flag(flag: &AtomicBool) { @@ -2184,6 +2488,8 @@ mod tests { let install_dropped = Arc::new(AtomicBool::new(false)); let connection_started = Arc::new(AtomicBool::new(false)); let connection_dropped = Arc::new(AtomicBool::new(false)); + let closed_after_task = Arc::new(AtomicBool::new(false)); + let shutdown_after_task = Arc::new(AtomicBool::new(false)); let build = complete_builder() .await .with_plugin(ScopedTaskPlugin { @@ -2191,6 +2497,8 @@ mod tests { install_dropped: install_dropped.clone(), connection_started: connection_started.clone(), connection_dropped: connection_dropped.clone(), + closed_after_task: closed_after_task.clone(), + shutdown_after_task: shutdown_after_task.clone(), }) .build() .await @@ -2208,11 +2516,20 @@ mod tests { .expect("plugin ready callback"); wait_for_flag(&connection_started).await; scope.cancel(); - wait_for_flag(&connection_dropped).await; + client + .plugin_host + .as_ref() + .expect("plugin host") + .on_closed(scope) + .await + .expect("plugin closed callback"); + assert!(connection_dropped.load(Ordering::Acquire)); + assert!(closed_after_task.load(Ordering::Acquire)); assert!(!install_dropped.load(Ordering::Acquire)); client.disconnect().await; - wait_for_flag(&install_dropped).await; + assert!(install_dropped.load(Ordering::Acquire)); + assert!(shutdown_after_task.load(Ordering::Acquire)); } #[tokio::test] @@ -2222,6 +2539,7 @@ mod tests { let tasks = PluginConnectionTasks { runtime: Arc::new(TokioRuntime), scope: scope.clone(), + tracker: TaskTracker::new(), }; let guard = DropFlag(task_dropped.clone()); tasks @@ -2269,6 +2587,7 @@ mod tests { let connection_tasks = PluginConnectionTasks { runtime: Arc::new(TokioRuntime), scope: scope.clone(), + tracker: TaskTracker::new(), }; let connection_sleeper = tokio::spawn(async move { connection_tasks.sleep(Duration::from_secs(60)).await }); From 47dcf10a8489f84c50f684aff3193ca4c6cdde9f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:33:05 -0300 Subject: [PATCH 22/46] fix(lifecycle): bound and compact callback backlog --- src/client/extension_lifecycle.rs | 199 ++++++++++++++++++++++++++---- 1 file changed, 172 insertions(+), 27 deletions(-) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index e34e039bc..d3fa1a7d9 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -16,6 +16,7 @@ const SCOPE_READY: u8 = 1; const SCOPE_CANCELLED: u8 = 2; const SCOPE_CLOSED: u8 = 3; const CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); +const CALLBACK_QUEUE_CAPACITY: usize = 64; std::thread_local! { static ACTIVE_CALLBACK: Cell<*const LifecycleRegistration> = const { Cell::new(std::ptr::null()) }; @@ -189,6 +190,37 @@ struct CallbackQueue { shutdown_requested: bool, shutdown_enqueued: bool, drain_scheduled: bool, + overflowed: bool, +} + +impl CallbackQueue { + fn push_bounded(&mut self, callback: LifecycleCallback) -> Option { + let dropped = (self.pending.len() >= CALLBACK_QUEUE_CAPACITY) + .then(|| self.pending.pop_front()) + .flatten(); + self.overflowed |= dropped.is_some(); + self.pending.push_back(callback); + dropped + } + + fn compact_for_shutdown(&mut self, keep_last_closed: bool) -> Vec { + if !self.overflowed { + return Vec::new(); + } + let last_closed = keep_last_closed + .then(|| { + self.pending + .iter() + .rposition(|callback| matches!(callback, LifecycleCallback::Closed(_))) + }) + .flatten() + .and_then(|position| self.pending.remove(position)); + let dropped = self.pending.drain(..).collect::>(); + if let Some(last_closed) = last_closed { + self.pending.push_back(last_closed); + } + dropped + } } struct CallbackContextGuard { @@ -356,7 +388,7 @@ impl LifecycleRegistration { /// Non-noop hooks are test-only, must run off-executor, and must not re-enter lifecycle APIs. fn close_scope_with(self: &Arc, generation: u64, after_remove: impl FnOnce()) { - let should_spawn = { + let (should_spawn, dropped) = { let mut scopes = self.scopes(); let scope = if scopes .active @@ -382,19 +414,23 @@ impl LifecycleRegistration { // shutdown cannot overtake the final on_closed callback. let no_open_scopes = scopes.active.is_none() && scopes.retired.is_empty(); let mut queue = self.callback_queue(); - queue.pending.push_back(LifecycleCallback::Closed(scope)); - if no_open_scopes && queue.shutdown_requested && !queue.shutdown_enqueued { - queue.shutdown_enqueued = true; - queue.pending.push_back(LifecycleCallback::Shutdown); - } - if queue.drain_scheduled { - false + let mut dropped = Vec::new(); + if queue.shutdown_requested { + dropped.extend(queue.push_bounded(LifecycleCallback::Closed(scope))); + dropped.extend(queue.compact_for_shutdown(no_open_scopes)); + if no_open_scopes && !queue.shutdown_enqueued { + queue.shutdown_enqueued = true; + queue.pending.push_back(LifecycleCallback::Shutdown); + } } else { - queue.drain_scheduled = true; - true + dropped.extend(queue.push_bounded(LifecycleCallback::Closed(scope))); } + let should_spawn = !queue.pending.is_empty() && !queue.drain_scheduled; + queue.drain_scheduled |= should_spawn; + (should_spawn, dropped) }; + warn_dropped_callbacks(dropped); self.spawn_callback_driver(should_spawn); } @@ -439,16 +475,18 @@ impl LifecycleRegistration { } fn enqueue_callback(self: &Arc, callback: LifecycleCallback) { - let should_spawn = { + let (should_spawn, dropped) = { let mut queue = self.callback_queue(); - queue.pending.push_back(callback); - if queue.drain_scheduled { - false + if queue.shutdown_requested || self.terminal.load(Ordering::Acquire) { + (false, Some(callback)) } else { + let dropped = queue.push_bounded(callback); + let should_spawn = !queue.drain_scheduled; queue.drain_scheduled = true; - true + (should_spawn, dropped) } }; + warn_dropped_callbacks(dropped.into_iter().collect()); self.spawn_callback_driver(should_spawn); } @@ -457,24 +495,21 @@ impl LifecycleRegistration { let scopes = self.scopes(); scopes.active.is_none() && scopes.retired.is_empty() }; - if !no_open_scopes { - return; - } - - let should_spawn = { + let (should_spawn, dropped) = { let mut queue = self.callback_queue(); if !queue.shutdown_requested || queue.shutdown_enqueued { return; } - queue.shutdown_enqueued = true; - queue.pending.push_back(LifecycleCallback::Shutdown); - if queue.drain_scheduled { - false - } else { - queue.drain_scheduled = true; - true + let dropped = queue.compact_for_shutdown(no_open_scopes); + if no_open_scopes { + queue.shutdown_enqueued = true; + queue.pending.push_back(LifecycleCallback::Shutdown); } + let should_spawn = !queue.pending.is_empty() && !queue.drain_scheduled; + queue.drain_scheduled |= should_spawn; + (should_spawn, dropped) }; + warn_dropped_callbacks(dropped); self.spawn_callback_driver(should_spawn); } @@ -498,6 +533,7 @@ impl LifecycleRegistration { Some(callback) => callback, None => { queue.drain_scheduled = false; + queue.overflowed = false; return; } } @@ -608,6 +644,15 @@ impl LifecycleRegistration { } } +fn warn_dropped_callbacks(dropped: Vec) { + if !dropped.is_empty() { + log::warn!( + "Dropped {} stale client lifecycle callback(s) under queue pressure or terminal shutdown", + dropped.len() + ); + } +} + #[cfg(test)] mod tests { use std::sync::atomic::AtomicUsize; @@ -789,6 +834,37 @@ mod tests { completed: AtomicBool, } + struct QueuePressureLifecycle { + ready_started: async_channel::Sender<()>, + release_ready: async_channel::Receiver<()>, + closed_calls: AtomicUsize, + shutdown_calls: AtomicUsize, + } + + impl ClientLifecycle for QueuePressureLifecycle { + fn on_ready<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + let _ = self.ready_started.try_send(()); + let _ = self.release_ready.recv().await; + Ok(()) + }) + } + + fn on_closed<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + self.closed_calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + self.shutdown_calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + } + } + impl ClientLifecycle for BlockingShutdownLifecycle { fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { Box::pin(async move { @@ -1355,6 +1431,75 @@ mod tests { ); } + #[tokio::test] + async fn callback_queue_is_bounded_and_terminal_shutdown_compacts_backlog() { + let (ready_started_tx, ready_started_rx) = async_channel::bounded(1); + let (release_ready_tx, release_ready_rx) = async_channel::bounded(1); + let lifecycle = Arc::new(QueuePressureLifecycle { + ready_started: ready_started_tx, + release_ready: release_ready_rx, + closed_calls: AtomicUsize::new(0), + shutdown_calls: AtomicUsize::new(0), + }); + let registration = Arc::new(LifecycleRegistration::new_with_timeout( + lifecycle.clone(), + Arc::new(TokioRuntime), + Duration::from_secs(1), + )); + + let active_scope = ConnectionScope::new(1); + assert!(active_scope.mark_ready()); + let (done, _done_rx) = async_channel::bounded(1); + registration.enqueue_callback(LifecycleCallback::Ready { + scope: active_scope, + done, + }); + ready_started_rx + .recv() + .await + .expect("ready callback started"); + + for generation in 2..(CALLBACK_QUEUE_CAPACITY as u64 * 4) { + let scope = ConnectionScope::new(generation); + scope.close(); + registration.enqueue_callback(LifecycleCallback::Closed(scope)); + } + assert_eq!( + registration.callback_queue().pending.len(), + CALLBACK_QUEUE_CAPACITY + ); + + let shutdown_registration = registration.clone(); + let shutdown = tokio::spawn(async move { shutdown_registration.shutdown().await }); + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let compacted = { + let queue = registration.callback_queue(); + if queue.shutdown_enqueued { + assert_eq!(queue.pending.len(), 2); + true + } else { + false + } + }; + if compacted { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("terminal backlog compaction"); + + release_ready_tx + .send(()) + .await + .expect("release ready callback"); + shutdown.await.expect("bounded terminal shutdown"); + assert_eq!(lifecycle.closed_calls.load(Ordering::SeqCst), 1); + assert_eq!(lifecycle.shutdown_calls.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn cancelled_shutdown_waiter_does_not_cancel_shutdown() { let (started_tx, started_rx) = async_channel::bounded(1); From 1ee952b9c9714e9e78fad214cbfc53c31342a2f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:36:52 -0300 Subject: [PATCH 23/46] fix(lifecycle): retire scopes when reconnect starts --- src/client/extension_lifecycle.rs | 130 ++++++++++++++++++++++++++++++ src/client/lifecycle.rs | 6 ++ 2 files changed, 136 insertions(+) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index d3fa1a7d9..fa7d06090 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -382,6 +382,15 @@ impl LifecycleRegistration { } } + pub(super) fn cancel_active_scope(&self) { + if ready_publication_active(self) { + self.cancel_active_scope_inner(); + } else { + let _publication = self.ready_publication(); + self.cancel_active_scope_inner(); + } + } + pub(super) fn close_scope(self: &Arc, generation: u64) { self.close_scope_with(generation, || {}); } @@ -616,6 +625,12 @@ impl LifecycleRegistration { } } + fn cancel_active_scope_inner(&self) { + if let Some(scope) = &self.scopes().active { + scope.cancel(); + } + } + fn mark_terminal_and_cancel_scopes(&self) -> bool { let first_signal = !self.terminal.swap(true, Ordering::AcqRel); let scopes = self.scopes(); @@ -773,6 +788,70 @@ mod tests { events: std::sync::Mutex>, } + struct ReentrantReconnectLifecycle { + client: std::sync::Mutex>>, + events: std::sync::Mutex>, + immediate: bool, + } + + impl ClientLifecycle for ReentrantReconnectLifecycle { + fn install<'a>(&'a self, client: Weak) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + *self + .client + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(client); + Ok(()) + }) + } + + fn on_ready<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("ready-started"); + let client = self + .client + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .as_ref() + .and_then(Weak::upgrade) + .expect("installed client"); + if self.immediate { + client.reconnect_immediately().await; + } else { + client.reconnect().await; + } + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("ready-finished"); + Ok(()) + }) + } + + fn on_closed<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("closed"); + Ok(()) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + self.events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push("shutdown"); + Ok(()) + }) + } + } + impl ClientLifecycle for ReentrantDisconnectLifecycle { fn install<'a>(&'a self, client: Weak) -> BoxFuture<'a, anyhow::Result<()>> { Box::pin(async move { @@ -1378,6 +1457,57 @@ mod tests { assert!(!client.is_ready.load(Ordering::Relaxed)); } + #[tokio::test] + async fn ready_callback_reconnect_requests_retire_the_scope_before_returning() { + for immediate in [false, true] { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let lifecycle = Arc::new(ReentrantReconnectLifecycle { + client: std::sync::Mutex::new(None), + events: std::sync::Mutex::new(Vec::new()), + immediate, + }); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await + .expect("client build") + .into_client(); + const GENERATION: u64 = 18; + client + .connection_generation + .store(GENERATION, Ordering::SeqCst); + let registration = client.lifecycle.as_ref().expect("lifecycle registration"); + assert!(registration.begin_scope_if_current(GENERATION, || true)); + + tokio::time::timeout(Duration::from_secs(2), client.dispatch_connected()) + .await + .expect("reentrant reconnect completed"); + let scope = registration + .scope_for(GENERATION) + .expect("cancelled connection scope"); + assert_eq!(scope.state(), ConnectionScopeState::Cancelled); + assert!(!client.is_ready.load(Ordering::Relaxed)); + + registration.close_scope(GENERATION); + registration.shutdown().await; + assert_eq!( + *lifecycle + .events + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["ready-started", "ready-finished", "closed", "shutdown"] + ); + } + } + #[tokio::test] async fn callback_timeout_does_not_hold_connection_cleanup() { let (ready_started_tx, ready_started_rx) = async_channel::bounded(1); diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 50c5f359d..087120339 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -782,6 +782,9 @@ impl Client { )] pub async fn reconnect(self: &Arc) { info!("Reconnecting: dropping transport for auto-reconnect."); + if let Some(lifecycle) = &self.lifecycle { + lifecycle.cancel_active_scope(); + } wacore::telemetry::reconnect(); self.intentional_reconnect.store(true, Ordering::Relaxed); self.auto_reconnect_errors @@ -818,6 +821,9 @@ impl Client { )] pub async fn reconnect_immediately(self: &Arc) { info!("Reconnecting immediately (expected disconnect)."); + if let Some(lifecycle) = &self.lifecycle { + lifecycle.cancel_active_scope(); + } self.expected_disconnect.store(true, Ordering::Relaxed); // Same durable-before-receipts gate as disconnect(). From c86f6d522ea24b47825da6731da6bac9a9e11714 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:04:44 -0300 Subject: [PATCH 24/46] fix(plugins): seal construction and rollback races --- src/client/builder.rs | 161 +++++++++++++++++++++++++++++- src/client/extension_lifecycle.rs | 65 +++++++++++- src/client/lifecycle.rs | 15 +++ src/plugins/mod.rs | 119 ++++++++++++++++++++-- 4 files changed, 347 insertions(+), 13 deletions(-) diff --git a/src/client/builder.rs b/src/client/builder.rs index 69c7528fb..d9af334b3 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -432,6 +432,7 @@ impl ClientBuilder { }, ); let client = assembly.client(); + let mut construction = ClientConstructionGuard::new(Arc::clone(&client)); if !self.custom_enc_handlers.is_empty() { let _ = client.custom_enc_handlers.set(self.custom_enc_handlers); @@ -471,9 +472,31 @@ impl ClientBuilder { let _ = build.client.saver_handle.set(saver_handle); } #[cfg(feature = "plugins")] - if let Some(plugin_host) = &client.plugin_host { - plugin_host.activate(); + if let Some(plugin_host) = &client.plugin_host + && !plugin_host.activate() + { + client.signal_shutdown_sync(); + client.shutdown_lifecycle().await; + return Err(ClientBuilderError::PluginInstall(anyhow::anyhow!( + "client shutdown began before plugin activation" + ))); + } + if let Some(lifecycle) = &client.lifecycle + && !lifecycle.activate() + { + client.signal_shutdown_sync(); + client.shutdown_lifecycle().await; + #[cfg(feature = "plugins")] + if client.plugin_host.is_some() { + return Err(ClientBuilderError::PluginInstall(anyhow::anyhow!( + "client shutdown raced plugin activation" + ))); + } + return Err(ClientBuilderError::LifecycleInstall(anyhow::anyhow!( + "client shutdown raced lifecycle activation" + ))); } + construction.disarm(); Ok(build) } } @@ -522,6 +545,32 @@ pub(super) struct ClientAssembly { sync_task_receiver: async_channel::Receiver, } +struct ClientConstructionGuard { + client: Arc, + armed: bool, +} + +impl ClientConstructionGuard { + fn new(client: Arc) -> Self { + Self { + client, + armed: true, + } + } + + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for ClientConstructionGuard { + fn drop(&mut self) { + if self.armed { + self.client.signal_shutdown_sync(); + } + } +} + #[derive(Default)] pub(super) struct ClientExtensions { pub(super) lifecycle: Option>, @@ -568,6 +617,43 @@ mod tests { installed_client: std::sync::Mutex>>, } + struct RunDuringInstallLifecycle { + client: async_channel::Sender>, + release: async_channel::Receiver<()>, + run_finished: async_channel::Sender<()>, + } + + impl ClientLifecycle for RunDuringInstallLifecycle { + fn install<'a>( + &'a self, + client: std::sync::Weak, + ) -> wacore::runtime::BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + let client = client + .upgrade() + .ok_or_else(|| anyhow::anyhow!("client unavailable during install"))?; + let run_client = client.clone(); + let run_finished = self.run_finished.clone(); + client + .runtime + .spawn(Box::pin(async move { + run_client.run().await; + let _ = run_finished.send(()).await; + })) + .detach(); + self.client + .send(client) + .await + .map_err(|_| anyhow::anyhow!("test client receiver closed"))?; + self.release + .recv() + .await + .map_err(|_| anyhow::anyhow!("test install release closed"))?; + Ok(()) + }) + } + } + impl ClientLifecycle for FailingLifecycle { fn install<'a>( &'a self, @@ -692,6 +778,77 @@ mod tests { build.into_client().signal_shutdown_sync(); } + #[tokio::test] + async fn run_leaked_during_install_waits_for_complete_construction() { + let (client_tx, client_rx) = async_channel::bounded(1); + let (release_tx, release_rx) = async_channel::bounded(1); + let (run_finished_tx, run_finished_rx) = async_channel::bounded(1); + let builder = complete_builder() + .await + .with_lifecycle(RunDuringInstallLifecycle { + client: client_tx, + release: release_rx, + run_finished: run_finished_tx, + }); + let build = tokio::spawn(async move { builder.build().await }); + let leaked_client = client_rx + .recv() + .await + .expect("client leaked during install"); + tokio::task::yield_now().await; + assert!(!leaked_client.is_running.load(Ordering::Acquire)); + + release_tx.send(()).await.expect("release installation"); + let client = build + .await + .expect("builder task") + .expect("successful build") + .into_client(); + tokio::time::timeout(Duration::from_secs(1), async { + while !client.is_running.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("run released after activation"); + client.signal_shutdown_sync(); + tokio::time::timeout(Duration::from_secs(5), run_finished_rx.recv()) + .await + .expect("run stop timeout") + .expect("run stopped"); + } + + #[tokio::test] + async fn shutdown_during_install_rejects_leaked_run_and_the_build() { + let (client_tx, client_rx) = async_channel::bounded(1); + let (release_tx, release_rx) = async_channel::bounded(1); + let (run_finished_tx, run_finished_rx) = async_channel::bounded(1); + let builder = complete_builder() + .await + .with_lifecycle(RunDuringInstallLifecycle { + client: client_tx, + release: release_rx, + run_finished: run_finished_tx, + }); + let build = tokio::spawn(async move { builder.build().await }); + let leaked_client = client_rx + .recv() + .await + .expect("client leaked during install"); + leaked_client.signal_shutdown_sync(); + release_tx.send(()).await.expect("release installation"); + + assert!(matches!( + build.await.expect("builder task"), + Err(ClientBuilderError::LifecycleInstall(_)) + )); + tokio::time::timeout(Duration::from_secs(5), run_finished_rx.recv()) + .await + .expect("rejected run stop timeout") + .expect("rejected run stopped"); + assert!(!leaked_client.is_running.load(Ordering::Acquire)); + } + #[tokio::test] async fn low_level_builder_installs_options_and_owned_services() { let persistence_manager = Arc::new( diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index fa7d06090..240ee3049 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -15,6 +15,9 @@ const SCOPE_OPEN: u8 = 0; const SCOPE_READY: u8 = 1; const SCOPE_CANCELLED: u8 = 2; const SCOPE_CLOSED: u8 = 3; +const CONSTRUCTION_INSTALLING: u8 = 0; +const CONSTRUCTION_ACTIVE: u8 = 1; +const CONSTRUCTION_REJECTED: u8 = 2; const CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); const CALLBACK_QUEUE_CAPACITY: usize = 64; @@ -167,6 +170,8 @@ pub(super) struct LifecycleRegistration { shutdown_notifier: ShutdownNotifier, callback_timeout: Duration, terminal: AtomicBool, + construction_state: AtomicU8, + construction_notifier: ShutdownNotifier, } #[derive(Default)] @@ -285,6 +290,8 @@ impl LifecycleRegistration { shutdown_notifier: ShutdownNotifier::new(), callback_timeout, terminal: AtomicBool::new(false), + construction_state: AtomicU8::new(CONSTRUCTION_INSTALLING), + construction_notifier: ShutdownNotifier::new(), } } @@ -293,10 +300,48 @@ impl LifecycleRegistration { self.handler.install(client) })) .map_err(|_| anyhow::anyhow!("lifecycle install panicked before returning a future"))?; - std::panic::AssertUnwindSafe(install) + let result = std::panic::AssertUnwindSafe(install) .catch_unwind() .await - .map_err(|_| anyhow::anyhow!("lifecycle install future panicked"))? + .map_err(|_| anyhow::anyhow!("lifecycle install future panicked"))?; + if result.is_ok() + && self.construction_state.load(Ordering::Acquire) == CONSTRUCTION_REJECTED + { + return Err(anyhow::anyhow!( + "client shutdown began during lifecycle installation" + )); + } + result + } + + pub(super) fn activate(&self) -> bool { + if self.terminal.load(Ordering::Acquire) { + self.reject_construction(); + return false; + } + match self.construction_state.compare_exchange( + CONSTRUCTION_INSTALLING, + CONSTRUCTION_ACTIVE, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => self.construction_notifier.notify(), + Err(CONSTRUCTION_ACTIVE) => {} + Err(_) => return false, + } + !self.terminal.load(Ordering::Acquire) + } + + pub(super) async fn wait_until_active(&self) -> bool { + let activated = self.construction_notifier.subscribe(); + match self.construction_state.load(Ordering::Acquire) { + CONSTRUCTION_ACTIVE => return !self.terminal.load(Ordering::Acquire), + CONSTRUCTION_REJECTED => return false, + _ => {} + } + wacore::runtime::wait_for_shutdown(&activated).await; + self.construction_state.load(Ordering::Acquire) == CONSTRUCTION_ACTIVE + && !self.terminal.load(Ordering::Acquire) } pub(super) fn begin_scope_if_current( @@ -467,6 +512,7 @@ impl LifecycleRegistration { } pub(super) fn signal_shutdown_sync(&self) { + self.reject_construction(); let first_signal = if ready_publication_active(self) { self.mark_terminal_and_cancel_scopes() } else { @@ -625,6 +671,21 @@ impl LifecycleRegistration { } } + fn reject_construction(&self) { + if self + .construction_state + .compare_exchange( + CONSTRUCTION_INSTALLING, + CONSTRUCTION_REJECTED, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + self.construction_notifier.notify(); + } + } + fn cancel_active_scope_inner(&self) { if let Some(scope) = &self.scopes().active { scope.cancel(); diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 087120339..b375ab517 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -406,10 +406,25 @@ impl Client { // keepalive-loop span. Identity (lid/pn) attribution comes from the // per-operation spans (send/request), which record it themselves. pub async fn run(self: &Arc) { + if let Some(lifecycle) = &self.lifecycle + && !lifecycle.wait_until_active().await + { + warn!("Client `run` rejected before construction completed."); + return; + } + let shutdown = self.shutdown_signal(); + if shutdown.is_fired() { + warn!("Client `run` called after shutdown."); + return; + } if self.is_running.swap(true, Ordering::SeqCst) { warn!("Client `run` method called while already running."); return; } + if shutdown.is_fired() { + self.is_running.store(false, Ordering::SeqCst); + return; + } // Reconnects are counted at iteration start: every pass after the // first is an attempt actually being made. Counting at the branches // below would also count a final pass that never reconnects (a user diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 18cefbb73..6876f7998 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -1031,6 +1031,7 @@ struct PluginInstallRollback { installed: Vec, current: Option, upstream: Option>, + staged_apis: Option>, armed: bool, } @@ -1041,6 +1042,7 @@ impl PluginInstallRollback { installed: Vec::with_capacity(capacity), current: None, upstream: None, + staged_apis: None, armed: true, } } @@ -1062,6 +1064,7 @@ impl PluginInstallRollback { let current = self.current.take(); let installed = std::mem::take(&mut self.installed); let upstream = self.upstream.take(); + let staged_apis = self.staged_apis.take(); self.armed = false; if current.is_none() && installed.is_empty() && upstream.is_none() { return None; @@ -1078,7 +1081,8 @@ impl PluginInstallRollback { let cleanup_runtime = runtime.clone(); runtime .spawn(Box::pin(async move { - shutdown_staged_plugins(cleanup_runtime, current, installed, upstream).await; + shutdown_staged_plugins(cleanup_runtime, current, installed, upstream, staged_apis) + .await; completed.notify(); })) .detach(); @@ -1102,6 +1106,7 @@ impl PluginInstallRollback { fn disarm(&mut self) { self.armed = false; self.upstream = None; + self.staged_apis = None; } } @@ -1120,6 +1125,7 @@ pub(crate) struct PluginHost { runtime: OnceLock>, event_router: Option, callback_timeout: Duration, + terminal: AtomicBool, } impl PluginHost { @@ -1157,6 +1163,7 @@ impl PluginHost { runtime: OnceLock::new(), event_router, callback_timeout, + terminal: AtomicBool::new(false), }) } @@ -1280,16 +1287,20 @@ impl PluginHost { .map_err(|_| anyhow::anyhow!("plugin host was installed more than once"))?; let mut rollback = PluginInstallRollback::new(runtime.clone(), self.ordered.len()); + self.abort_install_if_terminal(&mut rollback).await?; if let Some(upstream) = &self.upstream { if let Err(error) = plugin_callback(|| upstream.install(client.clone())).await { rollback.disarm(); return Err(error); } rollback.upstream = Some(upstream.clone()); + self.abort_install_if_terminal(&mut rollback).await?; } let staging = Arc::new(ApiRegistry::default()); + rollback.staged_apis = Some(Arc::clone(&staging)); for planned in &self.ordered { + self.abort_install_if_terminal(&mut rollback).await?; let resources = PluginResources::new(); let context = self.context( &client, @@ -1315,12 +1326,14 @@ impl PluginHost { } }; staging.insert(planned.plugin.marker_type_id(), api); + self.abort_install_if_terminal(&mut rollback).await?; let Some(installed) = rollback.current.take() else { rollback.rollback().await; anyhow::bail!("plugin installation rollback state was lost"); }; rollback.installed.push(installed); } + self.abort_install_if_terminal(&mut rollback).await?; self.apis .set(staging.snapshot()) @@ -1334,10 +1347,37 @@ impl PluginHost { Ok(()) } - pub(crate) fn activate(&self) { + async fn abort_install_if_terminal( + &self, + rollback: &mut PluginInstallRollback, + ) -> anyhow::Result<()> { + if !self.terminal.load(Ordering::Acquire) { + return Ok(()); + } + rollback.rollback().await; + anyhow::bail!("plugin host shut down during installation") + } + + pub(crate) fn activate(&self) -> bool { + if self.terminal.load(Ordering::Acquire) { + self.close_installed_resources(); + return false; + } for plugin in self.installed.get().into_iter().flatten() { plugin.resources.activate(); } + if self.terminal.load(Ordering::Acquire) { + self.close_installed_resources(); + false + } else { + true + } + } + + fn close_installed_resources(&self) { + for plugin in self.installed.get().into_iter().flatten().rev() { + plugin.resources.close(); + } } async fn run_callback<'a>( @@ -1425,12 +1465,11 @@ impl ClientLifecycle for PluginHost { } fn signal_shutdown(&self) { + self.terminal.store(true, Ordering::Release); if let Some(router) = &self.event_router { router.close(); } - for plugin in self.installed.get().into_iter().flatten().rev() { - plugin.resources.close(); - } + self.close_installed_resources(); if let Some(upstream) = &self.upstream && std::panic::catch_unwind(AssertUnwindSafe(|| upstream.signal_shutdown())).is_err() { @@ -1474,6 +1513,7 @@ async fn shutdown_staged_plugins( current: Option, mut installed: Vec, upstream: Option>, + staged_apis: Option>, ) { if let Some(plugin) = current { if let Err(error) = wait_for_plugin_tasks( @@ -1527,6 +1567,7 @@ async fn shutdown_staged_plugins( { log::warn!("Upstream lifecycle rollback failed: {error:#}"); } + drop(staged_apis); } async fn wait_for_plugin_tasks( @@ -1656,6 +1697,20 @@ mod tests { log: Log, } + struct ShutdownDuringPluginInstall; + + impl ClientLifecycle for ShutdownDuringPluginInstall { + fn install(&self, client: Weak) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + client + .upgrade() + .ok_or_else(|| anyhow::anyhow!("client unavailable during install"))? + .signal_shutdown_sync(); + Ok(()) + }) + } + } + impl ClientPlugin for FoundationPlugin { type Api = String; @@ -1683,6 +1738,24 @@ mod tests { } } + #[tokio::test] + async fn shutdown_during_upstream_install_prevents_plugin_installation() { + let log = Arc::new(Mutex::new(Vec::new())); + let result = complete_builder() + .await + .with_lifecycle(ShutdownDuringPluginInstall) + .with_plugin(FoundationPlugin { log: log.clone() }) + .build() + .await; + + assert!(matches!(result, Err(ClientBuilderError::PluginInstall(_)))); + assert!( + log.lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .is_empty() + ); + } + struct DependentPlugin { log: Log, } @@ -2017,10 +2090,15 @@ mod tests { struct RollbackPlugin { log: Log, task_dropped: Arc, + api_dropped: Arc, + } + + struct RollbackApi { + _drop_flag: DropFlag, } impl ClientPlugin for RollbackPlugin { - type Api = (); + type Api = RollbackApi; fn manifest(&self) -> PluginManifest { PluginManifest::new("rollback", "0.1.0").with_capability(PluginCapability::Tasks) @@ -2029,6 +2107,7 @@ mod tests { fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { let log = self.log.clone(); let task_dropped = self.task_dropped.clone(); + let api_dropped = self.api_dropped.clone(); Box::pin(async move { record(&log, "install:rollback"); let guard = DropFlag(task_dropped); @@ -2039,18 +2118,25 @@ mod tests { let _guard = guard; futures::future::pending::<()>().await; })?; - Ok(Arc::new(())) + Ok(Arc::new(RollbackApi { + _drop_flag: DropFlag(api_dropped), + })) }) } fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { let log = self.log.clone(); let task_dropped = self.task_dropped.clone(); + let api_dropped = self.api_dropped.clone(); Box::pin(async move { anyhow::ensure!( task_dropped.load(Ordering::Acquire), "rollback task still running during shutdown" ); + anyhow::ensure!( + !api_dropped.load(Ordering::Acquire), + "rollback API dropped before shutdown" + ); record(&log, "shutdown:rollback"); Ok(()) }) @@ -2092,6 +2178,7 @@ mod tests { async fn install_failure_rolls_back_resources_and_plugins_in_lifo_order() { let log = Arc::new(Mutex::new(Vec::new())); let task_dropped = Arc::new(AtomicBool::new(false)); + let api_dropped = Arc::new(AtomicBool::new(false)); let result = complete_builder() .await .with_lifecycle(UpstreamLifecycle { log: log.clone() }) @@ -2099,6 +2186,7 @@ mod tests { .with_plugin(RollbackPlugin { log: log.clone(), task_dropped: task_dropped.clone(), + api_dropped: api_dropped.clone(), }) .build() .await; @@ -2122,6 +2210,7 @@ mod tests { }) .await .expect("rollback aborted the install-scoped task"); + assert!(api_dropped.load(Ordering::Acquire)); } struct BlockingFailingPlugin { @@ -2199,6 +2288,7 @@ mod tests { async fn cancelled_explicit_rollback_finishes_detached_and_signals_upstream() { let log = Arc::new(Mutex::new(Vec::new())); let task_dropped = Arc::new(AtomicBool::new(false)); + let api_dropped = Arc::new(AtomicBool::new(false)); let signalled = Arc::new(AtomicBool::new(false)); let shutdown_saw_signal = Arc::new(AtomicBool::new(false)); let (started_tx, started_rx) = async_channel::bounded(1); @@ -2218,6 +2308,7 @@ mod tests { .with_plugin(RollbackPlugin { log: log.clone(), task_dropped: task_dropped.clone(), + api_dropped: api_dropped.clone(), }); let build = tokio::spawn(async move { builder.build().await }); @@ -2237,7 +2328,10 @@ mod tests { .unwrap_or_else(|poisoned| poisoned.into_inner()) .last() .is_some_and(|entry| entry == "shutdown:upstream"); - if complete && task_dropped.load(Ordering::Acquire) { + if complete + && task_dropped.load(Ordering::Acquire) + && api_dropped.load(Ordering::Acquire) + { break; } tokio::task::yield_now().await; @@ -2303,6 +2397,7 @@ mod tests { async fn cancelled_build_closes_resources_and_schedules_lifo_rollback() { let log = Arc::new(Mutex::new(Vec::new())); let task_dropped = Arc::new(AtomicBool::new(false)); + let api_dropped = Arc::new(AtomicBool::new(false)); let (started_tx, started_rx) = async_channel::bounded(1); let (_release_tx, release_rx) = async_channel::bounded(1); let builder = complete_builder() @@ -2315,6 +2410,7 @@ mod tests { .with_plugin(RollbackPlugin { log: log.clone(), task_dropped: task_dropped.clone(), + api_dropped: api_dropped.clone(), }); let build = tokio::spawn(async move { builder.build().await }); @@ -2329,7 +2425,10 @@ mod tests { .unwrap_or_else(|poisoned| poisoned.into_inner()) .last() .is_some_and(|entry| entry == "shutdown:rollback"); - if complete && task_dropped.load(Ordering::Acquire) { + if complete + && task_dropped.load(Ordering::Acquire) + && api_dropped.load(Ordering::Acquire) + { break; } tokio::task::yield_now().await; @@ -2386,6 +2485,7 @@ mod tests { .with_plugin(RollbackPlugin { log: log.clone(), task_dropped: Arc::new(AtomicBool::new(false)), + api_dropped: Arc::new(AtomicBool::new(false)), }) .build() .await; @@ -2933,6 +3033,7 @@ mod tests { .with_plugin(RollbackPlugin { log: Arc::new(Mutex::new(Vec::new())), task_dropped: task_dropped.clone(), + api_dropped: Arc::new(AtomicBool::new(false)), }) .with_plugin(EventSubscriptionPlugin) .build() From 63bb4d95cbfb9f78f52588eab5a130f16461a4df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:16:05 -0300 Subject: [PATCH 25/46] perf(plugins): make lifecycle integration opt in --- Cargo.toml | 5 +++- src/client.rs | 4 ++++ src/client/builder.rs | 30 ++++++++++++++++++++++- src/client/lifecycle.rs | 53 +++++++++++++++++++++++++++++++---------- src/client/node_io.rs | 1 + src/lib.rs | 14 +++++------ 6 files changed, 85 insertions(+), 22 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 0d4ca88d9..bfa2aacc9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -113,9 +113,12 @@ zlib-rs = { version = "0.6.5", default-features = false, features = ["std", "rus [features] debug-snapshots = ["wacore/debug-snapshots"] +# Generation-scoped extension lifecycle. Kept opt-in so ordinary clients do not +# retain lifecycle state or branches when no extension host is present. +client-lifecycle = [] # Build-time native plugin host. Kept opt-in so clients that do not use plugins # retain the pre-host binary footprint. -plugins = ["dep:bon"] +plugins = ["client-lifecycle", "dep:bon"] # Optional observability. Off by default: no `tracing` dep, zero overhead. # Emits tracing spans/events only; the application installs the subscriber # (and any OpenTelemetry bridge). See examples/observability.rs. diff --git a/src/client.rs b/src/client.rs index d488e1ee6..6981683f1 100644 --- a/src/client.rs +++ b/src/client.rs @@ -5,6 +5,7 @@ mod builder; mod context_impl; mod device_registry; pub(crate) mod device_topology; +#[cfg(feature = "client-lifecycle")] mod extension_lifecycle; mod iq_ops; mod lid_pn; @@ -17,7 +18,9 @@ mod sessions; mod voip; use builder::{ClientAssembly, ClientExtensions}; pub use builder::{ClientBuild, ClientBuilder, ClientBuilderError}; +#[cfg(feature = "client-lifecycle")] use extension_lifecycle::LifecycleRegistration; +#[cfg(feature = "client-lifecycle")] pub use extension_lifecycle::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; pub use voip::{CallError, Voip}; @@ -689,6 +692,7 @@ pub struct Client { /// (keepalive, request waiters, read loop, offline flush) observe this. pub(crate) connection_shutdown: std::sync::Mutex, /// Allocated only when an extension host installs lifecycle callbacks. + #[cfg(feature = "client-lifecycle")] lifecycle: Option>, /// Allocated only when at least one build-time plugin is registered. #[cfg(feature = "plugins")] diff --git a/src/client/builder.rs b/src/client/builder.rs index d9af334b3..34f4c426d 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -4,7 +4,9 @@ use std::time::Duration; use thiserror::Error; -use super::{Client, ClientLifecycle, LifecycleRegistration}; +use super::Client; +#[cfg(feature = "client-lifecycle")] +use super::{ClientLifecycle, LifecycleRegistration}; use crate::cache_config::CacheConfig; use crate::http::HttpClient; #[cfg(feature = "plugins")] @@ -67,6 +69,7 @@ pub enum ClientBuilderError { InvalidBackgroundSaverInterval, #[error("the configured backend does not support the inbound durability hook: {0}")] UnsupportedDurabilityBackend(String), + #[cfg(feature = "client-lifecycle")] #[error("client lifecycle installation failed: {0}")] LifecycleInstall(#[source] anyhow::Error), #[cfg(feature = "plugins")] @@ -97,6 +100,7 @@ pub struct ClientBuilder { task_instrument: Option>, alloc_meter: Option>, background_saver_interval: Option, + #[cfg(feature = "client-lifecycle")] lifecycle: Option>, #[cfg(feature = "plugins")] plugins: Vec, @@ -126,6 +130,7 @@ impl ClientBuilder { task_instrument: None, alloc_meter: None, background_saver_interval: None, + #[cfg(feature = "client-lifecycle")] lifecycle: None, #[cfg(feature = "plugins")] plugins: Vec::new(), @@ -278,6 +283,7 @@ impl ClientBuilder { } /// Install the aggregate lifecycle used by extensions of this client. + #[cfg(feature = "client-lifecycle")] pub fn with_lifecycle(mut self, lifecycle: L) -> Self where L: ClientLifecycle + 'static, @@ -287,6 +293,7 @@ impl ClientBuilder { } /// Install an already-shared aggregate lifecycle. + #[cfg(feature = "client-lifecycle")] pub fn with_lifecycle_arc(mut self, lifecycle: Arc) -> Self { self.lifecycle = Some(lifecycle); self @@ -396,6 +403,7 @@ impl ClientBuilder { None => runtime, }; + #[cfg(feature = "client-lifecycle")] let lifecycle_handler = self.lifecycle; #[cfg(feature = "plugins")] let (lifecycle_handler, plugin_host) = { @@ -407,6 +415,7 @@ impl ClientBuilder { }); (lifecycle_handler, plugin_host) }; + #[cfg(feature = "client-lifecycle")] let lifecycle = lifecycle_handler.map(|handler| { #[cfg(feature = "plugins")] if let Some(plugin_host) = &plugin_host { @@ -426,12 +435,14 @@ impl ClientBuilder { self.override_version, self.cache_config, ClientExtensions { + #[cfg(feature = "client-lifecycle")] lifecycle, #[cfg(feature = "plugins")] plugin_host, }, ); let client = assembly.client(); + #[cfg(feature = "client-lifecycle")] let mut construction = ClientConstructionGuard::new(Arc::clone(&client)); if !self.custom_enc_handlers.is_empty() { @@ -452,6 +463,7 @@ impl ClientBuilder { if let Some(meter) = self.alloc_meter { let _ = client.alloc_meter.set(meter); } + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &client.lifecycle && let Err(error) = lifecycle.install(Arc::downgrade(&client)).await { @@ -481,6 +493,7 @@ impl ClientBuilder { "client shutdown began before plugin activation" ))); } + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &client.lifecycle && !lifecycle.activate() { @@ -496,6 +509,7 @@ impl ClientBuilder { "client shutdown raced lifecycle activation" ))); } + #[cfg(feature = "client-lifecycle")] construction.disarm(); Ok(build) } @@ -545,11 +559,13 @@ pub(super) struct ClientAssembly { sync_task_receiver: async_channel::Receiver, } +#[cfg(feature = "client-lifecycle")] struct ClientConstructionGuard { client: Arc, armed: bool, } +#[cfg(feature = "client-lifecycle")] impl ClientConstructionGuard { fn new(client: Arc) -> Self { Self { @@ -563,6 +579,7 @@ impl ClientConstructionGuard { } } +#[cfg(feature = "client-lifecycle")] impl Drop for ClientConstructionGuard { fn drop(&mut self) { if self.armed { @@ -573,6 +590,7 @@ impl Drop for ClientConstructionGuard { #[derive(Default)] pub(super) struct ClientExtensions { + #[cfg(feature = "client-lifecycle")] pub(super) lifecycle: Option>, #[cfg(feature = "plugins")] pub(super) plugin_host: Option>, @@ -612,17 +630,20 @@ mod tests { use crate::transport::mock::MockTransportFactory; use wacore::runtime::AbortHandle; + #[cfg(feature = "client-lifecycle")] struct FailingLifecycle { spawns: Arc, installed_client: std::sync::Mutex>>, } + #[cfg(feature = "client-lifecycle")] struct RunDuringInstallLifecycle { client: async_channel::Sender>, release: async_channel::Receiver<()>, run_finished: async_channel::Sender<()>, } + #[cfg(feature = "client-lifecycle")] impl ClientLifecycle for RunDuringInstallLifecycle { fn install<'a>( &'a self, @@ -654,6 +675,7 @@ mod tests { } } + #[cfg(feature = "client-lifecycle")] impl ClientLifecycle for FailingLifecycle { fn install<'a>( &'a self, @@ -779,6 +801,7 @@ mod tests { } #[tokio::test] + #[cfg(feature = "client-lifecycle")] async fn run_leaked_during_install_waits_for_complete_construction() { let (client_tx, client_rx) = async_channel::bounded(1); let (release_tx, release_rx) = async_channel::bounded(1); @@ -819,6 +842,7 @@ mod tests { } #[tokio::test] + #[cfg(feature = "client-lifecycle")] async fn shutdown_during_install_rejects_leaked_run_and_the_build() { let (client_tx, client_rx) = async_channel::bounded(1); let (release_tx, release_rx) = async_channel::bounded(1); @@ -915,10 +939,12 @@ mod tests { )); } + #[cfg(feature = "client-lifecycle")] struct PanickingInstallLifecycle { when_polled: bool, } + #[cfg(feature = "client-lifecycle")] impl ClientLifecycle for PanickingInstallLifecycle { fn install( &self, @@ -932,6 +958,7 @@ mod tests { } #[tokio::test] + #[cfg(feature = "client-lifecycle")] async fn lifecycle_install_panics_are_typed_build_errors() { for when_polled in [false, true] { let result = complete_builder() @@ -947,6 +974,7 @@ mod tests { } #[tokio::test] + #[cfg(feature = "client-lifecycle")] async fn lifecycle_install_failure_publishes_nothing_and_starts_no_tasks() { let persistence_manager = Arc::new( PersistenceManager::new(crate::test_utils::create_test_backend().await) diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index b375ab517..c33e437d2 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -36,6 +36,7 @@ impl Client { self.expected_disconnect.store(true, Ordering::Relaxed); self.is_running.store(false, Ordering::Relaxed); self.shutdown_notifier.notify(); + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &self.lifecycle { lifecycle.signal_shutdown_sync(); } @@ -82,20 +83,23 @@ impl Client { /// Dispatch the Connected event and notify waiters. pub(crate) async fn dispatch_connected(&self) { - let generation = self.connection_generation.load(Ordering::SeqCst); - if let Some(lifecycle) = &self.lifecycle { - if !lifecycle.ready(generation).await { - debug!("Skipping Connected dispatch for retired generation {generation}"); - return; - } - if self.connection_generation.load(Ordering::SeqCst) != generation { - debug!("Skipping Connected dispatch after generation changed"); + #[cfg(feature = "client-lifecycle")] + { + let generation = self.connection_generation.load(Ordering::SeqCst); + if let Some(lifecycle) = &self.lifecycle { + if !lifecycle.ready(generation).await { + debug!("Skipping Connected dispatch for retired generation {generation}"); + return; + } + if self.connection_generation.load(Ordering::SeqCst) != generation { + debug!("Skipping Connected dispatch after generation changed"); + return; + } + if !lifecycle.publish_ready(generation, || self.publish_connected()) { + debug!("Skipping Connected dispatch after lifecycle cancellation"); + } return; } - if !lifecycle.publish_ready(generation, || self.publish_connected()) { - debug!("Skipping Connected dispatch after lifecycle cancellation"); - } - return; } self.publish_connected(); @@ -110,12 +114,14 @@ impl Client { self.connected_notifier.notify(usize::MAX); } + #[cfg(feature = "client-lifecycle")] pub(super) async fn shutdown_lifecycle(&self) { if let Some(lifecycle) = &self.lifecycle { lifecycle.shutdown().await; } } + #[cfg(feature = "client-lifecycle")] fn request_lifecycle_shutdown(&self) { if let Some(lifecycle) = &self.lifecycle { lifecycle.request_shutdown(); @@ -176,6 +182,7 @@ impl Client { extensions: ClientExtensions, ) -> ClientAssembly { let ClientExtensions { + #[cfg(feature = "client-lifecycle")] lifecycle, #[cfg(feature = "plugins")] plugin_host, @@ -207,6 +214,7 @@ impl Client { ik_handshake_failures: Arc::new(AtomicU32::new(0)), shutdown_notifier: wacore::runtime::ShutdownNotifier::new(), connection_shutdown: std::sync::Mutex::new(wacore::runtime::ShutdownNotifier::new()), + #[cfg(feature = "client-lifecycle")] lifecycle, #[cfg(feature = "plugins")] plugin_host, @@ -406,6 +414,7 @@ impl Client { // keepalive-loop span. Identity (lid/pn) attribution comes from the // per-operation spans (send/request), which record it themselves. pub async fn run(self: &Arc) { + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &self.lifecycle && !lifecycle.wait_until_active().await { @@ -536,6 +545,7 @@ impl Client { ); self.runtime.sleep(delay).await; } + #[cfg(feature = "client-lifecycle")] self.shutdown_lifecycle().await; info!("Client run loop has shut down."); } @@ -721,6 +731,7 @@ impl Client { self.expected_disconnect.store(true, Ordering::Relaxed); self.is_running.store(false, Ordering::Relaxed); self.shutdown_notifier.notify(); + #[cfg(feature = "client-lifecycle")] self.request_lifecycle_shutdown(); // Drain buffered offline receipts into the flush window before @@ -768,6 +779,7 @@ impl Client { // final flush below and then be acked. self.msg_secret_buffer.seal(); self.msg_secret_buffer.flush().await; + #[cfg(feature = "client-lifecycle")] self.shutdown_lifecycle().await; } @@ -797,6 +809,7 @@ impl Client { )] pub async fn reconnect(self: &Arc) { info!("Reconnecting: dropping transport for auto-reconnect."); + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &self.lifecycle { lifecycle.cancel_active_scope(); } @@ -836,6 +849,7 @@ impl Client { )] pub async fn reconnect_immediately(self: &Arc) { info!("Reconnecting immediately (expected disconnect)."); + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &self.lifecycle { lifecycle.cancel_active_scope(); } @@ -863,6 +877,16 @@ impl Client { feature = "tracing", tracing::instrument(name = "wa.conn.cleanup", level = "debug", skip_all) )] + #[cfg(not(feature = "client-lifecycle"))] + pub(crate) async fn cleanup_connection_state(self: &Arc) { + self.cleanup_connection_state_inner().await; + } + + #[cfg_attr( + feature = "tracing", + tracing::instrument(name = "wa.conn.cleanup", level = "debug", skip_all) + )] + #[cfg(feature = "client-lifecycle")] pub(crate) async fn cleanup_connection_state(self: &Arc) { if self.lifecycle.is_none() { self.cleanup_connection_state_inner().await; @@ -900,7 +924,11 @@ impl Client { // process_classified_message — no decrypt can START after the // permit-held cache settle below, so no rowless ratchet advances can // dirty the cache behind teardown's back. + #[cfg(feature = "client-lifecycle")] let closed_generation = self.connection_generation.fetch_add(1, Ordering::SeqCst); + #[cfg(not(feature = "client-lifecycle"))] + self.connection_generation.fetch_add(1, Ordering::SeqCst); + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &self.lifecycle { lifecycle.cancel_scope(closed_generation); } @@ -1060,6 +1088,7 @@ impl Client { if let Some(proc) = self.app_state_processor.lock().await.as_ref() { proc.clear_key_cache().await; } + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &self.lifecycle { lifecycle.close_scope(closed_generation); } diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 82b3fbd4d..6ae49df84 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -701,6 +701,7 @@ impl Client { // Increment connection generation to invalidate any stale post-login tasks // from previous connections (e.g., during 515 reconnect cycles). let current_generation = self.connection_generation.fetch_add(1, Ordering::SeqCst) + 1; + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &self.lifecycle { let opened = lifecycle.begin_scope_if_current(current_generation, || { self.connection_generation.load(Ordering::SeqCst) == current_generation diff --git a/src/lib.rs b/src/lib.rs index 5a7451e11..bf2944645 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -93,10 +93,9 @@ pub use client::{ StatsSnapshot, StorageResourceReport, TransportResourceReport, }; pub use client::{CallError, Voip}; -pub use client::{ - Client, ClientBuild, ClientBuilder, ClientBuilderError, ClientLifecycle, ConnectionScope, - ConnectionScopeState, RawNodeLease, -}; +pub use client::{Client, ClientBuild, ClientBuilder, ClientBuilderError, RawNodeLease}; +#[cfg(feature = "client-lifecycle")] +pub use client::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; pub use types::durability_hook::InboundDurabilityHook; pub use types::retry_admission::RetryAdmission; pub mod download; @@ -192,10 +191,9 @@ pub mod version; /// `use whatsapp_rust::prelude::*;`. pub mod prelude { pub use crate::bot::{Bot, BotBuilder, BotHandle, EventDelivery, MessageContext}; - pub use crate::client::{ - Client, ClientBuilder, ClientBuilderError, ClientError, ClientLifecycle, ConnectionScope, - ConnectionScopeState, RawNodeLease, - }; + pub use crate::client::{Client, ClientBuilder, ClientBuilderError, ClientError, RawNodeLease}; + #[cfg(feature = "client-lifecycle")] + pub use crate::client::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; #[cfg(feature = "plugins")] pub use crate::plugins::{ ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, From 54cbf95c4a75535770339c76e4a61a685594a53a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:01:32 -0300 Subject: [PATCH 26/46] fix(plugins): make terminal rollback unwind safe --- src/client/builder.rs | 8 +- src/client/extension_lifecycle.rs | 50 ++++++- src/client/lifecycle.rs | 17 ++- src/plugins/mod.rs | 232 +++++++++++++++++++++++++++++- 4 files changed, 290 insertions(+), 17 deletions(-) diff --git a/src/client/builder.rs b/src/client/builder.rs index 34f4c426d..a266ac754 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -845,7 +845,7 @@ mod tests { #[cfg(feature = "client-lifecycle")] async fn shutdown_during_install_rejects_leaked_run_and_the_build() { let (client_tx, client_rx) = async_channel::bounded(1); - let (release_tx, release_rx) = async_channel::bounded(1); + let (_release_tx, release_rx) = async_channel::bounded(1); let (run_finished_tx, run_finished_rx) = async_channel::bounded(1); let builder = complete_builder() .await @@ -860,10 +860,12 @@ mod tests { .await .expect("client leaked during install"); leaked_client.signal_shutdown_sync(); - release_tx.send(()).await.expect("release installation"); assert!(matches!( - build.await.expect("builder task"), + tokio::time::timeout(Duration::from_secs(2), build) + .await + .expect("lifecycle install ignored terminal shutdown") + .expect("builder task"), Err(ClientBuilderError::LifecycleInstall(_)) )); tokio::time::timeout(Duration::from_secs(5), run_finished_rx.recv()) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 240ee3049..3f5e0f7a6 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -8,7 +8,7 @@ use std::time::Duration; use super::Client; use futures::FutureExt; use wacore::runtime::{ - BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, timeout as rt_timeout, + BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, timeout as rt_timeout, wait_for_shutdown, }; const SCOPE_OPEN: u8 = 0; @@ -296,14 +296,32 @@ impl LifecycleRegistration { } pub(super) async fn install(&self, client: Weak) -> anyhow::Result<()> { + let rejected = self.construction_notifier.subscribe(); + if self.construction_state.load(Ordering::Acquire) == CONSTRUCTION_REJECTED { + return Err(anyhow::anyhow!( + "client shutdown began during lifecycle installation" + )); + } let install = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { self.handler.install(client) })) .map_err(|_| anyhow::anyhow!("lifecycle install panicked before returning a future"))?; - let result = std::panic::AssertUnwindSafe(install) - .catch_unwind() - .await - .map_err(|_| anyhow::anyhow!("lifecycle install future panicked"))?; + let cancelled = Box::pin(wait_for_shutdown(&rejected)); + let install = Box::pin(std::panic::AssertUnwindSafe(install).catch_unwind()); + let result = match futures::future::select(cancelled, install).await { + futures::future::Either::Left((_, install)) => { + if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(install))).is_err() + { + log::warn!("Lifecycle install future panicked while being cancelled"); + } + return Err(anyhow::anyhow!( + "client shutdown began during lifecycle installation" + )); + } + futures::future::Either::Right((result, _)) => { + result.map_err(|_| anyhow::anyhow!("lifecycle install future panicked"))? + } + }; if result.is_ok() && self.construction_state.load(Ordering::Acquire) == CONSTRUCTION_REJECTED { @@ -1852,16 +1870,27 @@ mod tests { .await .expect("persistence manager"), ); + let lifecycle = Arc::new(RecordingLifecycle::default()); let client = Client::builder() .with_runtime(TokioRuntime) .with_persistence_manager(persistence_manager) .with_transport_factory(MockTransportFactory::new()) .with_http_client(MockHttpClient) - .with_lifecycle(RecordingLifecycle::default()) + .with_lifecycle_arc(lifecycle.clone()) .build() .await .expect("client build") .into_client(); + const GENERATION: u64 = 41; + client + .connection_generation + .store(GENERATION, Ordering::SeqCst); + let registration = client.lifecycle.as_ref().expect("lifecycle registration"); + assert!(registration.begin_scope_if_current(GENERATION, || true)); + client.dispatch_connected().await; + let scope = registration + .scope_for(GENERATION) + .expect("connection scope"); *client.transport.lock().await = Some(Arc::new(PanickingDisconnect)); let cleanup_client = Arc::clone(&client); @@ -1874,6 +1903,15 @@ mod tests { .expect_err("cleanup panic should reach its waiter"); assert!(panic.is_panic()); + assert_eq!(scope.state(), ConnectionScopeState::Closed); + assert!(registration.scope_for(GENERATION).is_none()); + tokio::time::timeout(Duration::from_secs(2), registration.shutdown()) + .await + .expect("shutdown waited for the panicked cleanup scope"); + assert_eq!( + lifecycle.events(), + vec!["install", "ready:41", "closed:41", "shutdown"] + ); client.signal_shutdown_sync(); } diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index c33e437d2..d0bf8fa0f 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -929,6 +929,19 @@ impl Client { #[cfg(not(feature = "client-lifecycle"))] self.connection_generation.fetch_add(1, Ordering::SeqCst); #[cfg(feature = "client-lifecycle")] + let scope_close = self.lifecycle.as_ref().map(|lifecycle| { + let lifecycle = Arc::clone(lifecycle); + scopeguard::guard((lifecycle, closed_generation), |(lifecycle, generation)| { + if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + lifecycle.close_scope(generation); + })) + .is_err() + { + error!("Client lifecycle scope closure panicked"); + } + }) + }); + #[cfg(feature = "client-lifecycle")] if let Some(lifecycle) = &self.lifecycle { lifecycle.cancel_scope(closed_generation); } @@ -1089,9 +1102,7 @@ impl Client { proc.clear_key_cache().await; } #[cfg(feature = "client-lifecycle")] - if let Some(lifecycle) = &self.lifecycle { - lifecycle.close_scope(closed_generation); - } + drop(scope_close); } /// Waits for the noise socket to be established. diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 6876f7998..e30ced95f 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -1081,9 +1081,19 @@ impl PluginInstallRollback { let cleanup_runtime = runtime.clone(); runtime .spawn(Box::pin(async move { - shutdown_staged_plugins(cleanup_runtime, current, installed, upstream, staged_apis) - .await; + let result = AssertUnwindSafe(shutdown_staged_plugins( + cleanup_runtime, + current, + installed, + upstream, + staged_apis, + )) + .catch_unwind() + .await; completed.notify(); + if result.is_err() { + log::warn!("Plugin installation rollback panicked"); + } })) .detach(); Some(completion) @@ -1126,6 +1136,8 @@ pub(crate) struct PluginHost { event_router: Option, callback_timeout: Duration, terminal: AtomicBool, + terminal_notifier: ShutdownNotifier, + installing_resources: Mutex>>, } impl PluginHost { @@ -1164,6 +1176,8 @@ impl PluginHost { event_router, callback_timeout, terminal: AtomicBool::new(false), + terminal_notifier: ShutdownNotifier::new(), + installing_resources: Mutex::new(Vec::new()), }) } @@ -1286,6 +1300,13 @@ impl PluginHost { .set(runtime.clone()) .map_err(|_| anyhow::anyhow!("plugin host was installed more than once"))?; + let installing_resources = &self.installing_resources; + let _installing_resources = scopeguard::guard((), move |_| { + installing_resources + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clear(); + }); let mut rollback = PluginInstallRollback::new(runtime.clone(), self.ordered.len()); self.abort_install_if_terminal(&mut rollback).await?; if let Some(upstream) = &self.upstream { @@ -1313,9 +1334,29 @@ impl PluginHost { rollback.current = Some(InstalledPlugin { plugin: planned.plugin.clone(), manifest: planned.manifest.clone(), - resources, + resources: Arc::clone(&resources), }); - let api = match plugin_install(|| planned.plugin.install(context)).await { + if !self.track_installing_resources(&resources) { + rollback.rollback().await; + anyhow::bail!("plugin host shut down during installation"); + } + let terminal = self.terminal_notifier.subscribe(); + let cancelled = Box::pin(wait_for_shutdown(&terminal)); + let install = Box::pin(plugin_install(|| planned.plugin.install(context))); + let install_result = match futures::future::select(cancelled, install).await { + futures::future::Either::Left((_, install)) => { + if std::panic::catch_unwind(AssertUnwindSafe(|| drop(install))).is_err() { + log::warn!( + "Plugin `{}` install future panicked while being cancelled", + planned.manifest.id + ); + } + rollback.rollback().await; + anyhow::bail!("plugin host shut down during installation"); + } + futures::future::Either::Right((result, _)) => result, + }; + let api = match install_result { Ok(api) => api, Err(error) => { rollback.rollback().await; @@ -1358,6 +1399,33 @@ impl PluginHost { anyhow::bail!("plugin host shut down during installation") } + fn track_installing_resources(&self, resources: &Arc) -> bool { + let mut installing = self + .installing_resources + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.terminal.load(Ordering::Acquire) { + drop(installing); + resources.close(); + return false; + } + installing.push(Arc::downgrade(resources)); + true + } + + fn close_installing_resources(&self) { + let resources = self + .installing_resources + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .filter_map(Weak::upgrade) + .collect::>(); + for resources in resources { + resources.close(); + } + } + pub(crate) fn activate(&self) -> bool { if self.terminal.load(Ordering::Acquire) { self.close_installed_resources(); @@ -1466,6 +1534,8 @@ impl ClientLifecycle for PluginHost { fn signal_shutdown(&self) { self.terminal.store(true, Ordering::Release); + self.close_installing_resources(); + self.terminal_notifier.notify(); if let Some(router) = &self.event_router { router.close(); } @@ -1567,7 +1637,9 @@ async fn shutdown_staged_plugins( { log::warn!("Upstream lifecycle rollback failed: {error:#}"); } - drop(staged_apis); + if std::panic::catch_unwind(AssertUnwindSafe(|| drop(staged_apis))).is_err() { + log::warn!("Plugin API panicked while being dropped during rollback"); + } } async fn wait_for_plugin_tasks( @@ -1699,6 +1771,57 @@ mod tests { struct ShutdownDuringPluginInstall; + struct CaptureInstallClient { + client: async_channel::Sender>, + } + + impl ClientLifecycle for CaptureInstallClient { + fn install(&self, client: Weak) -> BoxFuture<'_, anyhow::Result<()>> { + let sender = self.client.clone(); + Box::pin(async move { + sender.send(client).await?; + Ok(()) + }) + } + } + + struct TerminalBlockingInstallPlugin { + started: async_channel::Sender, + install_dropped: Arc, + shutdown_called: Arc, + } + + impl ClientPlugin for TerminalBlockingInstallPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("terminal-blocking-install", "0.1.0") + .with_capability(PluginCapability::Tasks) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + let started = self.started.clone(); + let install_dropped = self.install_dropped.clone(); + Box::pin(async move { + let _drop = DropFlag(install_dropped); + let shutdown = context + .tasks() + .ok_or_else(|| anyhow::anyhow!("tasks capability missing"))? + .shutdown_signal(); + started.send(shutdown).await?; + futures::future::pending().await + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let shutdown_called = self.shutdown_called.clone(); + Box::pin(async move { + shutdown_called.store(true, Ordering::Release); + Ok(()) + }) + } + } + impl ClientLifecycle for ShutdownDuringPluginInstall { fn install(&self, client: Weak) -> BoxFuture<'_, anyhow::Result<()>> { Box::pin(async move { @@ -1756,6 +1879,43 @@ mod tests { ); } + #[tokio::test] + async fn shutdown_cancels_an_inflight_plugin_install_and_closes_its_resources() { + let (client_tx, client_rx) = async_channel::bounded(1); + let (started_tx, started_rx) = async_channel::bounded(1); + let install_dropped = Arc::new(AtomicBool::new(false)); + let shutdown_called = Arc::new(AtomicBool::new(false)); + let builder = complete_builder() + .await + .with_lifecycle(CaptureInstallClient { client: client_tx }) + .with_plugin(TerminalBlockingInstallPlugin { + started: started_tx, + install_dropped: install_dropped.clone(), + shutdown_called: shutdown_called.clone(), + }); + + let build = tokio::spawn(async move { builder.build().await }); + let client = client_rx + .recv() + .await + .expect("captured install client") + .upgrade() + .expect("client under construction"); + let resource_shutdown = started_rx.recv().await.expect("plugin install started"); + + client.signal_shutdown_sync(); + assert!(resource_shutdown.is_fired()); + let result = tokio::time::timeout(Duration::from_secs(2), build) + .await + .expect("plugin install ignored terminal shutdown") + .expect("build task"); + + assert!(matches!(result, Err(ClientBuilderError::PluginInstall(_)))); + assert!(install_dropped.load(Ordering::Acquire)); + assert!(shutdown_called.load(Ordering::Acquire)); + drop(client); + } + struct DependentPlugin { log: Log, } @@ -2035,6 +2195,68 @@ mod tests { } } + struct PanickingDropApi; + + impl Drop for PanickingDropApi { + fn drop(&mut self) { + panic!("injected API drop panic"); + } + } + + struct PanickingDropPlugin { + shutdown_called: Arc, + } + + impl ClientPlugin for PanickingDropPlugin { + type Api = PanickingDropApi; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("panicking-drop", "0.1.0") + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async { Ok(Arc::new(PanickingDropApi)) }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let shutdown_called = self.shutdown_called.clone(); + Box::pin(async move { + shutdown_called.store(true, Ordering::Release); + Ok(()) + }) + } + } + + #[tokio::test] + async fn panicking_staged_api_drop_cannot_strand_rollback_completion() { + let shutdown_called = Arc::new(AtomicBool::new(false)); + let plugin = Arc::new(PanickingDropPlugin { + shutdown_called: shutdown_called.clone(), + }); + let manifest = plugin.manifest(); + let erased_plugin: Arc = Arc::new(PluginAdapter(plugin)); + let resources = PluginResources::new(); + let registry = Arc::new(ApiRegistry::default()); + let api: ErasedApi = Arc::new(TypedApi(Arc::new(PanickingDropApi))); + registry.insert(TypeId::of::(), api); + + let mut rollback = PluginInstallRollback::new(Arc::new(TokioRuntime), 1); + rollback.installed.push(InstalledPlugin { + plugin: erased_plugin, + manifest, + resources, + }); + rollback.staged_apis = Some(registry); + + tokio::time::timeout(Duration::from_secs(2), rollback.rollback()) + .await + .expect("panicking API drop stranded rollback completion"); + assert!(shutdown_called.load(Ordering::Acquire)); + } + struct ContextRetainingApi { _context: PluginContext, _drop_flag: DropFlag, From 271d8423e21f195c2d58c113e7e206a1d5227cd4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:21:13 -0300 Subject: [PATCH 27/46] fix(plugins): isolate resource teardown panics --- src/plugins/mod.rs | 114 ++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 108 insertions(+), 6 deletions(-) diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index e30ced95f..a8e8b75f3 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -469,7 +469,17 @@ impl PluginResources { .unwrap_or_else(|poisoned| poisoned.into_inner()); std::mem::take(&mut *subscriptions) }; - drop(subscriptions); + for subscription in subscriptions { + if std::panic::catch_unwind(AssertUnwindSafe(|| drop(subscription))).is_err() { + log::warn!("Plugin core-event subscription panicked while being dropped"); + } + } + } +} + +fn close_plugin_resources(plugin_id: &str, resources: &PluginResources) { + if std::panic::catch_unwind(AssertUnwindSafe(|| resources.close())).is_err() { + log::warn!("Plugin `{plugin_id}` resource closure panicked"); } } @@ -1049,10 +1059,10 @@ impl PluginInstallRollback { fn close_resources(&self) { if let Some(current) = &self.current { - current.resources.close(); + close_plugin_resources(¤t.manifest.id, ¤t.resources); } for plugin in self.installed.iter().rev() { - plugin.resources.close(); + close_plugin_resources(&plugin.manifest.id, &plugin.resources); } } @@ -1406,7 +1416,9 @@ impl PluginHost { .unwrap_or_else(|poisoned| poisoned.into_inner()); if self.terminal.load(Ordering::Acquire) { drop(installing); - resources.close(); + if std::panic::catch_unwind(AssertUnwindSafe(|| resources.close())).is_err() { + log::warn!("Installing plugin resource closure panicked"); + } return false; } installing.push(Arc::downgrade(resources)); @@ -1422,7 +1434,9 @@ impl PluginHost { .filter_map(Weak::upgrade) .collect::>(); for resources in resources { - resources.close(); + if std::panic::catch_unwind(AssertUnwindSafe(|| resources.close())).is_err() { + log::warn!("Installing plugin resource closure panicked"); + } } } @@ -1444,7 +1458,7 @@ impl PluginHost { fn close_installed_resources(&self) { for plugin in self.installed.get().into_iter().flatten().rev() { - plugin.resources.close(); + close_plugin_resources(&plugin.manifest.id, &plugin.resources); } } @@ -3097,6 +3111,70 @@ mod tests { struct EventSubscriptionPlugin; + struct PanickingDropEventHandler; + + impl EventHandler for PanickingDropEventHandler { + fn handle_event(&self, _event: Arc) {} + } + + impl Drop for PanickingDropEventHandler { + fn drop(&mut self) { + panic!("injected event handler drop panic"); + } + } + + struct PanickingSubscriptionPlugin; + + impl ClientPlugin for PanickingSubscriptionPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("panicking-subscription", "0.1.0") + .with_capability(PluginCapability::CoreEvents) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async move { + context + .core_events() + .ok_or_else(|| anyhow::anyhow!("core events capability missing"))? + .subscribe( + EventInterest::of(&[EventKind::Connected]), + Arc::new(PanickingDropEventHandler), + )?; + Ok(Arc::new(())) + }) + } + } + + struct ShutdownSignalPlugin; + + impl ClientPlugin for ShutdownSignalPlugin { + type Api = ShutdownSignal; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("shutdown-signal", "0.1.0").with_capability(PluginCapability::Tasks) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async move { + context + .tasks() + .map(PluginTasks::shutdown_signal) + .map(Arc::new) + .ok_or_else(|| anyhow::anyhow!("tasks capability missing")) + }) + } + } + + struct ShutdownSignalLifecycle(Arc); + + impl ClientLifecycle for ShutdownSignalLifecycle { + fn signal_shutdown(&self) { + self.0.store(true, Ordering::Release); + } + } + impl ClientPlugin for EventSubscriptionPlugin { type Api = (); @@ -3190,6 +3268,30 @@ mod tests { assert!(!client.raw_node_forwarding_enabled()); } + #[tokio::test] + async fn panicking_handler_drop_does_not_strand_later_plugins_or_upstream() { + let upstream_signalled = Arc::new(AtomicBool::new(false)); + let client = complete_builder() + .await + .with_lifecycle(ShutdownSignalLifecycle(upstream_signalled.clone())) + .with_plugin(ShutdownSignalPlugin) + .with_plugin(PanickingSubscriptionPlugin) + .build() + .await + .expect("panicking subscription client") + .into_client(); + let plugin_shutdown = client + .plugin::() + .expect("shutdown signal API"); + + let result = std::panic::catch_unwind(AssertUnwindSafe(|| client.signal_shutdown_sync())); + + assert!(result.is_ok()); + assert!(plugin_shutdown.is_fired()); + assert!(upstream_signalled.load(Ordering::Acquire)); + client.disconnect().await; + } + #[tokio::test] async fn resource_close_drops_reentrant_handlers_outside_the_subscription_lock() { let client = complete_builder() From e2f6bc4cec10e9dabb8c9892f5f4db723f34f2be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:54:48 -0300 Subject: [PATCH 28/46] fix(plugins): harden lifecycle cleanup under faults --- src/client/extension_lifecycle.rs | 152 +++++++++++++++---- src/plugins/mod.rs | 234 ++++++++++++++++++++++++------ 2 files changed, 311 insertions(+), 75 deletions(-) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 3f5e0f7a6..8a8ec7c39 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -7,9 +7,7 @@ use std::time::Duration; use super::Client; use futures::FutureExt; -use wacore::runtime::{ - BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, timeout as rt_timeout, wait_for_shutdown, -}; +use wacore::runtime::{BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, wait_for_shutdown}; const SCOPE_OPEN: u8 = 0; const SCOPE_READY: u8 = 1; @@ -302,26 +300,35 @@ impl LifecycleRegistration { "client shutdown began during lifecycle installation" )); } - let install = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let mut install = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { self.handler.install(client) })) .map_err(|_| anyhow::anyhow!("lifecycle install panicked before returning a future"))?; let cancelled = Box::pin(wait_for_shutdown(&rejected)); - let install = Box::pin(std::panic::AssertUnwindSafe(install).catch_unwind()); - let result = match futures::future::select(cancelled, install).await { - futures::future::Either::Left((_, install)) => { - if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(install))).is_err() - { - log::warn!("Lifecycle install future panicked while being cancelled"); + let result = { + let install_poll = std::future::poll_fn(|context| install.as_mut().poll(context)); + let install_poll = Box::pin(std::panic::AssertUnwindSafe(install_poll).catch_unwind()); + match futures::future::select(cancelled, install_poll).await { + futures::future::Either::Left((_, install_poll)) => { + drop(install_poll); + None } - return Err(anyhow::anyhow!( - "client shutdown began during lifecycle installation" - )); - } - futures::future::Either::Right((result, _)) => { - result.map_err(|_| anyhow::anyhow!("lifecycle install future panicked"))? + futures::future::Either::Right((result, _)) => Some(result), } }; + let drop_panicked = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(install))).is_err(); + if drop_panicked { + return Err(anyhow::anyhow!( + "lifecycle install future panicked while being dropped" + )); + } + let Some(result) = result else { + return Err(anyhow::anyhow!( + "client shutdown began during lifecycle installation" + )); + }; + let result = result.map_err(|_| anyhow::anyhow!("lifecycle install future panicked"))?; if result.is_ok() && self.construction_state.load(Ordering::Acquire) == CONSTRUCTION_REJECTED { @@ -646,22 +653,37 @@ impl LifecycleRegistration { log::warn!("Client lifecycle {name} panicked"); return; }; - let callback = std::future::poll_fn(|context| { + let result = { + let callback_poll = std::future::poll_fn(|context| { + let _callback_context = CallbackContextGuard::enter(self); + callback.as_mut().poll(context) + }); + let callback_poll = + Box::pin(std::panic::AssertUnwindSafe(callback_poll).catch_unwind()); + match futures::future::select(callback_poll, self.runtime.sleep(self.callback_timeout)) + .await + { + futures::future::Either::Left((result, _)) => Some(result), + futures::future::Either::Right(((), callback_poll)) => { + drop(callback_poll); + None + } + } + }; + let drop_panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { let _callback_context = CallbackContextGuard::enter(self); - callback.as_mut().poll(context) - }); - let result = std::panic::AssertUnwindSafe(rt_timeout( - &*self.runtime, - self.callback_timeout, - callback, - )) - .catch_unwind() - .await; + drop(callback); + })) + .is_err(); + if drop_panicked { + log::warn!("Client lifecycle {name} panicked while being dropped"); + return; + } match result { - Ok(Ok(Ok(()))) => {} - Ok(Ok(Err(error))) => log::warn!("Client lifecycle {name} failed: {error:#}"), - Ok(Err(_)) => log::warn!("Client lifecycle {name} timed out"), - Err(_) => log::warn!("Client lifecycle {name} panicked"), + Some(Ok(Ok(()))) => {} + Some(Ok(Err(error))) => log::warn!("Client lifecycle {name} failed: {error:#}"), + Some(Err(_)) => log::warn!("Client lifecycle {name} panicked"), + None => log::warn!("Client lifecycle {name} timed out"), } } @@ -1062,6 +1084,31 @@ mod tests { shutdown_calls: AtomicUsize, } + #[derive(Default)] + struct DropPanickingFutureLifecycle { + closed_calls: AtomicUsize, + shutdown_calls: AtomicUsize, + } + + struct DropPanickingPendingFuture; + + impl Future for DropPanickingPendingFuture { + type Output = anyhow::Result<()>; + + fn poll( + self: std::pin::Pin<&mut Self>, + _context: &mut std::task::Context<'_>, + ) -> std::task::Poll { + std::task::Poll::Pending + } + } + + impl Drop for DropPanickingPendingFuture { + fn drop(&mut self) { + panic!("injected lifecycle callback drop panic"); + } + } + impl ClientLifecycle for SynchronousPanicLifecycle { fn on_ready<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { self.ready_calls.fetch_add(1, Ordering::SeqCst); @@ -1079,6 +1126,26 @@ mod tests { } } + impl ClientLifecycle for DropPanickingFutureLifecycle { + fn on_ready<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(DropPanickingPendingFuture) + } + + fn on_closed<'a>(&'a self, _scope: ConnectionScope) -> BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + self.closed_calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + self.shutdown_calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + } + } + struct LogoutOrderHandler { lifecycle: Arc, } @@ -1763,6 +1830,31 @@ mod tests { assert_eq!(lifecycle.shutdown_calls.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn callback_drop_panics_do_not_strand_the_driver() { + let lifecycle = Arc::new(DropPanickingFutureLifecycle::default()); + let registration = Arc::new(LifecycleRegistration::new_with_timeout( + lifecycle.clone(), + Arc::new(TokioRuntime), + Duration::from_millis(10), + )); + const GENERATION: u64 = 24; + assert!(registration.begin_scope_if_current(GENERATION, || true)); + + assert!( + tokio::time::timeout(Duration::from_secs(1), registration.ready(GENERATION)) + .await + .expect("ready callback cancellation completed") + ); + registration.close_scope(GENERATION); + tokio::time::timeout(Duration::from_secs(1), registration.shutdown()) + .await + .expect("callback driver recovered from a drop panic"); + + assert_eq!(lifecycle.closed_calls.load(Ordering::SeqCst), 1); + assert_eq!(lifecycle.shutdown_calls.load(Ordering::SeqCst), 1); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn final_scope_closure_is_published_before_shutdown() { let lifecycle = Arc::new(RecordingLifecycle::default()); diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index a8e8b75f3..29e86b0d4 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -389,31 +389,54 @@ impl PluginResources { } } - fn connection_task_tracker(&self, generation: u64) -> Arc { + fn connection_task_tracker(&self, generation: u64) -> (Arc, bool) { let mut registry = self .connection_tasks .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); if registry.closed { - return TaskTracker::closed(); + return (TaskTracker::closed(), false); + } + match registry.trackers.entry(generation) { + std::collections::hash_map::Entry::Occupied(entry) => (Arc::clone(entry.get()), false), + std::collections::hash_map::Entry::Vacant(entry) => { + let tracker = TaskTracker::new(); + entry.insert(Arc::clone(&tracker)); + (tracker, true) + } } - Arc::clone( - registry - .trackers - .entry(generation) - .or_insert_with(TaskTracker::new), - ) + } + + fn retire_connection_tasks_on_cancel( + self: &Arc, + runtime: &Arc, + generation: u64, + tracker: Arc, + cancellation: ShutdownSignal, + ) { + // Lifecycle queue pressure may discard on_closed, so retirement follows cancellation. + let resources = Arc::downgrade(self); + runtime + .spawn(Box::pin(async move { + wait_for_shutdown(&cancellation).await; + tracker.close(); + wait_for_shutdown(&tracker.completion_signal()).await; + if let Some(resources) = resources.upgrade() { + resources.forget_connection_tasks(generation, &tracker); + } + })) + .detach(); } fn close_connection_tasks(&self, generation: u64) -> Arc { - let tracker = Arc::clone( - self.connection_tasks - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .trackers - .entry(generation) - .or_insert_with(TaskTracker::closed), - ); + let tracker = self + .connection_tasks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .trackers + .get(&generation) + .cloned() + .unwrap_or_else(TaskTracker::closed); tracker.close(); tracker } @@ -1492,7 +1515,19 @@ impl ClientLifecycle for PluginHost { .manifest .capabilities .contains(PluginCapability::Tasks) - .then(|| plugin.resources.connection_task_tracker(scope.generation())); + .then(|| { + let (tracker, created) = + plugin.resources.connection_task_tracker(scope.generation()); + if created && let Some(runtime) = self.runtime.get() { + plugin.resources.retire_connection_tasks_on_cancel( + runtime, + scope.generation(), + Arc::clone(&tracker), + scope.cancellation_signal(), + ); + } + tracker + }); let plugin_scope = self.connection_scope(scope.clone(), &plugin.manifest, task_tracker); if let Err(error) = self @@ -1681,35 +1716,60 @@ async fn bounded_plugin_callback<'a>( timeout: Duration, make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result<()>>, ) -> anyhow::Result<()> { - match runtime_timeout(runtime, timeout, plugin_callback(make_future)).await { - Ok(result) => result, - Err(_) => anyhow::bail!( - "callback timed out after {:.3} seconds", - timeout.as_secs_f64() - ), + let callback = Box::pin(plugin_callback(make_future)); + match futures::future::select(callback, runtime.sleep(timeout)).await { + futures::future::Either::Left((result, _)) => result, + futures::future::Either::Right(((), callback)) => { + let cancellation_panicked = + std::panic::catch_unwind(AssertUnwindSafe(|| drop(callback))).is_err(); + if cancellation_panicked { + anyhow::bail!( + "callback timed out after {:.3} seconds and panicked while being cancelled", + timeout.as_secs_f64() + ); + } + anyhow::bail!( + "callback timed out after {:.3} seconds", + timeout.as_secs_f64() + ) + } } } async fn plugin_callback<'a>( make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result<()>>, ) -> anyhow::Result<()> { - let future = std::panic::catch_unwind(AssertUnwindSafe(make_future)) + let mut future = std::panic::catch_unwind(AssertUnwindSafe(make_future)) .map_err(|_| anyhow::anyhow!("callback panicked before returning a future"))?; - AssertUnwindSafe(future) - .catch_unwind() - .await - .map_err(|_| anyhow::anyhow!("callback future panicked"))? + let result = AssertUnwindSafe(std::future::poll_fn(|context| { + future.as_mut().poll(context) + })) + .catch_unwind() + .await + .map_err(|_| anyhow::anyhow!("callback future panicked")); + let drop_result = std::panic::catch_unwind(AssertUnwindSafe(|| drop(future))); + if drop_result.is_err() { + anyhow::bail!("callback future panicked while being dropped"); + } + result? } async fn plugin_install<'a>( make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result>, ) -> anyhow::Result { - let future = std::panic::catch_unwind(AssertUnwindSafe(make_future)) + let mut future = std::panic::catch_unwind(AssertUnwindSafe(make_future)) .map_err(|_| anyhow::anyhow!("install panicked before returning a future"))?; - AssertUnwindSafe(future) - .catch_unwind() - .await - .map_err(|_| anyhow::anyhow!("install future panicked"))? + let result = AssertUnwindSafe(std::future::poll_fn(|context| { + future.as_mut().poll(context) + })) + .catch_unwind() + .await + .map_err(|_| anyhow::anyhow!("install future panicked")); + let drop_result = std::panic::catch_unwind(AssertUnwindSafe(|| drop(future))); + if drop_result.is_err() { + anyhow::bail!("install future panicked while being dropped"); + } + result? } fn finish_callbacks(stage: &str, failures: Vec) -> anyhow::Result<()> { @@ -2841,24 +2901,33 @@ mod tests { .expect("scoped task plugin"); let client = build.into_client(); wait_for_flag(&install_started).await; + let host = client.plugin_host.as_ref().expect("plugin host").clone(); + let resources = Arc::clone(&host.installed.get().expect("installed plugins")[0].resources); let scope = ConnectionScope::new(88); - client - .plugin_host - .as_ref() - .expect("plugin host") - .on_ready(scope.clone()) + host.on_ready(scope.clone()) .await .expect("plugin ready callback"); wait_for_flag(&connection_started).await; scope.cancel(); - client - .plugin_host - .as_ref() - .expect("plugin host") - .on_closed(scope) - .await - .expect("plugin closed callback"); + wait_for_flag(&connection_dropped).await; + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let retained = resources + .connection_tasks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .trackers + .contains_key(&scope.generation()); + if !retained { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled generation tracker retired without on_closed"); + host.on_closed(scope).await.expect("plugin closed callback"); assert!(connection_dropped.load(Ordering::Acquire)); assert!(closed_after_task.load(Ordering::Acquire)); assert!(!install_dropped.load(Ordering::Acquire)); @@ -3007,6 +3076,49 @@ mod tests { } } + struct DropPanickingReadyPlugin { + log: Log, + } + + struct DropPanickingPendingFuture; + + impl Future for DropPanickingPendingFuture { + type Output = anyhow::Result<()>; + + fn poll( + self: std::pin::Pin<&mut Self>, + _context: &mut std::task::Context<'_>, + ) -> std::task::Poll { + std::task::Poll::Pending + } + } + + impl Drop for DropPanickingPendingFuture { + fn drop(&mut self) { + panic!("injected callback cancellation panic"); + } + } + + impl ClientPlugin for DropPanickingReadyPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("drop-panicking-ready", "0.1.0") + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async { Ok(Arc::new(())) }) + } + + fn on_ready(&self, _scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + record(&self.log, "ready:drop-panicking-ready"); + Box::pin(DropPanickingPendingFuture) + } + } + #[tokio::test] async fn upstream_ready_failure_does_not_suppress_plugins() { let log = Arc::new(Mutex::new(Vec::new())); @@ -3075,6 +3187,38 @@ mod tests { client.disconnect().await; } + #[tokio::test] + async fn panicking_timeout_cancellation_does_not_suppress_following_plugins() { + let log = Arc::new(Mutex::new(Vec::new())); + let plan = PluginPlan::prepare(vec![ + PluginRegistration::new(DropPanickingReadyPlugin { log: log.clone() }), + PluginRegistration::new(ReadyPlugin::<4> { + id: "following-drop-panic", + dependency: Some("drop-panicking-ready"), + log: log.clone(), + stalls: false, + }), + ]) + .expect("valid callback plan") + .expect("non-empty callback plan"); + let host = PluginHost::new_with_callback_timeout(plan, None, Duration::from_millis(10)); + let client = complete_builder() + .await + .with_lifecycle_arc(host.clone()) + .build() + .await + .expect("callback cancellation client") + .into_client(); + + let result = host.on_ready(ConnectionScope::new(93)).await; + assert!(result.is_err()); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["ready:drop-panicking-ready", "ready:following-drop-panic"] + ); + client.disconnect().await; + } + #[tokio::test] async fn composes_existing_lifecycle_outside_plugin_lifo_order() { let log = Arc::new(Mutex::new(Vec::new())); From e72f2729e2949c378de9adbdca3444f69cc27de5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:13:43 -0300 Subject: [PATCH 29/46] fix(lifecycle): preserve all scope closure callbacks --- src/client/extension_lifecycle.rs | 101 ++++++++++++++++++++---------- 1 file changed, 67 insertions(+), 34 deletions(-) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 8a8ec7c39..6be803182 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -17,7 +17,7 @@ const CONSTRUCTION_INSTALLING: u8 = 0; const CONSTRUCTION_ACTIVE: u8 = 1; const CONSTRUCTION_REJECTED: u8 = 2; const CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); -const CALLBACK_QUEUE_CAPACITY: usize = 64; +const CALLBACK_QUEUE_TARGET_CAPACITY: usize = 64; std::thread_local! { static ACTIVE_CALLBACK: Cell<*const LifecycleRegistration> = const { Cell::new(std::ptr::null()) }; @@ -130,10 +130,11 @@ impl fmt::Debug for ConnectionScope { /// Aggregate lifecycle seam installed during [`Client`](super::Client) construction. /// /// Implementations must make `install` transactional. Connection callbacks are -/// serialized and bounded; connection cleanup only schedules `on_closed` so a -/// stalled extension cannot block reconnect. A future plugin host owns -/// per-plugin ordering and isolation behind this client-level seam. `install` -/// receives a weak client reference so retaining it cannot create a cycle. +/// serialized, with bounded ready work; connection cleanup only schedules +/// `on_closed` so a stalled extension cannot block reconnect. Closure callbacks +/// are lossless and may temporarily exceed the target capacity. A future plugin +/// host owns per-plugin ordering and isolation behind this client-level seam. +/// `install` receives a weak client reference so retaining it cannot create a cycle. /// `signal_shutdown` is the non-blocking boundary for resources that must stop /// even when an FFI host cannot await `shutdown`. pub trait ClientLifecycle: wacore::sync_marker::MaybeSendSync { @@ -197,31 +198,49 @@ struct CallbackQueue { } impl CallbackQueue { - fn push_bounded(&mut self, callback: LifecycleCallback) -> Option { - let dropped = (self.pending.len() >= CALLBACK_QUEUE_CAPACITY) - .then(|| self.pending.pop_front()) - .flatten(); - self.overflowed |= dropped.is_some(); - self.pending.push_back(callback); - dropped + fn push_with_pressure_policy( + &mut self, + callback: LifecycleCallback, + ) -> Option { + if self.pending.len() < CALLBACK_QUEUE_TARGET_CAPACITY { + self.pending.push_back(callback); + return None; + } + + self.overflowed = true; + let ready_position = self + .pending + .iter() + .position(|pending| matches!(pending, LifecycleCallback::Ready { .. })); + match (callback, ready_position) { + (callback @ LifecycleCallback::Ready { .. }, None) => Some(callback), + (callback, Some(position)) => { + let dropped = self.pending.remove(position); + self.pending.push_back(callback); + dropped + } + (callback, None) => { + // Every closed scope must reach the extension even when a callback stalls. + self.pending.push_back(callback); + None + } + } } - fn compact_for_shutdown(&mut self, keep_last_closed: bool) -> Vec { + fn compact_for_shutdown(&mut self) -> Vec { if !self.overflowed { return Vec::new(); } - let last_closed = keep_last_closed - .then(|| { - self.pending - .iter() - .rposition(|callback| matches!(callback, LifecycleCallback::Closed(_))) - }) - .flatten() - .and_then(|position| self.pending.remove(position)); - let dropped = self.pending.drain(..).collect::>(); - if let Some(last_closed) = last_closed { - self.pending.push_back(last_closed); + + let mut retained = VecDeque::with_capacity(self.pending.len()); + let mut dropped = Vec::new(); + for callback in self.pending.drain(..) { + match callback { + callback @ LifecycleCallback::Closed(_) => retained.push_back(callback), + callback => dropped.push(callback), + } } + self.pending = retained; dropped } } @@ -495,14 +514,14 @@ impl LifecycleRegistration { let mut queue = self.callback_queue(); let mut dropped = Vec::new(); if queue.shutdown_requested { - dropped.extend(queue.push_bounded(LifecycleCallback::Closed(scope))); - dropped.extend(queue.compact_for_shutdown(no_open_scopes)); + dropped.extend(queue.push_with_pressure_policy(LifecycleCallback::Closed(scope))); + dropped.extend(queue.compact_for_shutdown()); if no_open_scopes && !queue.shutdown_enqueued { queue.shutdown_enqueued = true; queue.pending.push_back(LifecycleCallback::Shutdown); } } else { - dropped.extend(queue.push_bounded(LifecycleCallback::Closed(scope))); + dropped.extend(queue.push_with_pressure_policy(LifecycleCallback::Closed(scope))); } let should_spawn = !queue.pending.is_empty() && !queue.drain_scheduled; queue.drain_scheduled |= should_spawn; @@ -560,7 +579,7 @@ impl LifecycleRegistration { if queue.shutdown_requested || self.terminal.load(Ordering::Acquire) { (false, Some(callback)) } else { - let dropped = queue.push_bounded(callback); + let dropped = queue.push_with_pressure_policy(callback); let should_spawn = !queue.drain_scheduled; queue.drain_scheduled = true; (should_spawn, dropped) @@ -580,7 +599,7 @@ impl LifecycleRegistration { if !queue.shutdown_requested || queue.shutdown_enqueued { return; } - let dropped = queue.compact_for_shutdown(no_open_scopes); + let dropped = queue.compact_for_shutdown(); if no_open_scopes { queue.shutdown_enqueued = true; queue.pending.push_back(LifecycleCallback::Shutdown); @@ -1708,7 +1727,7 @@ mod tests { } #[tokio::test] - async fn callback_queue_is_bounded_and_terminal_shutdown_compacts_backlog() { + async fn callback_queue_preserves_every_scope_closure_before_shutdown() { let (ready_started_tx, ready_started_rx) = async_channel::bounded(1); let (release_ready_tx, release_ready_rx) = async_channel::bounded(1); let lifecycle = Arc::new(QueuePressureLifecycle { @@ -1735,14 +1754,21 @@ mod tests { .await .expect("ready callback started"); - for generation in 2..(CALLBACK_QUEUE_CAPACITY as u64 * 4) { + let closed_callbacks = CALLBACK_QUEUE_TARGET_CAPACITY as u64 * 4 - 2; + for generation in 2..(CALLBACK_QUEUE_TARGET_CAPACITY as u64 * 4) { let scope = ConnectionScope::new(generation); scope.close(); registration.enqueue_callback(LifecycleCallback::Closed(scope)); + + let (done, _done_rx) = async_channel::bounded(1); + registration.enqueue_callback(LifecycleCallback::Ready { + scope: ConnectionScope::new(generation + closed_callbacks), + done, + }); } assert_eq!( registration.callback_queue().pending.len(), - CALLBACK_QUEUE_CAPACITY + usize::try_from(closed_callbacks).expect("closure count fits usize") ); let shutdown_registration = registration.clone(); @@ -1752,7 +1778,11 @@ mod tests { let compacted = { let queue = registration.callback_queue(); if queue.shutdown_enqueued { - assert_eq!(queue.pending.len(), 2); + assert_eq!( + queue.pending.len(), + usize::try_from(closed_callbacks + 1) + .expect("terminal callback count fits usize") + ); true } else { false @@ -1772,7 +1802,10 @@ mod tests { .await .expect("release ready callback"); shutdown.await.expect("bounded terminal shutdown"); - assert_eq!(lifecycle.closed_calls.load(Ordering::SeqCst), 1); + assert_eq!( + lifecycle.closed_calls.load(Ordering::SeqCst), + usize::try_from(closed_callbacks).expect("closure count fits usize") + ); assert_eq!(lifecycle.shutdown_calls.load(Ordering::SeqCst), 1); } From 667689ab83a8ba421306eb9b6542314f6d669911 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:20:15 -0300 Subject: [PATCH 30/46] feat(plugins): add native observability --- agent_docs/observability.md | 28 +- plugins/metrics/src/lib.rs | 10 +- src/client.rs | 49 ++- src/client/accessors.rs | 56 +++ src/lib.rs | 11 +- src/plugins/events.rs | 360 +++++++++++++++++-- src/plugins/mod.rs | 695 ++++++++++++++++++++++++++++++++---- 7 files changed, 1102 insertions(+), 107 deletions(-) diff --git a/agent_docs/observability.md b/agent_docs/observability.md index 987e23a97..bb77c71f5 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -58,10 +58,30 @@ figures come from the `wacore::stats::HeapSize` trait: Semantics: honest estimates for attribution and leak detection, not byte-exact accounting. The e2e `memory_soak.rs` logs the byte totals next to RSS; its growth-bound assertions are on entry counts. -When a new cache is added to `Client`, add it to `memory_report()` (the -`MemoryReport::collections()` list keeps the total and `Display` in sync) and -— if it can dominate memory — implement `HeapSize` for its value type next to -that type's definition. +When a new cache is added to `Client`, add it to `memory_report()` (the common +`MemoryReport::collections()` list or its feature-gated report section) and — +if it can dominate memory — implement `HeapSize` for its value type next to that +type's definition. + +With the opt-in `plugins` feature, the report also includes installed plugins, +active install/connection tasks, retained connection generations, core-event +subscriptions, custom-event endpoints, and unique queued payload bytes. Fanout +shares one envelope, so queued payload memory is counted once even when several +endpoints retain it. + +### Plugin host snapshots (opt-in) + +`Client::plugin_stats()` is computed only when called and returns lifecycle, +health, task, subscription, and custom-event counters keyed by public manifest +ID. `PluginEventRouter::stats()` provides endpoint capacity, current unique +queue retention, and cumulative delivery/backpressure totals; publishers can +read their own totals through `PluginEvents::stats()`. + +Health is sticky for the lifetime of the host: lifecycle errors/panics, +timeouts, task-drain timeouts, isolated core-event panics, resource teardown +panics, publication failures, and queue drops mark only the responsible plugin +as degraded. Concurrent snapshots are intentionally approximate, and carry no +message content, JIDs, or phone numbers. ### 3. `BotBuilder::with_task_instrument` — CPU / custom attribution (opt-in) diff --git a/plugins/metrics/src/lib.rs b/plugins/metrics/src/lib.rs index caa507a9d..7eb85186c 100644 --- a/plugins/metrics/src/lib.rs +++ b/plugins/metrics/src/lib.rs @@ -359,7 +359,7 @@ mod tests { use whatsapp_rust::wacore::store::InMemoryBackend; use whatsapp_rust::{ Client, PluginEventEndpointConfig, PluginEventOverflow, PluginEventSubscribeError, - TokioRuntime, + PluginHealth, TokioRuntime, }; use super::*; @@ -469,6 +469,14 @@ mod tests { .await .expect("bounded endpoint reports pressure"); assert!(api.snapshot().install_ticks >= payload.install_ticks); + let stats = client.plugin_stats().expect("plugin host stats"); + let metrics = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == METRICS_PLUGIN_ID) + .expect("metrics plugin stats"); + assert_eq!(metrics.health, PluginHealth::Degraded); + assert!(metrics.events.expect("metrics event stats").dropped > 0); client.disconnect().await; assert!(api.snapshot().shutdown); diff --git a/src/client.rs b/src/client.rs index 6981683f1..9dcec9ac7 100644 --- a/src/client.rs +++ b/src/client.rs @@ -295,15 +295,31 @@ pub struct MemoryReport { pub signal_sessions: CollectionStats, pub signal_identities: CollectionStats, pub signal_sender_keys: CollectionStats, + #[cfg(feature = "plugins")] + pub plugins: u64, + #[cfg(feature = "plugins")] + pub plugin_install_tasks: u64, + #[cfg(feature = "plugins")] + pub plugin_connection_tasks: u64, + #[cfg(feature = "plugins")] + pub plugin_connection_generations: u64, + #[cfg(feature = "plugins")] + pub plugin_core_event_subscriptions: u64, + #[cfg(feature = "plugins")] + pub plugin_event_endpoints: u64, + #[cfg(feature = "plugins")] + pub plugin_event_endpoint_capacity: u64, + /// Unique custom-event envelopes and payload bytes still retained in endpoint queues. + #[cfg(feature = "plugins")] + pub plugin_event_queue: CollectionStats, // -- Misc -- pub chatstate_handlers: usize, pub custom_enc_handlers: usize, } impl MemoryReport { - /// Every byte-carrying collection with its display name — the single list - /// [`Self::total_estimated_bytes`] and `Display` derive from, so a new - /// collection cannot be summed but not shown (or vice versa). + /// Common byte-carrying collections used by both totals and `Display`. + /// Feature-specific collections stay beside their gated report section. fn collections(&self) -> [(&'static str, &CollectionStats); 11] { [ ("group_cache:", &self.group_cache), @@ -322,7 +338,10 @@ impl MemoryReport { /// Sum of every estimated byte figure in the report. pub fn total_estimated_bytes(&self) -> u64 { - self.collections().iter().map(|(_, c)| c.bytes).sum() + let total: u64 = self.collections().iter().map(|(_, c)| c.bytes).sum(); + #[cfg(feature = "plugins")] + let total = total.saturating_add(self.plugin_event_queue.bytes); + total } } @@ -404,6 +423,28 @@ impl std::fmt::Display for MemoryReport { " peak payload storage: {} B", self.history_sync_payload_bytes_peak )?; + #[cfg(feature = "plugins")] + { + writeln!(f, "--- Plugins ---")?; + writeln!(f, " installed: {}", self.plugins)?; + writeln!(f, " install tasks: {}", self.plugin_install_tasks)?; + writeln!( + f, + " connection tasks: {} (generations: {})", + self.plugin_connection_tasks, self.plugin_connection_generations + )?; + writeln!( + f, + " core subscriptions: {}", + self.plugin_core_event_subscriptions + )?; + writeln!( + f, + " event endpoints: {} (capacity: {})", + self.plugin_event_endpoints, self.plugin_event_endpoint_capacity + )?; + line(f, "event_queue:", &self.plugin_event_queue)?; + } writeln!(f, "--- Misc ---")?; writeln!(f, " chatstate_handlers: {}", self.chatstate_handlers)?; writeln!(f, " custom_enc_handlers: {}", self.custom_enc_handlers)?; diff --git a/src/client/accessors.rs b/src/client/accessors.rs index ac28e17a0..ec168b83f 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -193,6 +193,43 @@ impl Client { history_sync_activity.tasks as u64, history_sync_activity.payload_bytes as u64, ); + #[cfg(feature = "plugins")] + let plugin_stats = self.plugin_stats(); + #[cfg(feature = "plugins")] + let ( + plugins, + plugin_install_tasks, + plugin_connection_tasks, + plugin_connection_generations, + plugin_core_event_subscriptions, + ) = plugin_stats + .as_ref() + .map(|host| { + host.plugins.iter().fold( + ( + u64::try_from(host.plugins.len()).unwrap_or(u64::MAX), + 0u64, + 0u64, + 0u64, + 0u64, + ), + |(plugins, install, connection, generations, subscriptions), plugin| { + ( + plugins, + install.saturating_add(plugin.install_tasks), + connection.saturating_add(plugin.connection_tasks), + generations.saturating_add(plugin.connection_generations), + subscriptions.saturating_add(plugin.core_event_subscriptions), + ) + }, + ) + }) + .unwrap_or_default(); + #[cfg(feature = "plugins")] + let plugin_event_router = plugin_stats + .as_ref() + .and_then(|host| host.event_router) + .unwrap_or_default(); MemoryReport { group_cache, @@ -224,6 +261,25 @@ impl Client { signal_sessions, signal_identities, signal_sender_keys, + #[cfg(feature = "plugins")] + plugins, + #[cfg(feature = "plugins")] + plugin_install_tasks, + #[cfg(feature = "plugins")] + plugin_connection_tasks, + #[cfg(feature = "plugins")] + plugin_connection_generations, + #[cfg(feature = "plugins")] + plugin_core_event_subscriptions, + #[cfg(feature = "plugins")] + plugin_event_endpoints: plugin_event_router.active_endpoints, + #[cfg(feature = "plugins")] + plugin_event_endpoint_capacity: plugin_event_router.endpoint_capacity, + #[cfg(feature = "plugins")] + plugin_event_queue: CollectionStats::new( + plugin_event_router.queued_events, + plugin_event_router.queued_payload_bytes, + ), chatstate_handlers, custom_enc_handlers: self.custom_enc_handlers.get().map_or(0, |m| m.len()), } diff --git a/src/lib.rs b/src/lib.rs index bf2944645..7d2e31528 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -117,11 +117,12 @@ pub use plugins::{ ClientPlugin, PluginCapabilities, PluginCapability, PluginConnectionScope, PluginConnectionTasks, PluginContext, PluginCoreEvents, PluginEventEndpointConfig, PluginEventEndpointStats, PluginEventEnvelope, PluginEventOverflow, PluginEventPayloadEncoding, - PluginEventPublishError, PluginEventPublishReport, PluginEventReceiveError, - PluginEventRouteError, PluginEventRouter, PluginEventSelector, PluginEventSubscribeError, - PluginEventSubscription, PluginEventTopic, PluginEventTryReceiveError, PluginEvents, - PluginFuture, PluginIq, PluginIqError, PluginManifest, PluginMessaging, PluginMessagingError, - PluginPlanError, PluginResourceError, PluginTasks, + PluginEventPublishError, PluginEventPublishReport, PluginEventPublisherStats, + PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, PluginEventRouterStats, + PluginEventSelector, PluginEventSubscribeError, PluginEventSubscription, PluginEventTopic, + PluginEventTryReceiveError, PluginEvents, PluginFuture, PluginHealth, PluginHostStats, + PluginIq, PluginIqError, PluginManifest, PluginMessaging, PluginMessagingError, + PluginPlanError, PluginResourceError, PluginState, PluginStats, PluginTasks, }; pub mod request; pub(crate) mod signal_flush; diff --git a/src/plugins/events.rs b/src/plugins/events.rs index 6f30d5498..b622339fb 100644 --- a/src/plugins/events.rs +++ b/src/plugins/events.rs @@ -215,6 +215,46 @@ pub struct PluginEventEndpointStats { pub capacity: usize, } +/// Cumulative publication and fanout counters for one plugin namespace. +/// +/// `published` counts successful calls, including calls with no subscriber. Fanout fields count +/// endpoint outcomes; `delivered` advances only when a receiver removes an envelope from its queue. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +#[non_exhaustive] +pub struct PluginEventPublisherStats { + pub published: u64, + pub publish_failures: u64, + pub matched: u64, + pub enqueued: u64, + pub delivered: u64, + pub dropped: u64, + pub closed: u64, +} + +/// On-demand aggregate for the custom-event router. +/// +/// Current occupancy may move while a concurrent snapshot is being assembled; cumulative counters +/// remain monotonic but are not an atomic cross-publisher transaction. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +#[non_exhaustive] +pub struct PluginEventRouterStats { + pub registered_publishers: u64, + pub active_routes: u64, + pub active_endpoints: u64, + pub endpoint_capacity: u64, + /// Unique event envelopes retained by at least one endpoint queue. + pub queued_events: u64, + /// Payload bytes retained by those unique queued envelopes. + pub queued_payload_bytes: u64, + pub published: u64, + pub publish_failures: u64, + pub matched: u64, + pub enqueued: u64, + pub delivered: u64, + pub dropped: u64, + pub closed: u64, +} + #[derive(Debug, Error, Clone, Copy, PartialEq, Eq)] #[non_exhaustive] #[error("the plugin event endpoint is closed")] @@ -236,9 +276,98 @@ enum EnqueueOutcome { Closed, } +struct PluginEventPublication { + schema_version: u32, + payload_encoding: PluginEventPayloadEncoding, + payload: Bytes, + connection_generation: u64, +} + +#[derive(Default)] +struct PublisherCounters { + published: AtomicU64, + publish_failures: AtomicU64, + matched: AtomicU64, + enqueued: AtomicU64, + delivered: AtomicU64, + dropped: AtomicU64, + closed: AtomicU64, +} + +impl PublisherCounters { + fn record_publish(&self, result: &Result) { + match result { + Ok(report) => { + self.published.fetch_add(1, Ordering::Relaxed); + self.matched.fetch_add(report.matched, Ordering::Relaxed); + self.enqueued.fetch_add(report.enqueued, Ordering::Relaxed); + self.dropped.fetch_add(report.dropped, Ordering::Relaxed); + self.closed.fetch_add(report.closed, Ordering::Relaxed); + } + Err(_) => { + self.publish_failures.fetch_add(1, Ordering::Relaxed); + } + } + } + + fn snapshot(&self) -> PluginEventPublisherStats { + PluginEventPublisherStats { + published: self.published.load(Ordering::Relaxed), + publish_failures: self.publish_failures.load(Ordering::Relaxed), + matched: self.matched.load(Ordering::Relaxed), + enqueued: self.enqueued.load(Ordering::Relaxed), + delivered: self.delivered.load(Ordering::Relaxed), + dropped: self.dropped.load(Ordering::Relaxed), + closed: self.closed.load(Ordering::Relaxed), + } + } +} + +#[derive(Default)] +struct QueueMemory { + events: AtomicU64, + payload_bytes: AtomicU64, +} + +struct QueuedPluginEvent { + envelope: Arc, + publisher: Arc, + memory: Arc, + payload_bytes: u64, +} + +impl QueuedPluginEvent { + fn new( + envelope: Arc, + publisher: Arc, + memory: Arc, + ) -> Arc { + let payload_bytes = u64::try_from(envelope.payload.len()).unwrap_or(u64::MAX); + memory.events.fetch_add(1, Ordering::Relaxed); + memory + .payload_bytes + .fetch_add(payload_bytes, Ordering::Relaxed); + Arc::new(Self { + envelope, + publisher, + memory, + payload_bytes, + }) + } +} + +impl Drop for QueuedPluginEvent { + fn drop(&mut self) { + self.memory.events.fetch_sub(1, Ordering::Relaxed); + self.memory + .payload_bytes + .fetch_sub(self.payload_bytes, Ordering::Relaxed); + } +} + struct EventEndpoint { id: u64, - sender: Sender>, + sender: Sender>, overflow: PluginEventOverflow, capacity: usize, enqueued: AtomicU64, @@ -247,7 +376,7 @@ struct EventEndpoint { } impl EventEndpoint { - fn enqueue(&self, event: Arc) -> EnqueueOutcome { + fn enqueue(&self, event: Arc) -> EnqueueOutcome { match self.overflow { PluginEventOverflow::DropNewest => match self.sender.try_send(event) { Ok(()) => { @@ -306,7 +435,8 @@ struct RouterState { } struct PluginEventRouterInner { - plugin_ids: HashSet>, + publishers: HashMap, Arc>, + queue_memory: Arc, state: RwLock, next_endpoint_id: AtomicU64, closed: AtomicBool, @@ -373,9 +503,14 @@ pub struct PluginEventRouter { impl PluginEventRouter { pub(super) fn new(plugin_ids: impl IntoIterator) -> Self { + let publishers = plugin_ids + .into_iter() + .map(|plugin_id| (Arc::from(plugin_id), Arc::new(PublisherCounters::default()))) + .collect(); Self { inner: Arc::new(PluginEventRouterInner { - plugin_ids: plugin_ids.into_iter().map(Arc::from).collect(), + publishers, + queue_memory: Arc::new(QueueMemory::default()), state: RwLock::new(RouterState::default()), next_endpoint_id: AtomicU64::new(1), closed: AtomicBool::new(false), @@ -396,6 +531,59 @@ impl PluginEventRouter { .is_some_and(|route| !route.endpoints.is_empty()) } + /// Cumulative counters for one registered publisher. + pub fn publisher_stats(&self, plugin_id: &str) -> Option { + self.inner + .publishers + .get(plugin_id) + .map(|stats| stats.snapshot()) + } + + /// Aggregate counters and current queue occupancy. + pub fn stats(&self) -> PluginEventRouterStats { + let (active_routes, active_endpoints, endpoint_capacity) = { + let state = self + .inner + .state + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let endpoint_capacity = state.endpoints.values().fold(0u64, |total, endpoint| { + total.saturating_add(u64::try_from(endpoint.capacity).unwrap_or(u64::MAX)) + }); + ( + u64::try_from(state.routes.len()).unwrap_or(u64::MAX), + u64::try_from(state.endpoints.len()).unwrap_or(u64::MAX), + endpoint_capacity, + ) + }; + let mut snapshot = PluginEventRouterStats { + registered_publishers: u64::try_from(self.inner.publishers.len()).unwrap_or(u64::MAX), + active_routes, + active_endpoints, + endpoint_capacity, + queued_events: self.inner.queue_memory.events.load(Ordering::Relaxed), + queued_payload_bytes: self + .inner + .queue_memory + .payload_bytes + .load(Ordering::Relaxed), + ..PluginEventRouterStats::default() + }; + for publisher in self.inner.publishers.values() { + let publisher = publisher.snapshot(); + snapshot.published = snapshot.published.saturating_add(publisher.published); + snapshot.publish_failures = snapshot + .publish_failures + .saturating_add(publisher.publish_failures); + snapshot.matched = snapshot.matched.saturating_add(publisher.matched); + snapshot.enqueued = snapshot.enqueued.saturating_add(publisher.enqueued); + snapshot.delivered = snapshot.delivered.saturating_add(publisher.delivered); + snapshot.dropped = snapshot.dropped.saturating_add(publisher.dropped); + snapshot.closed = snapshot.closed.saturating_add(publisher.closed); + } + snapshot + } + pub fn subscribe( &self, selectors: impl IntoIterator, @@ -429,7 +617,7 @@ impl PluginEventRouter { return Err(PluginEventSubscribeError::EmptySelectors); } for selector in &selectors { - if !self.inner.plugin_ids.contains(selector.plugin_id()) { + if !self.inner.publishers.contains_key(selector.plugin_id()) { return Err(PluginEventSubscribeError::UnknownPublisher { plugin_id: selector.plugin_id().to_string(), }); @@ -492,14 +680,24 @@ impl PluginEventRouter { fn publish( &self, + publisher: Arc, plugin_id: &Arc, topic: &PluginEventTopic, - schema_version: u32, - payload_encoding: PluginEventPayloadEncoding, - payload: Bytes, - connection_generation: u64, + publication: PluginEventPublication, ) -> Result { - if schema_version == 0 { + let result = self.publish_inner(Arc::clone(&publisher), plugin_id, topic, publication); + publisher.record_publish(&result); + result + } + + fn publish_inner( + &self, + publisher: Arc, + plugin_id: &Arc, + topic: &PluginEventTopic, + publication: PluginEventPublication, + ) -> Result { + if publication.schema_version == 0 { return Err(PluginEventPublishError::InvalidSchemaVersion); } if self.inner.closed.load(Ordering::Acquire) { @@ -530,17 +728,19 @@ impl PluginEventRouter { .checked_add(1) .ok_or(PluginEventPublishError::SequenceExhausted)?; *sequence = next_sequence; - let event = Arc::new( + let envelope = Arc::new( PluginEventEnvelope::builder() .plugin_id(plugin_id.clone()) .topic(topic.clone()) - .schema_version(schema_version) - .payload_encoding(payload_encoding) - .payload(payload) - .connection_generation(connection_generation) + .schema_version(publication.schema_version) + .payload_encoding(publication.payload_encoding) + .payload(publication.payload) + .connection_generation(publication.connection_generation) .sequence(next_sequence) .build(), ); + let event = + QueuedPluginEvent::new(envelope, publisher, Arc::clone(&self.inner.queue_memory)); let mut report = PluginEventPublishReport { matched: u64::try_from(endpoints.len()).unwrap_or(u64::MAX), @@ -570,7 +770,7 @@ impl PluginEventRouter { pub struct PluginEventSubscription { router: PluginEventRouter, endpoint: Arc, - receiver: Receiver>, + receiver: Receiver>, selectors: Vec, } @@ -594,7 +794,8 @@ impl PluginEventSubscription { .await .map_err(|_| PluginEventReceiveError)?; self.endpoint.delivered.fetch_add(1, Ordering::Relaxed); - Ok(event) + event.publisher.delivered.fetch_add(1, Ordering::Relaxed); + Ok(event.envelope.clone()) } pub fn try_recv(&self) -> Result, PluginEventTryReceiveError> { @@ -603,7 +804,8 @@ impl PluginEventSubscription { TryRecvError::Closed => PluginEventTryReceiveError::Closed, })?; self.endpoint.delivered.fetch_add(1, Ordering::Relaxed); - Ok(event) + event.publisher.delivered.fetch_add(1, Ordering::Relaxed); + Ok(event.envelope.clone()) } } @@ -623,6 +825,7 @@ impl Drop for PluginEventSubscription { pub struct PluginEvents { plugin_id: Arc, router: PluginEventRouter, + stats: Arc, resources: Arc, connection_generation: Arc, } @@ -641,6 +844,10 @@ impl PluginEvents { self.router.has_subscribers(&self.selector(topic)) } + pub fn stats(&self) -> PluginEventPublisherStats { + self.stats.snapshot() + } + pub fn publish( &self, topic: &PluginEventTopic, @@ -648,14 +855,20 @@ impl PluginEvents { payload_encoding: PluginEventPayloadEncoding, payload: impl Into, ) -> Result { - self.resources.ensure_active()?; + if let Err(error) = self.resources.ensure_active() { + self.stats.publish_failures.fetch_add(1, Ordering::Relaxed); + return Err(error.into()); + } self.router.publish( + Arc::clone(&self.stats), &self.plugin_id, topic, - schema_version, - payload_encoding, - payload.into(), - self.connection_generation.load(Ordering::Acquire), + PluginEventPublication { + schema_version, + payload_encoding, + payload: payload.into(), + connection_generation: self.connection_generation.load(Ordering::Acquire), + }, ) } } @@ -665,13 +878,15 @@ pub(super) fn publisher( router: PluginEventRouter, resources: Arc, connection_generation: Arc, -) -> PluginEvents { - PluginEvents { +) -> Option { + let stats = router.inner.publishers.get(plugin_id)?.clone(); + Some(PluginEvents { plugin_id: Arc::from(plugin_id), router, + stats, resources, connection_generation, - } + }) } fn valid_topic(topic: &str) -> bool { @@ -698,18 +913,90 @@ mod tests { topic: &PluginEventTopic, value: u32, ) -> PluginEventPublishReport { + let publisher = router + .inner + .publishers + .get(plugin_id) + .cloned() + .expect("registered publisher"); router .publish( + publisher, &Arc::from(plugin_id), topic, - 1, - PluginEventPayloadEncoding::Binary, - Bytes::copy_from_slice(&value.to_be_bytes()), - 7, + PluginEventPublication { + schema_version: 1, + payload_encoding: PluginEventPayloadEncoding::Binary, + payload: Bytes::copy_from_slice(&value.to_be_bytes()), + connection_generation: 7, + }, ) .expect("event publication") } + #[test] + fn router_stats_count_shared_queue_payload_once_and_keep_cumulative_totals() { + let router = PluginEventRouter::new(["metrics".to_string()]); + let tick = topic("tick"); + let first = router + .subscribe( + [selector("metrics", &tick)], + PluginEventEndpointConfig::new(2, PluginEventOverflow::DropNewest), + ) + .expect("first endpoint"); + let second = router + .subscribe( + [selector("metrics", &tick)], + PluginEventEndpointConfig::new(3, PluginEventOverflow::DropNewest), + ) + .expect("second endpoint"); + + assert_eq!(publish(&router, "metrics", &tick, 1).enqueued, 2); + assert_eq!( + router.stats(), + PluginEventRouterStats { + registered_publishers: 1, + active_routes: 1, + active_endpoints: 2, + endpoint_capacity: 5, + queued_events: 1, + queued_payload_bytes: 4, + published: 1, + publish_failures: 0, + matched: 2, + enqueued: 2, + delivered: 0, + dropped: 0, + closed: 0, + } + ); + + first.try_recv().expect("first delivery"); + assert_eq!(router.stats().queued_events, 1); + second.try_recv().expect("second delivery"); + assert_eq!(router.stats().queued_events, 0); + assert_eq!(router.stats().queued_payload_bytes, 0); + assert_eq!(router.stats().delivered, 2); + + drop(first); + drop(second); + let stats = router.stats(); + assert_eq!(stats.active_routes, 0); + assert_eq!(stats.active_endpoints, 0); + assert_eq!(stats.published, 1); + assert_eq!(stats.delivered, 2); + assert_eq!( + router.publisher_stats("metrics"), + Some(PluginEventPublisherStats { + published: 1, + matched: 2, + enqueued: 2, + delivered: 2, + ..PluginEventPublisherStats::default() + }) + ); + } + #[test] fn routes_only_exact_plugin_and_topic_matches() { let router = PluginEventRouter::new(["metrics".to_string(), "audit".to_string()]); @@ -764,6 +1051,17 @@ mod tests { capacity: 2, } ); + assert_eq!( + router.publisher_stats("metrics"), + Some(PluginEventPublisherStats { + published: 4, + matched: 4, + enqueued: 2, + delivered: 2, + dropped: 2, + ..PluginEventPublisherStats::default() + }) + ); } #[test] @@ -785,6 +1083,8 @@ mod tests { assert_eq!(subscription.try_recv().expect("fourth").sequence, 4); assert_eq!(subscription.stats().enqueued, 4); assert_eq!(subscription.stats().dropped, 2); + assert_eq!(router.stats().dropped, 2); + assert_eq!(router.stats().delivered, 2); } #[test] diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 29e86b0d4..4077d4687 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -5,9 +5,9 @@ mod events; pub use events::{ PluginEventEndpointConfig, PluginEventEndpointStats, PluginEventEnvelope, PluginEventOverflow, PluginEventPayloadEncoding, PluginEventPublishError, PluginEventPublishReport, - PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, PluginEventSelector, - PluginEventSubscribeError, PluginEventSubscription, PluginEventTopic, - PluginEventTryReceiveError, PluginEvents, + PluginEventPublisherStats, PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, + PluginEventRouterStats, PluginEventSelector, PluginEventSubscribeError, + PluginEventSubscription, PluginEventTopic, PluginEventTryReceiveError, PluginEvents, }; use std::any::{Any, TypeId}; @@ -226,6 +226,60 @@ pub enum PluginIqError { Iq(#[from] IqError), } +/// Lifecycle state of one installed plugin. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum PluginState { + Installing, + Active, + ShuttingDown, + /// The bounded shutdown attempt completed; health and task counts show incomplete cleanup. + Stopped, +} + +/// Sticky health derived from cumulative host and event-router failures. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum PluginHealth { + Healthy, + Degraded, +} + +/// On-demand runtime snapshot for one plugin, identified only by its public manifest ID. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub struct PluginStats { + pub plugin_id: String, + pub state: PluginState, + pub health: PluginHealth, + /// Lifecycle hooks that returned successfully. + pub callbacks_completed: u64, + /// Lifecycle hook errors and isolated panics. + pub callback_failures: u64, + pub callback_timeouts: u64, + pub task_drain_timeouts: u64, + /// Core-event handler calls that returned without panicking. + pub core_events_delivered: u64, + /// Panics isolated before they could unwind through the client's event dispatcher. + pub core_event_panics: u64, + pub resource_teardown_panics: u64, + pub install_tasks: u64, + pub connection_tasks: u64, + pub connection_generations: u64, + pub core_event_subscriptions: u64, + pub events: Option, +} + +/// On-demand aggregate for the native plugin host. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub struct PluginHostStats { + pub terminal: bool, + pub health: PluginHealth, + pub plugins: Vec, + pub event_router: Option, +} + struct PluginResources { active: AtomicBool, closed: AtomicBool, @@ -234,6 +288,7 @@ struct PluginResources { install_tasks: Arc, connection_tasks: Mutex, subscriptions: Mutex>, + teardown_panics: AtomicU64, } #[derive(Default)] @@ -301,6 +356,13 @@ impl TaskTracker { fn completion_signal(&self) -> ShutdownSignal { self.idle.subscribe() } + + fn active(&self) -> usize { + self.state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .active + } } struct TaskLease { @@ -339,6 +401,7 @@ impl PluginResources { install_tasks: TaskTracker::new(), connection_tasks: Mutex::new(ConnectionTaskRegistry::default()), subscriptions: Mutex::new(Vec::new()), + teardown_panics: AtomicU64::new(0), }) } @@ -494,6 +557,7 @@ impl PluginResources { }; for subscription in subscriptions { if std::panic::catch_unwind(AssertUnwindSafe(|| drop(subscription))).is_err() { + self.teardown_panics.fetch_add(1, Ordering::Relaxed); log::warn!("Plugin core-event subscription panicked while being dropped"); } } @@ -502,10 +566,177 @@ impl PluginResources { fn close_plugin_resources(plugin_id: &str, resources: &PluginResources) { if std::panic::catch_unwind(AssertUnwindSafe(|| resources.close())).is_err() { + resources.teardown_panics.fetch_add(1, Ordering::Relaxed); log::warn!("Plugin `{plugin_id}` resource closure panicked"); } } +impl PluginResources { + fn stats(&self) -> PluginResourceStats { + let (connection_generations, connection_trackers) = { + let registry = self + .connection_tasks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + ( + registry.trackers.len(), + registry.trackers.values().cloned().collect::>(), + ) + }; + let connection_tasks = connection_trackers.iter().fold(0usize, |total, tracker| { + total.saturating_add(tracker.active()) + }); + PluginResourceStats { + active: self.active.load(Ordering::Acquire), + closed: self.closed.load(Ordering::Acquire), + install_tasks: self.install_tasks.active(), + connection_tasks, + connection_generations, + core_event_subscriptions: self + .subscriptions + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .len(), + teardown_panics: self.teardown_panics.load(Ordering::Relaxed), + } + } +} + +#[derive(Default)] +struct PluginResourceStats { + active: bool, + closed: bool, + install_tasks: usize, + connection_tasks: usize, + connection_generations: usize, + core_event_subscriptions: usize, + teardown_panics: u64, +} + +struct PluginDiagnostics { + resources: Mutex>, + callbacks_completed: AtomicU64, + callback_failures: AtomicU64, + callback_timeouts: AtomicU64, + task_drain_timeouts: AtomicU64, + core_events_delivered: AtomicU64, + core_event_panics: AtomicU64, + shutdown_complete: AtomicBool, +} + +impl PluginDiagnostics { + fn new() -> Arc { + Arc::new(Self { + resources: Mutex::new(Weak::new()), + callbacks_completed: AtomicU64::new(0), + callback_failures: AtomicU64::new(0), + callback_timeouts: AtomicU64::new(0), + task_drain_timeouts: AtomicU64::new(0), + core_events_delivered: AtomicU64::new(0), + core_event_panics: AtomicU64::new(0), + shutdown_complete: AtomicBool::new(false), + }) + } + + fn attach_resources(&self, resources: &Arc) { + *self + .resources + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Arc::downgrade(resources); + } + + fn record_callback(&self, result: &Result<(), PluginCallbackError>) { + match result { + Ok(()) => { + self.callbacks_completed.fetch_add(1, Ordering::Relaxed); + } + Err(PluginCallbackError::Timeout { .. }) => { + self.callback_timeouts.fetch_add(1, Ordering::Relaxed); + } + Err(PluginCallbackError::TimeoutCancellationPanic { .. }) => { + self.callback_timeouts.fetch_add(1, Ordering::Relaxed); + self.callback_failures.fetch_add(1, Ordering::Relaxed); + } + Err(PluginCallbackError::Callback(_)) => { + self.callback_failures.fetch_add(1, Ordering::Relaxed); + } + } + } + + fn record_task_drain(&self, result: &Result<(), PluginTaskDrainError>) { + if matches!(result, Err(PluginTaskDrainError::Timeout { .. })) { + self.task_drain_timeouts.fetch_add(1, Ordering::Relaxed); + } + } + + fn mark_stopped(&self) { + self.shutdown_complete.store(true, Ordering::Release); + } + + fn snapshot( + &self, + plugin_id: &str, + terminal: bool, + events: Option, + ) -> PluginStats { + let resources = self + .resources + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .upgrade() + .map(|resources| resources.stats()) + .unwrap_or_default(); + let callbacks_completed = self.callbacks_completed.load(Ordering::Relaxed); + let callback_failures = self.callback_failures.load(Ordering::Relaxed); + let callback_timeouts = self.callback_timeouts.load(Ordering::Relaxed); + let task_drain_timeouts = self.task_drain_timeouts.load(Ordering::Relaxed); + let core_events_delivered = self.core_events_delivered.load(Ordering::Relaxed); + let core_event_panics = self.core_event_panics.load(Ordering::Relaxed); + let state = if self.shutdown_complete.load(Ordering::Acquire) { + PluginState::Stopped + } else if resources.closed || terminal { + PluginState::ShuttingDown + } else if resources.active { + PluginState::Active + } else { + PluginState::Installing + }; + let event_degraded = events + .as_ref() + .is_some_and(|events| events.publish_failures > 0 || events.dropped > 0); + let health = if callback_failures > 0 + || callback_timeouts > 0 + || task_drain_timeouts > 0 + || core_event_panics > 0 + || resources.teardown_panics > 0 + || event_degraded + { + PluginHealth::Degraded + } else { + PluginHealth::Healthy + }; + PluginStats { + plugin_id: plugin_id.to_string(), + state, + health, + callbacks_completed, + callback_failures, + callback_timeouts, + task_drain_timeouts, + core_events_delivered, + core_event_panics, + resource_teardown_panics: resources.teardown_panics, + install_tasks: u64::try_from(resources.install_tasks).unwrap_or(u64::MAX), + connection_tasks: u64::try_from(resources.connection_tasks).unwrap_or(u64::MAX), + connection_generations: u64::try_from(resources.connection_generations) + .unwrap_or(u64::MAX), + core_event_subscriptions: u64::try_from(resources.core_event_subscriptions) + .unwrap_or(u64::MAX), + events, + } + } +} + /// Install-scoped task capability. Work starts after the complete plugin set is published and /// stops during rollback or shutdown. #[derive(Clone)] @@ -549,6 +780,49 @@ impl PluginTasks { pub struct PluginCoreEvents { client: Weak, resources: Arc, + plugin_id: Arc, + diagnostics: Arc, +} + +struct PluginCoreEventHandler { + plugin_id: Arc, + inner: Option>, + resources: Weak, + diagnostics: Arc, +} + +impl EventHandler for PluginCoreEventHandler { + fn handle_event(&self, event: Arc) { + let Some(inner) = &self.inner else { + return; + }; + if std::panic::catch_unwind(AssertUnwindSafe(|| inner.handle_event(event))).is_err() { + self.diagnostics + .core_event_panics + .fetch_add(1, Ordering::Relaxed); + log::warn!("Plugin `{}` core-event handler panicked", self.plugin_id); + } else { + self.diagnostics + .core_events_delivered + .fetch_add(1, Ordering::Relaxed); + } + } +} + +impl Drop for PluginCoreEventHandler { + fn drop(&mut self) { + if let Some(inner) = self.inner.take() + && std::panic::catch_unwind(AssertUnwindSafe(|| drop(inner))).is_err() + { + if let Some(resources) = self.resources.upgrade() { + resources.teardown_panics.fetch_add(1, Ordering::Relaxed); + } + log::warn!( + "Plugin `{}` core-event handler panicked while being dropped", + self.plugin_id + ); + } + } } impl PluginCoreEvents { @@ -564,6 +838,12 @@ impl PluginCoreEvents { let raw_node_lease = interest .wants(EventKind::RawNode) .then(|| client.acquire_raw_node_forwarding()); + let handler = Arc::new(PluginCoreEventHandler { + plugin_id: Arc::clone(&self.plugin_id), + inner: Some(handler), + resources: Arc::downgrade(&self.resources), + diagnostics: Arc::clone(&self.diagnostics), + }); let subscription = client.subscribe(interest, handler); self.resources .retain_subscription(subscription, raw_node_lease) @@ -1057,6 +1337,7 @@ struct InstalledPlugin { plugin: Arc, manifest: PluginManifest, resources: Arc, + diagnostics: Arc, } struct PluginInstallRollback { @@ -1159,9 +1440,18 @@ impl Drop for PluginInstallRollback { } } +struct PluginContextParts { + resources: Arc, + apis: Arc, + runtime: Arc, + connection_generation: Arc, + diagnostics: Arc, +} + pub(crate) struct PluginHost { ordered: Vec, manifests: Vec, + diagnostics: Vec>, upstream: Option>, installed: OnceLock>, apis: OnceLock>, @@ -1199,9 +1489,13 @@ impl PluginHost { .collect::>(); let event_router = (!event_publishers.is_empty()).then(|| PluginEventRouter::new(event_publishers)); + let diagnostics = (0..manifests.len()) + .map(|_| PluginDiagnostics::new()) + .collect(); Arc::new(Self { ordered: plan.ordered, manifests, + diagnostics, upstream, installed: OnceLock::new(), apis: OnceLock::new(), @@ -1222,6 +1516,36 @@ impl PluginHost { &self.manifests } + pub(crate) fn stats(&self) -> PluginHostStats { + let terminal = self.terminal.load(Ordering::Acquire); + let plugins = self + .manifests + .iter() + .zip(&self.diagnostics) + .map(|(manifest, diagnostics)| { + let events = self + .event_router + .as_ref() + .and_then(|router| router.publisher_stats(&manifest.id)); + diagnostics.snapshot(&manifest.id, terminal, events) + }) + .collect::>(); + let health = if plugins + .iter() + .any(|plugin| plugin.health == PluginHealth::Degraded) + { + PluginHealth::Degraded + } else { + PluginHealth::Healthy + }; + PluginHostStats { + terminal, + health, + plugins, + event_router: self.event_router.as_ref().map(PluginEventRouter::stats), + } + } + pub(crate) fn lifecycle_callback_timeout(&self) -> Duration { let callback_count = self.ordered.len() + usize::from(self.upstream.is_some()); let task_barrier_count = self @@ -1243,11 +1567,15 @@ impl PluginHost { &self, client: &Weak, planned: &PlannedPlugin, - resources: Arc, - apis: Arc, - runtime: Arc, - connection_generation: Arc, + parts: PluginContextParts, ) -> PluginContext { + let PluginContextParts { + resources, + apis, + runtime, + connection_generation, + diagnostics, + } = parts; let manifest = &planned.manifest; let capabilities = manifest.capabilities; PluginContext { @@ -1258,6 +1586,8 @@ impl PluginHost { .then(|| PluginCoreEvents { client: client.clone(), resources: Arc::clone(&resources), + plugin_id: Arc::from(manifest.id.as_str()), + diagnostics: Arc::clone(&diagnostics), }), tasks: capabilities .contains(PluginCapability::Tasks) @@ -1281,7 +1611,7 @@ impl PluginHost { .event_router .as_ref() .filter(|_| capabilities.contains(PluginCapability::PluginEvents)) - .map(|router| { + .and_then(|router| { events::publisher( &manifest.id, router.clone(), @@ -1314,11 +1644,14 @@ impl PluginHost { PluginConnectionScope { scope, tasks } } - async fn wait_for_tasks(&self, completion_signals: Vec) -> anyhow::Result<()> { + async fn wait_for_tasks( + &self, + completion_signals: Vec, + ) -> Result<(), PluginTaskDrainError> { let runtime = self .runtime .get() - .ok_or_else(|| anyhow::anyhow!("plugin runtime is unavailable"))?; + .ok_or(PluginTaskDrainError::RuntimeUnavailable)?; wait_for_plugin_tasks(&**runtime, self.callback_timeout, completion_signals).await } @@ -1353,21 +1686,26 @@ impl PluginHost { let staging = Arc::new(ApiRegistry::default()); rollback.staged_apis = Some(Arc::clone(&staging)); - for planned in &self.ordered { + for (planned, diagnostics) in self.ordered.iter().zip(&self.diagnostics) { self.abort_install_if_terminal(&mut rollback).await?; let resources = PluginResources::new(); + diagnostics.attach_resources(&resources); let context = self.context( &client, planned, - Arc::clone(&resources), - Arc::clone(&staging), - runtime.clone(), - connection_generation.clone(), + PluginContextParts { + resources: Arc::clone(&resources), + apis: Arc::clone(&staging), + runtime: runtime.clone(), + connection_generation: connection_generation.clone(), + diagnostics: Arc::clone(diagnostics), + }, ); rollback.current = Some(InstalledPlugin { plugin: planned.plugin.clone(), manifest: planned.manifest.clone(), resources: Arc::clone(&resources), + diagnostics: Arc::clone(diagnostics), }); if !self.track_installing_resources(&resources) { rollback.rollback().await; @@ -1440,6 +1778,7 @@ impl PluginHost { if self.terminal.load(Ordering::Acquire) { drop(installing); if std::panic::catch_unwind(AssertUnwindSafe(|| resources.close())).is_err() { + resources.teardown_panics.fetch_add(1, Ordering::Relaxed); log::warn!("Installing plugin resource closure panicked"); } return false; @@ -1458,6 +1797,7 @@ impl PluginHost { .collect::>(); for resources in resources { if std::panic::catch_unwind(AssertUnwindSafe(|| resources.close())).is_err() { + resources.teardown_panics.fetch_add(1, Ordering::Relaxed); log::warn!("Installing plugin resource closure panicked"); } } @@ -1488,11 +1828,10 @@ impl PluginHost { async fn run_callback<'a>( &'a self, make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result<()>>, - ) -> anyhow::Result<()> { - let runtime = self - .runtime - .get() - .ok_or_else(|| anyhow::anyhow!("plugin runtime is unavailable"))?; + ) -> Result<(), PluginCallbackError> { + let runtime = self.runtime.get().ok_or_else(|| { + PluginCallbackError::Callback(anyhow::anyhow!("plugin runtime is unavailable")) + })?; bounded_plugin_callback(&**runtime, self.callback_timeout, make_future).await } } @@ -1530,10 +1869,11 @@ impl ClientLifecycle for PluginHost { }); let plugin_scope = self.connection_scope(scope.clone(), &plugin.manifest, task_tracker); - if let Err(error) = self + let result = self .run_callback(|| plugin.plugin.on_ready(plugin_scope)) - .await - { + .await; + plugin.diagnostics.record_callback(&result); + if let Err(error) = result { failures.push(format!("{}: {error:#}", plugin.manifest.id)); } } @@ -1551,10 +1891,11 @@ impl ClientLifecycle for PluginHost { .contains(PluginCapability::Tasks) .then(|| plugin.resources.close_connection_tasks(scope.generation())); if let Some(task_tracker) = &task_tracker { - match self + let result = self .wait_for_tasks(vec![task_tracker.completion_signal()]) - .await - { + .await; + plugin.diagnostics.record_task_drain(&result); + match result { Ok(()) => plugin .resources .forget_connection_tasks(scope.generation(), task_tracker), @@ -1565,10 +1906,11 @@ impl ClientLifecycle for PluginHost { } let plugin_scope = self.connection_scope(scope.clone(), &plugin.manifest, task_tracker); - if let Err(error) = self + let result = self .run_callback(|| plugin.plugin.on_closed(plugin_scope)) - .await - { + .await; + plugin.diagnostics.record_callback(&result); + if let Err(error) = result { failures.push(format!("{}: {error:#}", plugin.manifest.id)); } } @@ -1601,13 +1943,17 @@ impl ClientLifecycle for PluginHost { let mut failures = Vec::new(); self.signal_shutdown(); for plugin in self.installed.get().into_iter().flatten().rev() { - if let Err(error) = self + let task_result = self .wait_for_tasks(plugin.resources.task_completion_signals()) - .await - { + .await; + plugin.diagnostics.record_task_drain(&task_result); + if let Err(error) = task_result { failures.push(format!("{} tasks: {error:#}", plugin.manifest.id)); } - if let Err(error) = self.run_callback(|| plugin.plugin.shutdown()).await { + let callback_result = self.run_callback(|| plugin.plugin.shutdown()).await; + plugin.diagnostics.record_callback(&callback_result); + plugin.diagnostics.mark_stopped(); + if let Err(error) = callback_result { failures.push(format!("{}: {error:#}", plugin.manifest.id)); } } @@ -1627,6 +1973,26 @@ impl Drop for PluginHost { } } +#[derive(Debug, Error)] +enum PluginCallbackError { + #[error("callback timed out after {timeout_seconds:.3} seconds")] + Timeout { timeout_seconds: f64 }, + #[error( + "callback timed out after {timeout_seconds:.3} seconds and panicked while being cancelled" + )] + TimeoutCancellationPanic { timeout_seconds: f64 }, + #[error(transparent)] + Callback(#[from] anyhow::Error), +} + +#[derive(Debug, Error)] +enum PluginTaskDrainError { + #[error("plugin runtime is unavailable")] + RuntimeUnavailable, + #[error("plugin tasks did not stop within {timeout_seconds:.3} seconds")] + Timeout { timeout_seconds: f64 }, +} + async fn shutdown_staged_plugins( runtime: Arc, current: Option, @@ -1635,23 +2001,26 @@ async fn shutdown_staged_plugins( staged_apis: Option>, ) { if let Some(plugin) = current { - if let Err(error) = wait_for_plugin_tasks( + let task_result = wait_for_plugin_tasks( &*runtime, PLUGIN_CALLBACK_TIMEOUT, plugin.resources.task_completion_signals(), ) - .await - { + .await; + plugin.diagnostics.record_task_drain(&task_result); + if let Err(error) = task_result { log::warn!( "Plugin `{}` failed-install task cleanup failed: {error:#}", plugin.manifest.id ); } - if let Err(error) = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { + let callback_result = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { plugin.plugin.shutdown() }) - .await - { + .await; + plugin.diagnostics.record_callback(&callback_result); + plugin.diagnostics.mark_stopped(); + if let Err(error) = callback_result { log::warn!( "Plugin `{}` failed-install rollback failed: {error:#}", plugin.manifest.id @@ -1659,23 +2028,26 @@ async fn shutdown_staged_plugins( } } while let Some(plugin) = installed.pop() { - if let Err(error) = wait_for_plugin_tasks( + let task_result = wait_for_plugin_tasks( &*runtime, PLUGIN_CALLBACK_TIMEOUT, plugin.resources.task_completion_signals(), ) - .await - { + .await; + plugin.diagnostics.record_task_drain(&task_result); + if let Err(error) = task_result { log::warn!( "Plugin `{}` rollback task cleanup failed: {error:#}", plugin.manifest.id ); } - if let Err(error) = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { + let callback_result = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { plugin.plugin.shutdown() }) - .await - { + .await; + plugin.diagnostics.record_callback(&callback_result); + plugin.diagnostics.mark_stopped(); + if let Err(error) = callback_result { log::warn!("Plugin `{}` rollback failed: {error:#}", plugin.manifest.id); } } @@ -1695,7 +2067,7 @@ async fn wait_for_plugin_tasks( runtime: &dyn Runtime, timeout: Duration, completion_signals: Vec, -) -> anyhow::Result<()> { +) -> Result<(), PluginTaskDrainError> { let wait_for_all = async move { for signal in completion_signals { wait_for_shutdown(&signal).await; @@ -1703,11 +2075,8 @@ async fn wait_for_plugin_tasks( }; runtime_timeout(runtime, timeout, wait_for_all) .await - .map_err(|_| { - anyhow::anyhow!( - "plugin tasks did not stop within {:.3} seconds", - timeout.as_secs_f64() - ) + .map_err(|_| PluginTaskDrainError::Timeout { + timeout_seconds: timeout.as_secs_f64(), }) } @@ -1715,23 +2084,21 @@ async fn bounded_plugin_callback<'a>( runtime: &dyn Runtime, timeout: Duration, make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result<()>>, -) -> anyhow::Result<()> { +) -> Result<(), PluginCallbackError> { let callback = Box::pin(plugin_callback(make_future)); match futures::future::select(callback, runtime.sleep(timeout)).await { - futures::future::Either::Left((result, _)) => result, + futures::future::Either::Left((result, _)) => result.map_err(PluginCallbackError::Callback), futures::future::Either::Right(((), callback)) => { let cancellation_panicked = std::panic::catch_unwind(AssertUnwindSafe(|| drop(callback))).is_err(); if cancellation_panicked { - anyhow::bail!( - "callback timed out after {:.3} seconds and panicked while being cancelled", - timeout.as_secs_f64() - ); + return Err(PluginCallbackError::TimeoutCancellationPanic { + timeout_seconds: timeout.as_secs_f64(), + }); } - anyhow::bail!( - "callback timed out after {:.3} seconds", - timeout.as_secs_f64() - ) + Err(PluginCallbackError::Timeout { + timeout_seconds: timeout.as_secs_f64(), + }) } } } @@ -1794,6 +2161,11 @@ impl Client { .unwrap_or_default() } + /// Snapshot lifecycle, task, subscription, and custom-event health for installed plugins. + pub fn plugin_stats(&self) -> Option { + self.plugin_host.as_ref().map(|host| host.stats()) + } + /// Subscribe to custom events emitted by installed plugins. /// /// Returns `None` when no manifest requested custom-event publication. @@ -2322,6 +2694,7 @@ mod tests { plugin: erased_plugin, manifest, resources, + diagnostics: PluginDiagnostics::new(), }); rollback.staged_apis = Some(registry); @@ -2903,12 +3276,22 @@ mod tests { wait_for_flag(&install_started).await; let host = client.plugin_host.as_ref().expect("plugin host").clone(); let resources = Arc::clone(&host.installed.get().expect("installed plugins")[0].resources); + let stats = client.plugin_stats().expect("plugin stats"); + assert_eq!(stats.health, PluginHealth::Healthy); + assert_eq!(stats.plugins[0].state, PluginState::Active); + assert_eq!(stats.plugins[0].install_tasks, 1); + assert_eq!(stats.plugins[0].connection_tasks, 0); let scope = ConnectionScope::new(88); host.on_ready(scope.clone()) .await .expect("plugin ready callback"); wait_for_flag(&connection_started).await; + let stats = client.plugin_stats().expect("ready plugin stats"); + assert_eq!(stats.plugins[0].install_tasks, 1); + assert_eq!(stats.plugins[0].connection_tasks, 1); + assert_eq!(stats.plugins[0].connection_generations, 1); + assert_eq!(stats.plugins[0].callbacks_completed, 1); scope.cancel(); wait_for_flag(&connection_dropped).await; tokio::time::timeout(Duration::from_secs(1), async { @@ -2931,10 +3314,18 @@ mod tests { assert!(connection_dropped.load(Ordering::Acquire)); assert!(closed_after_task.load(Ordering::Acquire)); assert!(!install_dropped.load(Ordering::Acquire)); + let stats = client.plugin_stats().expect("closed plugin stats"); + assert_eq!(stats.plugins[0].connection_tasks, 0); + assert_eq!(stats.plugins[0].connection_generations, 0); + assert_eq!(stats.plugins[0].callbacks_completed, 2); client.disconnect().await; assert!(install_dropped.load(Ordering::Acquire)); assert!(shutdown_after_task.load(Ordering::Acquire)); + let stats = client.plugin_stats().expect("stopped plugin stats"); + assert_eq!(stats.plugins[0].state, PluginState::Stopped); + assert_eq!(stats.plugins[0].install_tasks, 0); + assert_eq!(stats.plugins[0].callbacks_completed, 3); } #[tokio::test] @@ -3184,6 +3575,23 @@ mod tests { *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), vec!["ready:stalling-ready", "ready:following-ready"] ); + let stats = host.stats(); + assert_eq!(stats.health, PluginHealth::Degraded); + let stalling = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == "stalling-ready") + .expect("stalling plugin stats"); + assert_eq!(stalling.health, PluginHealth::Degraded); + assert_eq!(stalling.callback_timeouts, 1); + assert_eq!(stalling.callback_failures, 0); + let following = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == "following-ready") + .expect("following plugin stats"); + assert_eq!(following.health, PluginHealth::Healthy); + assert_eq!(following.callbacks_completed, 1); client.disconnect().await; } @@ -3216,6 +3624,22 @@ mod tests { *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), vec!["ready:drop-panicking-ready", "ready:following-drop-panic"] ); + let stats = host.stats(); + let panicking = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == "drop-panicking-ready") + .expect("drop-panicking plugin stats"); + assert_eq!(panicking.health, PluginHealth::Degraded); + assert_eq!(panicking.callback_timeouts, 1); + assert_eq!(panicking.callback_failures, 1); + let following = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == "following-drop-panic") + .expect("following plugin stats"); + assert_eq!(following.health, PluginHealth::Healthy); + assert_eq!(following.callbacks_completed, 1); client.disconnect().await; } @@ -3267,6 +3691,38 @@ mod tests { } } + struct PanickingCoreEventHandler; + + impl EventHandler for PanickingCoreEventHandler { + fn handle_event(&self, _event: Arc) { + panic!("injected core-event handler panic"); + } + } + + struct PanickingCoreEventPlugin; + + impl ClientPlugin for PanickingCoreEventPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("panicking-core-event", "0.1.0") + .with_capability(PluginCapability::CoreEvents) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async move { + context + .core_events() + .ok_or_else(|| anyhow::anyhow!("core events capability missing"))? + .subscribe( + EventInterest::of(&[EventKind::Connected]), + Arc::new(PanickingCoreEventHandler), + )?; + Ok(Arc::new(())) + }) + } + } + struct PanickingSubscriptionPlugin; impl ClientPlugin for PanickingSubscriptionPlugin { @@ -3412,6 +3868,66 @@ mod tests { assert!(!client.raw_node_forwarding_enabled()); } + #[tokio::test] + async fn panicking_core_event_handler_is_isolated_and_degrades_only_its_plugin() { + let client = complete_builder() + .await + .with_plugin(PanickingCoreEventPlugin) + .with_plugin(ShutdownSignalPlugin) + .build() + .await + .expect("panicking core-event client") + .into_client(); + + let result = std::panic::catch_unwind(AssertUnwindSafe(|| { + client + .core + .event_bus + .dispatch(wacore::types::events::Event::Connected( + wacore::types::events::Connected::builder().build(), + )); + })); + + assert!(result.is_ok()); + let stats = client.plugin_stats().expect("plugin stats"); + assert_eq!(stats.health, PluginHealth::Degraded); + let panicking = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == "panicking-core-event") + .expect("panicking plugin stats"); + assert_eq!(panicking.health, PluginHealth::Degraded); + assert_eq!(panicking.core_event_panics, 1); + assert_eq!(panicking.core_events_delivered, 0); + let unaffected = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == "shutdown-signal") + .expect("unaffected plugin stats"); + assert_eq!(unaffected.health, PluginHealth::Healthy); + client.disconnect().await; + } + + #[test] + fn delayed_plugin_handler_drop_is_isolated_and_counted() { + let resources = PluginResources::new(); + let diagnostics = PluginDiagnostics::new(); + diagnostics.attach_resources(&resources); + let handler = Arc::new(PluginCoreEventHandler { + plugin_id: Arc::from("delayed-drop"), + inner: Some(Arc::new(PanickingDropEventHandler)), + resources: Arc::downgrade(&resources), + diagnostics, + }); + let delayed_snapshot = handler.clone(); + drop(handler); + + let result = std::panic::catch_unwind(AssertUnwindSafe(|| drop(delayed_snapshot))); + + assert!(result.is_ok()); + assert_eq!(resources.stats().teardown_panics, 1); + } + #[tokio::test] async fn panicking_handler_drop_does_not_strand_later_plugins_or_upstream() { let upstream_signalled = Arc::new(AtomicBool::new(false)); @@ -3433,6 +3949,14 @@ mod tests { assert!(result.is_ok()); assert!(plugin_shutdown.is_fired()); assert!(upstream_signalled.load(Ordering::Acquire)); + let stats = client.plugin_stats().expect("plugin stats"); + let panicking = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == "panicking-subscription") + .expect("panicking plugin stats"); + assert_eq!(panicking.health, PluginHealth::Degraded); + assert_eq!(panicking.resource_teardown_panics, 1); client.disconnect().await; } @@ -3617,6 +4141,7 @@ mod tests { .expect("bounded event endpoint"); assert!(publisher.has_subscribers(&tick)); + const TICK_PAYLOAD: &[u8] = br#"{"messages":1}"#; let generation = client.connection_generation.load(Ordering::Acquire); assert_eq!( publisher @@ -3624,7 +4149,7 @@ mod tests { &tick, 2, PluginEventPayloadEncoding::Json, - r#"{"messages":1}"#, + Bytes::from_static(TICK_PAYLOAD), ) .expect("publish tick"), PluginEventPublishReport { @@ -3634,12 +4159,31 @@ mod tests { closed: 0, } ); + let stats = client.plugin_stats().expect("plugin host stats"); + assert_eq!(stats.health, PluginHealth::Healthy); + let publisher_stats = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == "event-publisher") + .expect("publisher stats"); + assert_eq!(publisher_stats.state, PluginState::Active); + assert_eq!(publisher_stats.events.expect("event stats").published, 1); + let memory = client.memory_report().await; + assert_eq!(memory.plugins, 2); + assert_eq!(memory.plugin_event_endpoints, 1); + assert_eq!(memory.plugin_event_endpoint_capacity, 1); + assert_eq!(memory.plugin_event_queue.entries, 1); + assert_eq!( + memory.plugin_event_queue.bytes, + u64::try_from(TICK_PAYLOAD.len()).expect("payload length") + ); + assert!(memory.total_estimated_bytes() >= memory.plugin_event_queue.bytes); let event = subscription.recv().await.expect("routed tick"); assert_eq!(&*event.plugin_id, "event-publisher"); assert_eq!(event.topic, tick); assert_eq!(event.schema_version, 2); assert_eq!(event.payload_encoding, PluginEventPayloadEncoding::Json); - assert_eq!(event.payload, Bytes::from_static(br#"{"messages":1}"#)); + assert_eq!(event.payload, Bytes::from_static(TICK_PAYLOAD)); assert_eq!(event.connection_generation, generation); assert_eq!(event.sequence, 1); @@ -3669,5 +4213,30 @@ mod tests { ), Err(PluginEventSubscribeError::Closed) )); + let stats = client.plugin_stats().expect("terminal plugin stats"); + assert_eq!(stats.health, PluginHealth::Degraded); + let publisher_stats = stats + .plugins + .iter() + .find(|plugin| plugin.plugin_id == "event-publisher") + .expect("terminal publisher stats"); + assert_eq!(publisher_stats.state, PluginState::Stopped); + assert_eq!(publisher_stats.health, PluginHealth::Degraded); + assert_eq!( + publisher_stats.events, + Some(PluginEventPublisherStats { + published: 2, + publish_failures: 1, + matched: 2, + enqueued: 2, + delivered: 2, + dropped: 0, + closed: 0, + }) + ); + let router_stats = stats.event_router.expect("terminal router stats"); + assert_eq!(router_stats.active_endpoints, 0); + assert_eq!(router_stats.queued_events, 0); + assert_eq!(router_stats.delivered, 2); } } From 3244e32ccd9dc33158f9d7de553d3cbe2c1ceec8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:50:49 -0300 Subject: [PATCH 31/46] docs(lifecycle): clarify scope closure lock contract --- src/client/extension_lifecycle.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 6be803182..4334fa525 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -484,7 +484,9 @@ impl LifecycleRegistration { self.close_scope_with(generation, || {}); } - /// Non-noop hooks are test-only, must run off-executor, and must not re-enter lifecycle APIs. + /// `after_remove` runs with `scopes` held and before `callback_queue` is acquired. + /// It must not block an async executor or re-enter lifecycle APIs; blocking test hooks run + /// on a dedicated thread. fn close_scope_with(self: &Arc, generation: u64, after_remove: impl FnOnce()) { let (should_spawn, dropped) = { let mut scopes = self.scopes(); From dc55c1e3ec477e45a06ddb75a4ab81184bc90d45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 18:01:36 -0300 Subject: [PATCH 32/46] docs(plugins): define host and adapter contracts --- AGENTS.md | 1 + agent_docs/plugin_architecture.md | 277 ++++++++++++++++++++++++++++++ 2 files changed, 278 insertions(+) create mode 100644 agent_docs/plugin_architecture.md diff --git a/AGENTS.md b/AGENTS.md index b17d5c72d..13b9cf0f3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -44,6 +44,7 @@ Read these when working on the relevant area: - `agent_docs/debugging.md` — evcxr REPL, binary protocol debugging - `agent_docs/binary_size_ci.md` — size-tracking CI: metrics, budgets, baseline semantics - `agent_docs/observability.md` — per-session stats (I/O, memory report, TaskInstrument/CPU), design rules +- `agent_docs/plugin_architecture.md` — native plugin host, lifecycle/capability invariants, future foreign adapter seam - `agent_docs/signal_durability.md` — Signal counter leases, pre-wire gates, crash recovery, review checklist When adding comments to the code, dont be so verbose, also only explain why, not what diff --git a/agent_docs/plugin_architecture.md b/agent_docs/plugin_architecture.md new file mode 100644 index 000000000..352881407 --- /dev/null +++ b/agent_docs/plugin_architecture.md @@ -0,0 +1,277 @@ +# Native Plugin Architecture + +This document defines the native plugin contract, its lifecycle and ownership +invariants, and the boundary a future foreign-language adapter must preserve. +The implementation is intentionally native-only today: bridge, sidecar, and +wire-protocol work starts only with a concrete consumer. + +## Scope + +The initial plugin model supports: + +- build-time registration and transactional installation; +- type-safe Rust APIs exposed through a plugin marker; +- capability-shaped access to core events, tasks, messaging, IQ, and custom + events; +- install-scoped and connection-generation-scoped work; +- bounded custom-event delivery with explicit backpressure; +- on-demand health, resource, and queue snapshots. + +It does not support dynamic installation, ingress interception, pre-ack +decisions, a foreign wire protocol, process isolation, or sandboxing. Those +features must be justified by a real use case because they add materially +stronger compatibility and durability contracts. + +The host lives in the main `whatsapp-rust` crate, not `wacore`: plugins need +high-level client operations and lifecycle coordination. It is enabled by the +opt-in `plugins` feature, which enables `client-lifecycle`. A default build has +neither plugin/lifecycle fields nor their runtime branches. + +## Construction boundary + +`ClientBuilder` is the canonical low-level construction path. It validates +dependencies at runtime and returns `Result`; `BotBuilder` remains the +typestate-preserving facade and delegates to the same path. + +Construction follows one publication boundary: + +1. validate client dependencies and all plugin manifests; +2. resolve plugin dependencies topologically; +3. assemble an inert `Arc`; +4. install the upstream lifecycle and plugins while staging their APIs; +5. start client services; +6. atomically activate lifecycle/plugin resources and publish the completed + build. + +Plugin tasks requested during installation remain parked until activation. A +client leaked through an installation `Weak` cannot run before the +construction gate opens. + +Installation is transactional. Duplicate IDs or marker types, malformed +versions, missing/duplicate dependencies, and cycles fail before client +assembly. If an install fails, is cancelled, panics, or races terminal +shutdown, staged resources close synchronously and asynchronous shutdown hooks +run in reverse installation order. Staged APIs stay alive through task drains +and shutdown hooks. + +Plugins are install-once for one `Client`; reconnecting does not replace the +plugin instance or its exposed API. + +## Type-safe APIs + +Each plugin chooses an associated API type: + +```rust +use std::sync::Arc; + +use anyhow::Result; +use whatsapp_rust::{ClientPlugin, PluginContext, PluginFuture, PluginManifest}; + +struct SearchPlugin; + +struct SearchApi { + // Clone capability handles or plugin-owned state here. +} + +impl SearchApi { + async fn search(&self, query: &str) -> Result> { + // Plugin-specific behavior. + Ok(vec![query.to_owned()]) + } +} + +impl ClientPlugin for SearchPlugin { + type Api = SearchApi; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("example.search", "0.1.0") + } + + fn install(&self, _context: PluginContext) -> PluginFuture<'_, Result>> { + Box::pin(async { Ok(Arc::new(SearchApi {})) }) + } +} +``` + +Register and consume it as follows: + +```rust +let client = Client::builder() + // platform dependencies... + .with_plugin(SearchPlugin) + .build() + .await? + .into_client(); + +let search: Arc = client + .plugin::() + .expect("search plugin is installed"); +let matches = search.search("hello").await?; +``` + +The registry is keyed by `TypeId` of the plugin marker, not the API type. Two +plugins may therefore expose the same API type without colliding. +`Client::plugin::

()` returns `Option>` because the set of plugins +is selected at runtime by the builder. Encoding that set in `Client` generics +would make the client type viral and substantially increase monomorphization. + +During installation, `PluginContext::plugin::

()` exposes only directly +declared dependencies. The context keeps a weak dependency view so an API that +retains its context cannot create a registry ownership cycle. APIs should keep +plugin-owned state and cloned capability handles; they do not receive the raw +backend or Signal stores. + +The workspace crate `plugins/metrics` is the public-API conformance example. It +must remain buildable without private access to the main crate. + +## Capabilities and trust + +A native plugin is trusted in-process Rust code. Its manifest requests +capabilities, and `PluginContext` exposes only the corresponding small handles: + +| Capability identifier | Native handle | Boundary | +| --- | --- | --- | +| `events.core.observe` | `PluginCoreEvents` | Selective observation of sealed core events | +| `tasks.spawn` | `PluginTasks` / `PluginConnectionTasks` | Runtime-agnostic, cancellation-tracked work | +| `messaging.send` | `PluginMessaging` | High-level message sends | +| `iq.execute` | `PluginIq` | Typed `IqSpec` execution | +| `events.plugin.publish` | `PluginEvents` | Publication only in the plugin's own namespace | + +This is API shaping, not a security boundary: native code can use any other +crate dependency available to its process. Runtime grant enforcement belongs +at a future FFI/sidecar boundary, where every foreign command must be checked. +Do not add a per-call native capability checker that suggests sandboxing it +cannot provide. + +Capability handles keep `Weak` internally and reject calls before +activation or after shutdown. This avoids `Client -> plugin API -> Client` +cycles and gives terminal resource invalidation a synchronous boundary. + +## Lifecycle and task ownership + +The host maps the client's existing `connection_generation` to +`PluginConnectionScope`: + +```text +install once + | + +-- install-scoped tasks ---------------------------> terminal shutdown + | + +-- generation N: ready -> cancel -> closed + +-- generation N+1: ready -> cancel -> closed +``` + +- `install` runs once while the client is inert. +- `on_ready` runs in dependency order after authentication for that generation. +- scope cancellation is synchronous when a reconnect or terminal teardown + starts. +- generation tasks drain before `on_closed`; `on_closed` runs in reverse + dependency order after authoritative cleanup. +- install tasks drain before terminal `shutdown`; shutdown hooks run in reverse + dependency order. +- the separately configured upstream lifecycle wraps the plugin order: it is + readied first and closed/shut down last. + +`PluginTasks` survives reconnects and ends only on rollback or terminal +shutdown. `PluginConnectionTasks` is tied to one generation and must not leak +work into a later connection. + +Callbacks are serialized, bounded by timeouts, and isolated from panics, +including panics while constructing, polling, cancelling, or destroying their +futures. One faulty plugin must not suppress later callbacks. Stale `Ready` +work is bounded under reconnect pressure; every accepted `Closed` callback is +lossless and precedes terminal `Shutdown`, so the queue may temporarily exceed +its target to preserve cleanup. + +`signal_shutdown_sync()` closes tasks, subscriptions, event routes, and +capability handles promptly. `disconnect().await` remains required for async +task barriers, hooks, durability flushing, and transport teardown. `Drop` can +only provide the synchronous signal. + +## Event boundaries + +Core events remain the sealed `wacore::types::events::Event` contract. +Subscriptions use explicit `EventInterest`; interest changes go through the +retained `Subscription`, and the aggregate 128-bit mask provides the producer +fast path. Plugin core handlers run inline, must not block, and should hand work +to a task capability. `PluginCoreEvents` retains its subscriptions for the +plugin lifetime; early removal is not part of the initial plugin API. Requesting +`RawNode` retains a forwarding lease for the same lifetime. + +Custom events never enter the core enum or consume an `EventInterest` bit. +`PluginEventRouter` routes exact `(plugin_id, topic)` selectors and gives each +consumer an independently bounded queue with `DropNewest` or `DropOldest`. +Fanout shares one immutable envelope/payload across matching queues. + +The native envelope carries: + +- plugin ID and validated topic; +- schema version and payload encoding; +- opaque payload bytes; +- connection generation at publication; +- route-local monotonic sequence. + +Dropped events consume sequence numbers so consumers can detect gaps. A route +clock resets only after its final subscriber leaves. Publishers can check the +exact route before serializing, and the router filters before constructing an +envelope; a future adapter must preserve that filter before waking or crossing +FFI. + +`PluginEventSubscription` is an RAII endpoint. Dropping it atomically removes +all its selectors. Router shutdown rejects new work but lets receivers drain +already queued envelopes. + +## Diagnostics + +`Client::plugin_stats()` reports per-plugin lifecycle state, sticky health, +callback/task failures, active task scopes, subscriptions, and publisher +counters. `PluginEventRouter::stats()` and `PluginEvents::stats()` expose queue +and backpressure totals. `Client::memory_report()` includes plugin resources +and counts a shared queued payload once across fanout. + +Snapshots are on-demand, approximate under concurrency, and contain no JIDs, +phone numbers, or message bodies. See `observability.md` for accounting rules. + +## Future foreign-language adapter seam + +A future bridge should be a Rust adapter at the host boundary, not a second +client lifecycle. It can map a foreign endpoint onto the existing semantics: + +- build-time registration and stable install-scoped handles across reconnects; +- explicit capability grants checked for every foreign command; +- exact core/custom-event subscriptions before serialization or FFI wake-up; +- one bounded queue and overflow policy per endpoint; +- lifecycle events keyed by connection generation; +- event sequences, drop counters, timeouts, and typed failures; +- synchronous terminal invalidation followed by bounded asynchronous cleanup. + +The native structs are not the wire schema. When a bridge consumer is in scope, +define a separate versioned protocol from a working native + foreign vertical +slice. It must specify payload/command size limits, batching, unknown-field +behavior, removed-field reservations, error codes, lifecycle deadlines, schema +generation, and drift tests. Capability identifiers may be reused, but native +traits, `Client`, runtime objects, stores, and raw backend access must not cross +that protocol. + +Sidecars, WASM Components/WIT, and sandboxing remain separate decisions. A +sidecar is justified only when process isolation or a non-FFI runtime is a real +consumer requirement. + +## Review checklist + +When extending the host: + +- keep default builds free of plugin fields, branches, and linked code; +- keep `wacore` independent of the high-level plugin host; +- add capabilities narrowly; never expose the raw backend or Signal stores; +- choose install- or connection-scoped ownership explicitly for every task; +- keep core-event handlers non-blocking and custom-event queues bounded; +- preserve LIFO rollback/shutdown and per-generation close ordering; +- isolate faults so one plugin cannot strand unrelated cleanup; +- add plugin resource accounting and health degradation for new retained state + or failure modes; +- validate native and `wasm32-unknown-unknown` builds; +- measure both feature-disabled and enabled-with-no-plugin paths before claiming + zero overhead; +- defer interception/pre-ack work until its interaction with + `signal_durability.md` has a dedicated design. From 69aaca226ce90a89706013f25a4b464967129b66 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 19:50:23 -0300 Subject: [PATCH 33/46] fix(lifecycle): bind ready publication to its generation --- src/client/extension_lifecycle.rs | 65 +++++++++++++++++++++++++++---- src/client/lifecycle.rs | 38 ++++++++++++++---- src/client/node_io.rs | 4 +- 3 files changed, 89 insertions(+), 18 deletions(-) diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 4334fa525..6ff902d0e 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -1368,7 +1368,7 @@ mod tests { .store(GENERATION, Ordering::SeqCst); let registration = client.lifecycle.as_ref().expect("lifecycle registration"); assert!(registration.begin_scope_if_current(GENERATION, || true)); - client.dispatch_connected().await; + client.dispatch_connected(GENERATION).await; let scope = lifecycle .scopes @@ -1532,7 +1532,7 @@ mod tests { let ready_client = Arc::clone(&client); let ready_task = tokio::spawn(async move { - ready_client.dispatch_connected().await; + ready_client.dispatch_connected(GENERATION).await; }); ready_started_rx .recv() @@ -1607,7 +1607,7 @@ mod tests { tokio::time::timeout( std::time::Duration::from_secs(2), - client.dispatch_connected(), + client.dispatch_connected(GENERATION), ) .await .expect("reentrant disconnect completed"); @@ -1654,9 +1654,12 @@ mod tests { let registration = client.lifecycle.as_ref().expect("lifecycle registration"); assert!(registration.begin_scope_if_current(GENERATION, || true)); - tokio::time::timeout(Duration::from_secs(2), client.dispatch_connected()) - .await - .expect("reentrant reconnect completed"); + tokio::time::timeout( + Duration::from_secs(2), + client.dispatch_connected(GENERATION), + ) + .await + .expect("reentrant reconnect completed"); let scope = registration .scope_for(GENERATION) .expect("cancelled connection scope"); @@ -1951,7 +1954,7 @@ mod tests { .store(GENERATION, Ordering::SeqCst); let registration = client.lifecycle.as_ref().expect("lifecycle registration"); assert!(registration.begin_scope_if_current(GENERATION, || true)); - client.dispatch_connected().await; + client.dispatch_connected(GENERATION).await; let scope = registration .scope_for(GENERATION) .expect("connection scope"); @@ -2014,7 +2017,7 @@ mod tests { .store(GENERATION, Ordering::SeqCst); let registration = client.lifecycle.as_ref().expect("lifecycle registration"); assert!(registration.begin_scope_if_current(GENERATION, || true)); - client.dispatch_connected().await; + client.dispatch_connected(GENERATION).await; let scope = registration .scope_for(GENERATION) .expect("connection scope"); @@ -2063,6 +2066,52 @@ mod tests { assert_eq!(lifecycle.shutdowns.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn stale_connected_dispatch_cannot_claim_a_new_scope() { + let persistence_manager = Arc::new( + PersistenceManager::new(crate::test_utils::create_test_backend().await) + .await + .expect("persistence manager"), + ); + let lifecycle = Arc::new(RecordingLifecycle::default()); + let client = Client::builder() + .with_runtime(TokioRuntime) + .with_persistence_manager(persistence_manager) + .with_transport_factory(MockTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_lifecycle_arc(lifecycle.clone()) + .build() + .await + .expect("client build") + .into_client(); + const STALE_GENERATION: u64 = 60; + const CURRENT_GENERATION: u64 = 62; + client + .connection_generation + .store(CURRENT_GENERATION, Ordering::SeqCst); + let registration = client.lifecycle.as_ref().expect("lifecycle registration"); + assert!(registration.begin_scope_if_current(CURRENT_GENERATION, || true)); + + client.dispatch_connected(STALE_GENERATION).await; + + let scope = registration + .scope_for(CURRENT_GENERATION) + .expect("current scope"); + assert_eq!(scope.state(), ConnectionScopeState::Open); + assert!(!client.is_ready.load(Ordering::Relaxed)); + assert_eq!(lifecycle.events(), vec!["install"]); + + client.dispatch_connected(CURRENT_GENERATION).await; + assert_eq!(scope.state(), ConnectionScopeState::Ready); + assert!(client.is_ready.load(Ordering::Relaxed)); + assert_eq!(lifecycle.events(), vec!["install", "ready:62"]); + + registration.cancel_scope(CURRENT_GENERATION); + registration.close_scope(CURRENT_GENERATION); + registration.shutdown().await; + client.signal_shutdown_sync(); + } + #[tokio::test] async fn rejected_success_restores_logged_out_state() { let persistence_manager = Arc::new( diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index d0bf8fa0f..ee597c183 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -81,27 +81,49 @@ impl Client { self.is_connected() && self.is_logged_in() && self.is_ready.load(Ordering::Relaxed) } - /// Dispatch the Connected event and notify waiters. - pub(crate) async fn dispatch_connected(&self) { + /// Dispatch the Connected event and notify waiters for the originating connection. + pub(crate) async fn dispatch_connected(&self, expected_generation: u64) { #[cfg(feature = "client-lifecycle")] { - let generation = self.connection_generation.load(Ordering::SeqCst); if let Some(lifecycle) = &self.lifecycle { - if !lifecycle.ready(generation).await { - debug!("Skipping Connected dispatch for retired generation {generation}"); + if !lifecycle.ready(expected_generation).await { + debug!( + "Skipping Connected dispatch for retired generation {expected_generation}" + ); return; } - if self.connection_generation.load(Ordering::SeqCst) != generation { - debug!("Skipping Connected dispatch after generation changed"); + + // Cleanup takes the same lock before retiring the generation, so the final + // validation and publication form one transition with its generation bump. + let _login_transition = self + .login_transition + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.connection_generation.load(Ordering::SeqCst) != expected_generation + || self.expected_disconnect.load(Ordering::Acquire) + { + debug!( + "Skipping Connected dispatch after generation {expected_generation} retired" + ); return; } - if !lifecycle.publish_ready(generation, || self.publish_connected()) { + if !lifecycle.publish_ready(expected_generation, || self.publish_connected()) { debug!("Skipping Connected dispatch after lifecycle cancellation"); } return; } } + let _login_transition = self + .login_transition + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.connection_generation.load(Ordering::SeqCst) != expected_generation + || self.expected_disconnect.load(Ordering::Acquire) + { + debug!("Skipping Connected dispatch after its connection retired"); + return; + } self.publish_connected(); } diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 6ae49df84..49bf1b4ab 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -1121,7 +1121,7 @@ impl Client { // Presence is NOT sent here — WhatsApp Web sends presence from the // setting_pushName mutation handler (WAWebPushNameSync), not from // criticalSyncDone. Our setting_pushName handler already does this. - client_clone.dispatch_connected().await; + client_clone.dispatch_connected(task_generation).await; } Err(e) => { client_clone.log_sync_error("critical app state sync", &e); @@ -1183,7 +1183,7 @@ impl Client { // for an outdated connection that was replaced mid-await. check_generation!(); - client_clone.dispatch_connected().await; + client_clone.dispatch_connected(task_generation).await; } })).detach(); } From fb263ec25e2495fb59ff0679c8b62d6f00aaa532 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 19:56:49 -0300 Subject: [PATCH 34/46] fix(plugins): attribute spawned task panics --- src/plugins/mod.rs | 233 +++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 223 insertions(+), 10 deletions(-) diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 4077d4687..acb00bd84 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -258,6 +258,8 @@ pub struct PluginStats { pub callback_failures: u64, pub callback_timeouts: u64, pub task_drain_timeouts: u64, + /// Spawned workers that panicked while running or being cancelled. + pub task_panics: u64, /// Core-event handler calls that returned without panicking. pub core_events_delivered: u64, /// Panics isolated before they could unwind through the client's event dispatcher. @@ -619,6 +621,7 @@ struct PluginDiagnostics { callback_failures: AtomicU64, callback_timeouts: AtomicU64, task_drain_timeouts: AtomicU64, + task_panics: AtomicU64, core_events_delivered: AtomicU64, core_event_panics: AtomicU64, shutdown_complete: AtomicBool, @@ -632,6 +635,7 @@ impl PluginDiagnostics { callback_failures: AtomicU64::new(0), callback_timeouts: AtomicU64::new(0), task_drain_timeouts: AtomicU64::new(0), + task_panics: AtomicU64::new(0), core_events_delivered: AtomicU64::new(0), core_event_panics: AtomicU64::new(0), shutdown_complete: AtomicBool::new(false), @@ -690,6 +694,7 @@ impl PluginDiagnostics { let callback_failures = self.callback_failures.load(Ordering::Relaxed); let callback_timeouts = self.callback_timeouts.load(Ordering::Relaxed); let task_drain_timeouts = self.task_drain_timeouts.load(Ordering::Relaxed); + let task_panics = self.task_panics.load(Ordering::Relaxed); let core_events_delivered = self.core_events_delivered.load(Ordering::Relaxed); let core_event_panics = self.core_event_panics.load(Ordering::Relaxed); let state = if self.shutdown_complete.load(Ordering::Acquire) { @@ -707,6 +712,7 @@ impl PluginDiagnostics { let health = if callback_failures > 0 || callback_timeouts > 0 || task_drain_timeouts > 0 + || task_panics > 0 || core_event_panics > 0 || resources.teardown_panics > 0 || event_degraded @@ -723,6 +729,7 @@ impl PluginDiagnostics { callback_failures, callback_timeouts, task_drain_timeouts, + task_panics, core_events_delivered, core_event_panics, resource_teardown_panics: resources.teardown_panics, @@ -743,6 +750,8 @@ impl PluginDiagnostics { pub struct PluginTasks { runtime: Arc, resources: Arc, + diagnostics: Arc, + plugin_id: Arc, } impl PluginTasks { @@ -754,7 +763,14 @@ impl PluginTasks { return Err(PluginResourceError::ShuttingDown); } let lease = self.resources.install_tasks.register()?; - spawn_after_activation(&self.runtime, Arc::clone(&self.resources), lease, future); + spawn_after_activation( + &self.runtime, + Arc::clone(&self.resources), + Arc::clone(&self.diagnostics), + Arc::clone(&self.plugin_id), + lease, + future, + ); Ok(()) } @@ -985,6 +1001,8 @@ pub struct PluginConnectionTasks { runtime: Arc, scope: ConnectionScope, tracker: Arc, + diagnostics: Arc, + plugin_id: Arc, } impl PluginConnectionTasks { @@ -999,6 +1017,8 @@ impl PluginConnectionTasks { spawn_until_cancelled( &self.runtime, self.scope.cancellation_signal(), + Arc::clone(&self.diagnostics), + Arc::clone(&self.plugin_id), lease, future, ); @@ -1022,19 +1042,100 @@ impl PluginConnectionTasks { } } +struct GuardedPluginTask> { + future: Option>>, + diagnostics: Arc, + plugin_id: Arc, + failure_recorded: bool, +} + +impl GuardedPluginTask +where + F: Future, +{ + fn new(future: F, diagnostics: Arc, plugin_id: Arc) -> Self { + Self { + future: Some(Box::pin(future)), + diagnostics, + plugin_id, + failure_recorded: false, + } + } + + fn record_panic(&mut self, stage: &str) { + if !self.failure_recorded { + self.failure_recorded = true; + self.diagnostics.task_panics.fetch_add(1, Ordering::Relaxed); + } + log::warn!("Plugin `{}` task panicked {stage}", self.plugin_id); + } + + fn drop_future(&mut self) -> bool { + let future = self.future.take(); + std::panic::catch_unwind(AssertUnwindSafe(|| drop(future))).is_err() + } +} + +impl Unpin for GuardedPluginTask where F: Future {} + +impl Future for GuardedPluginTask +where + F: Future, +{ + type Output = (); + + fn poll( + self: std::pin::Pin<&mut Self>, + context: &mut std::task::Context<'_>, + ) -> std::task::Poll { + let this = self.get_mut(); + let Some(future) = this.future.as_mut() else { + return std::task::Poll::Ready(()); + }; + let result = std::panic::catch_unwind(AssertUnwindSafe(|| future.as_mut().poll(context))); + match result { + Ok(std::task::Poll::Pending) => std::task::Poll::Pending, + Ok(std::task::Poll::Ready(())) => { + if this.drop_future() { + this.record_panic("after completion"); + } + std::task::Poll::Ready(()) + } + Err(_) => { + this.record_panic("while running"); + if this.drop_future() { + this.record_panic("while cleaning up after failure"); + } + std::task::Poll::Ready(()) + } + } + } +} + +impl> Drop for GuardedPluginTask { + fn drop(&mut self) { + if self.drop_future() { + self.record_panic("while being cancelled"); + } + } +} + fn spawn_until_cancelled( runtime: &Arc, cancellation: ShutdownSignal, + diagnostics: Arc, + plugin_id: Arc, lease: TaskLease, future: F, ) where F: Future + Spawnable, { + let work = GuardedPluginTask::new(future, diagnostics, plugin_id); runtime .spawn(Box::pin(async move { let _lease = lease; let cancelled = Box::pin(wait_for_shutdown(&cancellation)); - let work = Box::pin(future); + let work = Box::pin(work); let _ = futures::future::select(cancelled, work).await; })) .detach(); @@ -1043,6 +1144,8 @@ fn spawn_until_cancelled( fn spawn_after_activation( runtime: &Arc, resources: Arc, + diagnostics: Arc, + plugin_id: Arc, lease: TaskLease, future: F, ) where @@ -1050,6 +1153,7 @@ fn spawn_after_activation( { let activation = resources.activation.subscribe(); let cancellation = resources.shutdown.subscribe(); + let work = GuardedPluginTask::new(future, diagnostics, plugin_id); runtime .spawn(Box::pin(async move { let _lease = lease; @@ -1065,7 +1169,7 @@ fn spawn_after_activation( return; } let cancelled = Box::pin(wait_for_shutdown(&cancellation)); - let work = Box::pin(future); + let work = Box::pin(work); let _ = futures::future::select(cancelled, work).await; })) .detach(); @@ -1578,6 +1682,7 @@ impl PluginHost { } = parts; let manifest = &planned.manifest; let capabilities = manifest.capabilities; + let plugin_id: Arc = Arc::from(manifest.id.as_str()); PluginContext { plugin_id: manifest.id.clone(), dependencies: apis.dependency_view(&planned.dependency_markers), @@ -1586,7 +1691,7 @@ impl PluginHost { .then(|| PluginCoreEvents { client: client.clone(), resources: Arc::clone(&resources), - plugin_id: Arc::from(manifest.id.as_str()), + plugin_id: Arc::clone(&plugin_id), diagnostics: Arc::clone(&diagnostics), }), tasks: capabilities @@ -1594,6 +1699,8 @@ impl PluginHost { .then(|| PluginTasks { runtime: Arc::clone(&runtime), resources: Arc::clone(&resources), + diagnostics: Arc::clone(&diagnostics), + plugin_id, }), messaging: capabilities.contains(PluginCapability::Messaging).then(|| { PluginMessaging { @@ -1625,10 +1732,14 @@ impl PluginHost { fn connection_scope( &self, scope: ConnectionScope, - manifest: &PluginManifest, + plugin: &InstalledPlugin, task_tracker: Option>, ) -> PluginConnectionScope { - let tasks = if manifest.capabilities.contains(PluginCapability::Tasks) { + let tasks = if plugin + .manifest + .capabilities + .contains(PluginCapability::Tasks) + { self.runtime .get() .cloned() @@ -1637,6 +1748,8 @@ impl PluginHost { runtime, scope: scope.clone(), tracker, + diagnostics: Arc::clone(&plugin.diagnostics), + plugin_id: Arc::from(plugin.manifest.id.as_str()), }) } else { None @@ -1867,8 +1980,7 @@ impl ClientLifecycle for PluginHost { } tracker }); - let plugin_scope = - self.connection_scope(scope.clone(), &plugin.manifest, task_tracker); + let plugin_scope = self.connection_scope(scope.clone(), plugin, task_tracker); let result = self .run_callback(|| plugin.plugin.on_ready(plugin_scope)) .await; @@ -1904,8 +2016,7 @@ impl ClientLifecycle for PluginHost { } } } - let plugin_scope = - self.connection_scope(scope.clone(), &plugin.manifest, task_tracker); + let plugin_scope = self.connection_scope(scope.clone(), plugin, task_tracker); let result = self .run_callback(|| plugin.plugin.on_closed(plugin_scope)) .await; @@ -2641,6 +2752,25 @@ mod tests { } } + struct PendingDropPanic; + + impl Future for PendingDropPanic { + type Output = (); + + fn poll( + self: std::pin::Pin<&mut Self>, + _context: &mut std::task::Context<'_>, + ) -> std::task::Poll { + std::task::Poll::Pending + } + } + + impl Drop for PendingDropPanic { + fn drop(&mut self) { + panic!("injected task cancellation panic"); + } + } + struct PanickingDropApi; impl Drop for PanickingDropApi { @@ -3336,6 +3466,8 @@ mod tests { runtime: Arc::new(TokioRuntime), scope: scope.clone(), tracker: TaskTracker::new(), + diagnostics: PluginDiagnostics::new(), + plugin_id: Arc::from("connection-task-test"), }; let guard = DropFlag(task_dropped.clone()); tasks @@ -3359,6 +3491,83 @@ mod tests { )); } + #[tokio::test] + async fn spawned_task_panics_are_isolated_and_degrade_health() { + let diagnostics = PluginDiagnostics::new(); + let plugin_id: Arc = Arc::from("panicking-task-test"); + let resources = PluginResources::new(); + diagnostics.attach_resources(&resources); + resources.activate(); + let install_tasks = PluginTasks { + runtime: Arc::new(TokioRuntime), + resources: resources.clone(), + diagnostics: diagnostics.clone(), + plugin_id: plugin_id.clone(), + }; + install_tasks + .spawn(async { panic!("injected install task panic") }) + .expect("spawn install task"); + + tokio::time::timeout(Duration::from_secs(1), async { + while diagnostics.task_panics.load(Ordering::Relaxed) < 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("install task panic was not recorded"); + assert_eq!(resources.install_tasks.active(), 0); + + let connection_scope = ConnectionScope::new(101); + let connection_tracker = TaskTracker::new(); + let connection_tasks = PluginConnectionTasks { + runtime: Arc::new(TokioRuntime), + scope: connection_scope, + tracker: connection_tracker.clone(), + diagnostics: diagnostics.clone(), + plugin_id: plugin_id.clone(), + }; + connection_tasks + .spawn(async { panic!("injected connection task panic") }) + .expect("spawn connection task"); + + tokio::time::timeout(Duration::from_secs(1), async { + while diagnostics.task_panics.load(Ordering::Relaxed) < 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("connection task panic was not recorded"); + assert_eq!(connection_tracker.active(), 0); + + let cancellation_scope = ConnectionScope::new(102); + let cancellation_tracker = TaskTracker::new(); + let cancellation_tasks = PluginConnectionTasks { + runtime: Arc::new(TokioRuntime), + scope: cancellation_scope.clone(), + tracker: cancellation_tracker.clone(), + diagnostics: diagnostics.clone(), + plugin_id, + }; + cancellation_tasks + .spawn(PendingDropPanic) + .expect("spawn cancellation task"); + cancellation_scope.cancel(); + + tokio::time::timeout(Duration::from_secs(1), async { + while diagnostics.task_panics.load(Ordering::Relaxed) < 3 { + tokio::task::yield_now().await; + } + }) + .await + .expect("task cancellation panic was not recorded"); + assert_eq!(cancellation_tracker.active(), 0); + + let stats = diagnostics.snapshot("panicking-task-test", false, None); + assert_eq!(stats.task_panics, 3); + assert_eq!(stats.health, PluginHealth::Degraded); + resources.close(); + } + #[tokio::test] async fn task_sleeps_return_when_their_owner_is_cancelled() { let resources = PluginResources::new(); @@ -3366,6 +3575,8 @@ mod tests { let install_tasks = PluginTasks { runtime: Arc::new(TokioRuntime), resources: resources.clone(), + diagnostics: PluginDiagnostics::new(), + plugin_id: Arc::from("install-sleep-test"), }; let install_sleeper = tokio::spawn(async move { install_tasks.sleep(Duration::from_secs(60)).await }); @@ -3384,6 +3595,8 @@ mod tests { runtime: Arc::new(TokioRuntime), scope: scope.clone(), tracker: TaskTracker::new(), + diagnostics: PluginDiagnostics::new(), + plugin_id: Arc::from("connection-sleep-test"), }; let connection_sleeper = tokio::spawn(async move { connection_tasks.sleep(Duration::from_secs(60)).await }); From a09f00706e3725bb53e789f9f664276742faee3e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 20:01:35 -0300 Subject: [PATCH 35/46] docs(plugins): clarify feature and task contracts --- Cargo.toml | 7 ++++++ README.md | 7 ++++++ agent_docs/observability.md | 8 +++---- agent_docs/plugin_architecture.md | 37 ++++++++++++++++++++++--------- src/bot.rs | 2 ++ src/client.rs | 1 + src/client/builder.rs | 7 ++++++ src/lib.rs | 6 +++++ src/plugins/mod.rs | 2 +- 9 files changed, 61 insertions(+), 16 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index bfa2aacc9..be907c3f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,10 @@ repository = "https://github.com/jlucaso1/whatsapp-rust" readme = "README.md" description = "Rust client for WhatsApp Web" +[package.metadata.docs.rs] +features = ["plugins"] +rustdoc-args = ["--cfg", "docsrs"] + [workspace] members = [ ".", @@ -50,6 +54,9 @@ disallowed_methods = "deny" # portable_atomic. Host-only test/bench counters carry an inline allow. disallowed_types = "deny" +[workspace.lints.rust] +unexpected_cfgs = { level = "warn", check-cfg = ["cfg(docsrs)"] } + [workspace.dependencies] # Shared dependencies aes = "0.9.1" diff --git a/README.md b/README.md index a8f393e98..9c9361dc1 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ A high-performance, async Rust library for the WhatsApp Web API. Inspired by [wh - **Profile** — Set push name, status text, profile picture - **Privacy** — Fetch/set privacy settings, disappearing messages - **Modular** — Pluggable storage, transport, HTTP client, and async runtime; SQLite, Tokio WebSocket, and ureq ship as the defaults, swap any of them with `default-features = false` +- **Native plugins** — Build-time, type-safe extensions with scoped capabilities and lifecycle ownership behind the `plugins` feature - **Runtime agnostic** — Bring your own async runtime via the `Runtime` trait (Tokio included by default) For the full API reference and guides, see the **[documentation](https://whatsapp-rust.jlucaso.com)**. @@ -60,6 +61,12 @@ async fn main() -> Result<(), Box> { The default cargo features wire up the Tokio WebSocket transport, the ureq HTTP client, the SQLite store, and the Tokio runtime; only the storage backend has to be chosen explicitly. Every piece is replaceable through the builder (`with_transport_factory`, `with_http_client`, `with_runtime`) for custom environments such as wasm or embedded targets. +Native plugin APIs are opt-in: use `features = ["plugins"]` when implementing a +plugin in the application. Published plugin crates can enable that feature in +their own `whatsapp-rust` dependency, and Cargo feature unification activates it +for the consumer. See [`agent_docs/plugin_architecture.md`](agent_docs/plugin_architecture.md) +for the host contract and type-safe API example. + ### One dependency is enough `whatsapp-rust` re-exports the whole stack, so you never need to declare the sibling crates (`wacore`, `wacore-binary`, `waproto`, `whatsapp-rust-tokio-transport`, `whatsapp-rust-ureq-http-client`, `whatsapp-rust-sqlite-storage`) yourself, including when pinning a git revision: diff --git a/agent_docs/observability.md b/agent_docs/observability.md index bb77c71f5..f0030c034 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -78,10 +78,10 @@ queue retention, and cumulative delivery/backpressure totals; publishers can read their own totals through `PluginEvents::stats()`. Health is sticky for the lifetime of the host: lifecycle errors/panics, -timeouts, task-drain timeouts, isolated core-event panics, resource teardown -panics, publication failures, and queue drops mark only the responsible plugin -as degraded. Concurrent snapshots are intentionally approximate, and carry no -message content, JIDs, or phone numbers. +timeouts, spawned-task panics, task-drain timeouts, isolated core-event panics, +resource teardown panics, publication failures, and queue drops mark only the +responsible plugin as degraded. Concurrent snapshots are intentionally +approximate, and carry no message content, JIDs, or phone numbers. ### 3. `BotBuilder::with_task_instrument` — CPU / custom attribution (opt-in) diff --git a/agent_docs/plugin_architecture.md b/agent_docs/plugin_architecture.md index 352881407..d09832b34 100644 --- a/agent_docs/plugin_architecture.md +++ b/agent_docs/plugin_architecture.md @@ -27,6 +27,15 @@ high-level client operations and lifecycle coordination. It is enabled by the opt-in `plugins` feature, which enables `client-lifecycle`. A default build has neither plugin/lifecycle fields nor their runtime branches. +The public feature surface stays intentionally small: `plugins` is the normal +opt-in and `client-lifecycle` is the advanced low-level seam for hosts that need +lifecycle integration without the native plugin host. Capabilities and +individual plugins do not become Cargo features. An external plugin crate can +enable `whatsapp-rust/plugins` in its own dependency, so Cargo feature unification +activates the host for its consumer. The host remains opt-in because LTO is not +a compatibility guarantee for client layout, reachable branches, dependencies, +compile time, or final binary size. + ## Construction boundary `ClientBuilder` is the canonical low-level construction path. It validates @@ -165,16 +174,20 @@ install once - `on_ready` runs in dependency order after authentication for that generation. - scope cancellation is synchronous when a reconnect or terminal teardown starts. -- generation tasks drain before `on_closed`; `on_closed` runs in reverse - dependency order after authoritative cleanup. -- install tasks drain before terminal `shutdown`; shutdown hooks run in reverse - dependency order. +- generation task cancellation is signalled before `on_closed`; the host waits + for the drain up to its configured timeout, then continues in reverse + dependency order and marks the plugin degraded if work remains. +- install task cancellation follows the same bounded drain before terminal + `shutdown`; shutdown hooks run in reverse dependency order. - the separately configured upstream lifecycle wraps the plugin order: it is readied first and closed/shut down last. -`PluginTasks` survives reconnects and ends only on rollback or terminal -shutdown. `PluginConnectionTasks` is tied to one generation and must not leak -work into a later connection. +`PluginTasks` survives reconnects and receives cancellation only on rollback or +terminal shutdown. `PluginConnectionTasks` is tied to one generation: +cancellation is signalled synchronously, while actual future destruction is +cooperative at executor poll boundaries. Plugin tasks must not block an executor +thread or detach untracked work. A non-cooperative task may outlive the bounded +drain, in which case teardown proceeds and diagnostics remain degraded. Callbacks are serialized, bounded by timeouts, and isolated from panics, including panics while constructing, polling, cancelling, or destroying their @@ -224,10 +237,12 @@ already queued envelopes. ## Diagnostics `Client::plugin_stats()` reports per-plugin lifecycle state, sticky health, -callback/task failures, active task scopes, subscriptions, and publisher -counters. `PluginEventRouter::stats()` and `PluginEvents::stats()` expose queue -and backpressure totals. `Client::memory_report()` includes plugin resources -and counts a shared queued payload once across fanout. +callback failures, spawned-task panics, drain timeouts, active task scopes, +subscriptions, and publisher counters. Spawned task panics are isolated during +polling and cancellation so a dead worker cannot remain falsely healthy. +`PluginEventRouter::stats()` and `PluginEvents::stats()` expose queue and +backpressure totals. `Client::memory_report()` includes plugin resources and +counts a shared queued payload once across fanout. Snapshots are on-demand, approximate under concurrency, and contain no JIDs, phone numbers, or message bodies. See `observability.md` for accounting rules. diff --git a/src/bot.rs b/src/bot.rs index ac36db814..c772d40a4 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -815,6 +815,7 @@ impl BotBuilder { /// Register a native plugin without changing the builder's typestate. #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] pub fn with_plugin(mut self, plugin: P) -> Self { self.plugins.push(PluginRegistration::new(plugin)); self @@ -822,6 +823,7 @@ impl BotBuilder { /// Register an already-shared native plugin without changing its marker type. #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] pub fn with_plugin_arc(mut self, plugin: Arc

) -> Self { self.plugins.push(PluginRegistration::new_arc(plugin)); self diff --git a/src/client.rs b/src/client.rs index 9dcec9ac7..ed7898597 100644 --- a/src/client.rs +++ b/src/client.rs @@ -21,6 +21,7 @@ pub use builder::{ClientBuild, ClientBuilder, ClientBuilderError}; #[cfg(feature = "client-lifecycle")] use extension_lifecycle::LifecycleRegistration; #[cfg(feature = "client-lifecycle")] +#[cfg_attr(docsrs, doc(cfg(feature = "client-lifecycle")))] pub use extension_lifecycle::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; pub use voip::{CallError, Voip}; diff --git a/src/client/builder.rs b/src/client/builder.rs index a266ac754..734e8e17c 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -70,12 +70,15 @@ pub enum ClientBuilderError { #[error("the configured backend does not support the inbound durability hook: {0}")] UnsupportedDurabilityBackend(String), #[cfg(feature = "client-lifecycle")] + #[cfg_attr(docsrs, doc(cfg(feature = "client-lifecycle")))] #[error("client lifecycle installation failed: {0}")] LifecycleInstall(#[source] anyhow::Error), #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] #[error("plugin host installation failed: {0}")] PluginInstall(#[source] anyhow::Error), #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] #[error("invalid plugin plan: {0}")] PluginPlan(#[from] PluginPlanError), } @@ -284,6 +287,7 @@ impl ClientBuilder { /// Install the aggregate lifecycle used by extensions of this client. #[cfg(feature = "client-lifecycle")] + #[cfg_attr(docsrs, doc(cfg(feature = "client-lifecycle")))] pub fn with_lifecycle(mut self, lifecycle: L) -> Self where L: ClientLifecycle + 'static, @@ -294,6 +298,7 @@ impl ClientBuilder { /// Install an already-shared aggregate lifecycle. #[cfg(feature = "client-lifecycle")] + #[cfg_attr(docsrs, doc(cfg(feature = "client-lifecycle")))] pub fn with_lifecycle_arc(mut self, lifecycle: Arc) -> Self { self.lifecycle = Some(lifecycle); self @@ -301,6 +306,7 @@ impl ClientBuilder { /// Register a native plugin for transactional installation before services start. #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] pub fn with_plugin(mut self, plugin: P) -> Self { self.plugins.push(PluginRegistration::new(plugin)); self @@ -308,6 +314,7 @@ impl ClientBuilder { /// Register an already-shared native plugin without changing its marker type. #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] pub fn with_plugin_arc(mut self, plugin: Arc

) -> Self { self.plugins.push(PluginRegistration::new_arc(plugin)); self diff --git a/src/lib.rs b/src/lib.rs index 7d2e31528..012d2d845 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,7 @@ // Compile-checks the README examples as doctests, so the advertised quick // start can never silently rot. #![doc = include_str!("../README.md")] +#![cfg_attr(docsrs, feature(doc_cfg))] // Instrumenting large async fns (e.g. process_sync_task) wraps them in deep // `Instrumented` future types; the default depth limit overflows when the // `tracing` + `tracing-pii` paths combine. Raise it (compile-time only). @@ -95,6 +96,7 @@ pub use client::{ pub use client::{CallError, Voip}; pub use client::{Client, ClientBuild, ClientBuilder, ClientBuilderError, RawNodeLease}; #[cfg(feature = "client-lifecycle")] +#[cfg_attr(docsrs, doc(cfg(feature = "client-lifecycle")))] pub use client::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; pub use types::durability_hook::InboundDurabilityHook; pub use types::retry_admission::RetryAdmission; @@ -111,8 +113,10 @@ pub mod pair; pub mod pair_code; pub mod passkey; #[cfg(feature = "plugins")] +#[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] pub mod plugins; #[cfg(feature = "plugins")] +#[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] pub use plugins::{ ClientPlugin, PluginCapabilities, PluginCapability, PluginConnectionScope, PluginConnectionTasks, PluginContext, PluginCoreEvents, PluginEventEndpointConfig, @@ -194,8 +198,10 @@ pub mod prelude { pub use crate::bot::{Bot, BotBuilder, BotHandle, EventDelivery, MessageContext}; pub use crate::client::{Client, ClientBuilder, ClientBuilderError, ClientError, RawNodeLease}; #[cfg(feature = "client-lifecycle")] + #[cfg_attr(docsrs, doc(cfg(feature = "client-lifecycle")))] pub use crate::client::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] pub use crate::plugins::{ ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, PluginEventEndpointConfig, PluginEventOverflow, PluginEventPayloadEncoding, diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index acb00bd84..2adac129e 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -995,7 +995,7 @@ impl PluginConnectionScope { } } -/// Task capability whose work is aborted synchronously when its generation retires. +/// Task capability whose cancellation is signalled synchronously when its generation retires. #[derive(Clone)] pub struct PluginConnectionTasks { runtime: Arc, From 7f025d32e38c69680267e35b198365113a8a3e1c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 21:00:20 -0300 Subject: [PATCH 36/46] fix(lifecycle): gate direct connect during construction --- src/client/builder.rs | 187 ++++++++++++++++++++++++++++++++++++++++ src/client/lifecycle.rs | 6 ++ 2 files changed, 193 insertions(+) diff --git a/src/client/builder.rs b/src/client/builder.rs index 734e8e17c..5b4d7ed62 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -650,6 +650,45 @@ mod tests { run_finished: async_channel::Sender<()>, } + #[cfg(feature = "client-lifecycle")] + struct ConnectDuringInstallLifecycle { + client: async_channel::Sender>, + release: async_channel::Receiver<()>, + connect_invoked: async_channel::Sender<()>, + connect_finished: async_channel::Sender, + } + + #[cfg(feature = "client-lifecycle")] + struct BlockingTransportFactory { + started: async_channel::Sender<()>, + release: async_channel::Receiver<()>, + } + + #[cfg(feature = "client-lifecycle")] + #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] + #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] + impl TransportFactory for BlockingTransportFactory { + async fn create_transport( + &self, + ) -> Result< + ( + Arc, + async_channel::Receiver, + ), + anyhow::Error, + > { + self.started + .send(()) + .await + .map_err(|_| anyhow::anyhow!("transport-start receiver closed"))?; + self.release + .recv() + .await + .map_err(|_| anyhow::anyhow!("transport release closed"))?; + Err(anyhow::anyhow!("injected transport stop")) + } + } + #[cfg(feature = "client-lifecycle")] impl ClientLifecycle for RunDuringInstallLifecycle { fn install<'a>( @@ -682,6 +721,40 @@ mod tests { } } + #[cfg(feature = "client-lifecycle")] + impl ClientLifecycle for ConnectDuringInstallLifecycle { + fn install<'a>( + &'a self, + client: std::sync::Weak, + ) -> wacore::runtime::BoxFuture<'a, anyhow::Result<()>> { + Box::pin(async move { + let client = client + .upgrade() + .ok_or_else(|| anyhow::anyhow!("client unavailable during install"))?; + let connect_client = client.clone(); + let connect_invoked = self.connect_invoked.clone(); + let connect_finished = self.connect_finished.clone(); + client + .runtime + .spawn(Box::pin(async move { + let _ = connect_invoked.send(()).await; + let failed = connect_client.connect().await.is_err(); + let _ = connect_finished.send(failed).await; + })) + .detach(); + self.client + .send(client) + .await + .map_err(|_| anyhow::anyhow!("test client receiver closed"))?; + self.release + .recv() + .await + .map_err(|_| anyhow::anyhow!("test install release closed"))?; + Ok(()) + }) + } + } + #[cfg(feature = "client-lifecycle")] impl ClientLifecycle for FailingLifecycle { fn install<'a>( @@ -848,6 +921,120 @@ mod tests { .expect("run stopped"); } + #[tokio::test] + #[cfg(feature = "client-lifecycle")] + async fn connect_leaked_during_install_waits_for_complete_construction() { + let (client_tx, client_rx) = async_channel::bounded(1); + let (install_release_tx, install_release_rx) = async_channel::bounded(1); + let (connect_invoked_tx, connect_invoked_rx) = async_channel::bounded(1); + let (connect_finished_tx, connect_finished_rx) = async_channel::bounded(1); + let (transport_started_tx, transport_started_rx) = async_channel::bounded(1); + let (transport_release_tx, transport_release_rx) = async_channel::bounded(1); + let builder = complete_builder() + .await + .with_transport_factory(BlockingTransportFactory { + started: transport_started_tx, + release: transport_release_rx, + }) + .with_lifecycle(ConnectDuringInstallLifecycle { + client: client_tx, + release: install_release_rx, + connect_invoked: connect_invoked_tx, + connect_finished: connect_finished_tx, + }); + let build = tokio::spawn(async move { builder.build().await }); + let leaked_client = client_rx + .recv() + .await + .expect("client leaked during install"); + connect_invoked_rx + .recv() + .await + .expect("direct connect invoked"); + + assert!( + tokio::time::timeout(Duration::from_millis(100), transport_started_rx.recv()) + .await + .is_err(), + "transport started before construction activation" + ); + assert!(!leaked_client.is_connecting.load(Ordering::Acquire)); + + install_release_tx + .send(()) + .await + .expect("release installation"); + let client = build + .await + .expect("builder task") + .expect("successful build") + .into_client(); + tokio::time::timeout(Duration::from_secs(1), transport_started_rx.recv()) + .await + .expect("connect remained gated after activation") + .expect("transport-start sender closed"); + transport_release_tx + .send(()) + .await + .expect("release transport"); + assert!( + tokio::time::timeout(Duration::from_secs(1), connect_finished_rx.recv()) + .await + .expect("direct connect did not finish") + .expect("connect-finished sender closed") + ); + client.signal_shutdown_sync(); + } + + #[tokio::test] + #[cfg(feature = "client-lifecycle")] + async fn shutdown_during_install_rejects_leaked_connect() { + let (client_tx, client_rx) = async_channel::bounded(1); + let (_install_release_tx, install_release_rx) = async_channel::bounded(1); + let (connect_invoked_tx, connect_invoked_rx) = async_channel::bounded(1); + let (connect_finished_tx, connect_finished_rx) = async_channel::bounded(1); + let (transport_started_tx, transport_started_rx) = async_channel::bounded(1); + let (_transport_release_tx, transport_release_rx) = async_channel::bounded(1); + let builder = complete_builder() + .await + .with_transport_factory(BlockingTransportFactory { + started: transport_started_tx, + release: transport_release_rx, + }) + .with_lifecycle(ConnectDuringInstallLifecycle { + client: client_tx, + release: install_release_rx, + connect_invoked: connect_invoked_tx, + connect_finished: connect_finished_tx, + }); + let build = tokio::spawn(async move { builder.build().await }); + let leaked_client = client_rx + .recv() + .await + .expect("client leaked during install"); + connect_invoked_rx + .recv() + .await + .expect("direct connect invoked"); + leaked_client.signal_shutdown_sync(); + + assert!(matches!( + tokio::time::timeout(Duration::from_secs(2), build) + .await + .expect("lifecycle install ignored terminal shutdown") + .expect("builder task"), + Err(ClientBuilderError::LifecycleInstall(_)) + )); + assert!( + tokio::time::timeout(Duration::from_secs(1), connect_finished_rx.recv()) + .await + .expect("direct connect did not observe rejection") + .expect("connect-finished sender closed") + ); + assert!(transport_started_rx.try_recv().is_err()); + assert!(!leaked_client.is_connecting.load(Ordering::Acquire)); + } + #[tokio::test] #[cfg(feature = "client-lifecycle")] async fn shutdown_during_install_rejects_leaked_run_and_the_build() { diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index ee597c183..a2d9beee8 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -576,6 +576,12 @@ impl Client { /// across crates, so consumers awaiting the connect graph directly would /// re-codegen it; the box makes them poll through a vtable instead. pub async fn connect(self: &Arc) -> Result<(), anyhow::Error> { + #[cfg(feature = "client-lifecycle")] + if let Some(lifecycle) = &self.lifecycle + && !lifecycle.wait_until_active().await + { + return Err(anyhow!("client construction did not activate")); + } self.connect_boxed().await } From b4b9c63cdd5b3d9985087a0b719b407729db543f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 21:02:05 -0300 Subject: [PATCH 37/46] fix(plugins): attribute evictions to event owners --- src/plugins/events.rs | 42 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/src/plugins/events.rs b/src/plugins/events.rs index b622339fb..f70408dad 100644 --- a/src/plugins/events.rs +++ b/src/plugins/events.rs @@ -195,6 +195,10 @@ pub enum PluginEventPublishError { } /// Result of one non-blocking fan-out attempt. +/// +/// `dropped` counts queue entries discarded while processing this call. Under `DropOldest`, the +/// discarded entry may belong to an earlier publication from another namespace; cumulative +/// publisher statistics attribute that loss to the discarded envelope's owner. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] #[non_exhaustive] pub struct PluginEventPublishReport { @@ -219,6 +223,7 @@ pub struct PluginEventEndpointStats { /// /// `published` counts successful calls, including calls with no subscriber. Fanout fields count /// endpoint outcomes; `delivered` advances only when a receiver removes an envelope from its queue. +/// `dropped` follows the discarded envelope, including cross-namespace `DropOldest` eviction. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] #[non_exhaustive] pub struct PluginEventPublisherStats { @@ -301,7 +306,6 @@ impl PublisherCounters { self.published.fetch_add(1, Ordering::Relaxed); self.matched.fetch_add(report.matched, Ordering::Relaxed); self.enqueued.fetch_add(report.enqueued, Ordering::Relaxed); - self.dropped.fetch_add(report.dropped, Ordering::Relaxed); self.closed.fetch_add(report.closed, Ordering::Relaxed); } Err(_) => { @@ -383,8 +387,9 @@ impl EventEndpoint { self.enqueued.fetch_add(1, Ordering::Relaxed); EnqueueOutcome::Enqueued } - Err(TrySendError::Full(_)) => { + Err(TrySendError::Full(dropped)) => { self.dropped.fetch_add(1, Ordering::Relaxed); + dropped.publisher.dropped.fetch_add(1, Ordering::Relaxed); EnqueueOutcome::Dropped } Err(TrySendError::Closed(_)) => EnqueueOutcome::Closed, @@ -392,8 +397,9 @@ impl EventEndpoint { PluginEventOverflow::DropOldest => match self.sender.force_send(event) { Ok(evicted) => { self.enqueued.fetch_add(1, Ordering::Relaxed); - if evicted.is_some() { + if let Some(evicted) = evicted { self.dropped.fetch_add(1, Ordering::Relaxed); + evicted.publisher.dropped.fetch_add(1, Ordering::Relaxed); EnqueueOutcome::EnqueuedAfterDrop } else { EnqueueOutcome::Enqueued @@ -1087,6 +1093,36 @@ mod tests { assert_eq!(router.stats().delivered, 2); } + #[test] + fn drop_oldest_charges_the_evicted_publisher_across_namespaces() { + let router = PluginEventRouter::new(["alpha".to_string(), "beta".to_string()]); + let tick = topic("tick"); + let subscription = router + .subscribe( + [selector("alpha", &tick), selector("beta", &tick)], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropOldest), + ) + .expect("subscription"); + + assert_eq!(publish(&router, "alpha", &tick, 1).dropped, 0); + assert_eq!(publish(&router, "beta", &tick, 2).dropped, 1); + + let event = subscription.try_recv().expect("newest event"); + assert_eq!(&*event.plugin_id, "beta"); + assert_eq!(subscription.stats().dropped, 1); + assert_eq!( + router + .publisher_stats("alpha") + .expect("alpha stats") + .dropped, + 1 + ); + let beta = router.publisher_stats("beta").expect("beta stats"); + assert_eq!(beta.dropped, 0); + assert_eq!(beta.delivered, 1); + assert_eq!(router.stats().dropped, 1); + } + #[test] fn backpressure_is_isolated_per_endpoint() { let router = PluginEventRouter::new(["metrics".to_string()]); From f1a8138b772f44d9caa1123ba55e3bec44e2523f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 21 Jul 2026 23:11:41 -0300 Subject: [PATCH 38/46] fix(plugins): publish APIs at construction commit --- agent_docs/plugin_architecture.md | 3 +- src/client/builder.rs | 10 ++ src/plugins/mod.rs | 211 ++++++++++++++++++++++++++++-- 3 files changed, 211 insertions(+), 13 deletions(-) diff --git a/agent_docs/plugin_architecture.md b/agent_docs/plugin_architecture.md index d09832b34..b8f36d081 100644 --- a/agent_docs/plugin_architecture.md +++ b/agent_docs/plugin_architecture.md @@ -54,7 +54,8 @@ Construction follows one publication boundary: Plugin tasks requested during installation remain parked until activation. A client leaked through an installation `Weak` cannot run before the -construction gate opens. +construction gate opens. Plugin APIs, manifests, diagnostics, and custom-event +routing also remain hidden until that final publication succeeds. Installation is transactional. Duplicate IDs or marker types, malformed versions, missing/duplicate dependencies, and cycles fail before client diff --git a/src/client/builder.rs b/src/client/builder.rs index 5b4d7ed62..f560d2e48 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -516,6 +516,16 @@ impl ClientBuilder { "client shutdown raced lifecycle activation" ))); } + #[cfg(feature = "plugins")] + if let Some(plugin_host) = &client.plugin_host + && !plugin_host.publish_apis() + { + client.signal_shutdown_sync(); + client.shutdown_lifecycle().await; + return Err(ClientBuilderError::PluginInstall(anyhow::anyhow!( + "plugin APIs could not be published" + ))); + } #[cfg(feature = "client-lifecycle")] construction.disarm(); Ok(build) diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 2adac129e..ecd40ad32 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -1552,12 +1552,17 @@ struct PluginContextParts { diagnostics: Arc, } +struct InstalledPlugins { + plugins: Vec, + staged_apis: Mutex>>, +} + pub(crate) struct PluginHost { ordered: Vec, manifests: Vec, diagnostics: Vec>, upstream: Option>, - installed: OnceLock>, + installed: OnceLock, apis: OnceLock>, runtime: OnceLock>, event_router: Option, @@ -1616,6 +1621,17 @@ impl PluginHost { downcast_api::(self.apis.get()?.get(&TypeId::of::

())?) } + fn is_published(&self) -> bool { + self.apis.get().is_some() + } + + fn installed_plugins(&self) -> &[InstalledPlugin] { + self.installed + .get() + .map(|installed| installed.plugins.as_slice()) + .unwrap_or_default() + } + pub(crate) fn manifests(&self) -> &[PluginManifest] { &self.manifests } @@ -1860,12 +1876,13 @@ impl PluginHost { } self.abort_install_if_terminal(&mut rollback).await?; - self.apis - .set(staging.snapshot()) - .map_err(|_| anyhow::anyhow!("plugin APIs were published more than once"))?; let installed = rollback.take_installed(); + let installed = InstalledPlugins { + plugins: installed, + staged_apis: Mutex::new(Some(staging.snapshot())), + }; if let Err(installed) = self.installed.set(installed) { - rollback.restore_installed(installed); + rollback.restore_installed(installed.plugins); anyhow::bail!("plugins were installed more than once"); } rollback.disarm(); @@ -1921,7 +1938,7 @@ impl PluginHost { self.close_installed_resources(); return false; } - for plugin in self.installed.get().into_iter().flatten() { + for plugin in self.installed_plugins() { plugin.resources.activate(); } if self.terminal.load(Ordering::Acquire) { @@ -1932,8 +1949,28 @@ impl PluginHost { } } + pub(crate) fn publish_apis(&self) -> bool { + let Some(installed) = self.installed.get() else { + return false; + }; + let mut staged = installed + .staged_apis + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some(apis) = staged.take() else { + return self.apis.get().is_some(); + }; + match self.apis.set(apis) { + Ok(()) => true, + Err(apis) => { + *staged = Some(apis); + false + } + } + } + fn close_installed_resources(&self) { - for plugin in self.installed.get().into_iter().flatten().rev() { + for plugin in self.installed_plugins().iter().rev() { close_plugin_resources(&plugin.manifest.id, &plugin.resources); } } @@ -1962,7 +1999,7 @@ impl ClientLifecycle for PluginHost { { failures.push(format!("upstream: {error:#}")); } - for plugin in self.installed.get().into_iter().flatten() { + for plugin in self.installed_plugins() { let task_tracker = plugin .manifest .capabilities @@ -1996,7 +2033,7 @@ impl ClientLifecycle for PluginHost { fn on_closed(&self, scope: ConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { Box::pin(async move { let mut failures = Vec::new(); - for plugin in self.installed.get().into_iter().flatten().rev() { + for plugin in self.installed_plugins().iter().rev() { let task_tracker = plugin .manifest .capabilities @@ -2053,7 +2090,7 @@ impl ClientLifecycle for PluginHost { Box::pin(async move { let mut failures = Vec::new(); self.signal_shutdown(); - for plugin in self.installed.get().into_iter().flatten().rev() { + for plugin in self.installed_plugins().iter().rev() { let task_result = self .wait_for_tasks(plugin.resources.task_completion_signals()) .await; @@ -2268,13 +2305,17 @@ impl Client { pub fn plugin_manifests(&self) -> &[PluginManifest] { self.plugin_host .as_ref() + .filter(|host| host.is_published()) .map(|host| host.manifests()) .unwrap_or_default() } /// Snapshot lifecycle, task, subscription, and custom-event health for installed plugins. pub fn plugin_stats(&self) -> Option { - self.plugin_host.as_ref().map(|host| host.stats()) + self.plugin_host + .as_ref() + .filter(|host| host.is_published()) + .map(|host| host.stats()) } /// Subscribe to custom events emitted by installed plugins. @@ -2283,12 +2324,15 @@ impl Client { pub fn plugin_event_router(&self) -> Option { self.plugin_host .as_ref() + .filter(|host| host.is_published()) .and_then(|host| host.event_router.clone()) } } #[cfg(test)] mod tests { + use std::pin::Pin; + use std::sync::Barrier; use std::sync::atomic::AtomicBool; use std::time::Duration; @@ -2332,6 +2376,41 @@ mod tests { client: async_channel::Sender>, } + struct BlockingFirstSpawnRuntime { + blocked: AtomicBool, + entered: async_channel::Sender<()>, + release: Arc, + } + + #[async_trait::async_trait] + impl Runtime for BlockingFirstSpawnRuntime { + fn spawn( + &self, + future: Pin + Send + 'static>>, + ) -> wacore::runtime::AbortHandle { + if !self.blocked.swap(true, Ordering::AcqRel) { + self.entered.try_send(()).expect("first spawn observer"); + self.release.wait(); + } + TokioRuntime.spawn(future) + } + + fn sleep(&self, duration: Duration) -> Pin + Send>> { + TokioRuntime.sleep(duration) + } + + fn spawn_blocking( + &self, + f: Box, + ) -> Pin + Send>> { + TokioRuntime.spawn_blocking(f) + } + + fn yield_now(&self) -> Option + Send>>> { + TokioRuntime.yield_now() + } + } + impl ClientLifecycle for CaptureInstallClient { fn install(&self, client: Weak) -> BoxFuture<'_, anyhow::Result<()>> { let sender = self.client.clone(); @@ -2342,6 +2421,24 @@ mod tests { } } + struct PublicationProbePlugin; + + impl ClientPlugin for PublicationProbePlugin { + type Api = String; + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("publication-probe", "0.1.0") + .with_capability(PluginCapability::PluginEvents) + } + + fn install( + &self, + _context: PluginContext, + ) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async { Ok(Arc::new("published-api".to_string())) }) + } + } + struct TerminalBlockingInstallPlugin { started: async_channel::Sender, install_dropped: Arc, @@ -2436,6 +2533,95 @@ mod tests { ); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn plugin_surfaces_publish_only_after_final_activation() { + let (client_tx, client_rx) = async_channel::bounded(1); + let (entered_tx, entered_rx) = async_channel::bounded(1); + let release = Arc::new(Barrier::new(2)); + let builder = complete_builder() + .await + .with_runtime(BlockingFirstSpawnRuntime { + blocked: AtomicBool::new(false), + entered: entered_tx, + release: Arc::clone(&release), + }) + .with_lifecycle(CaptureInstallClient { client: client_tx }) + .with_plugin(PublicationProbePlugin); + + let build = tokio::spawn(async move { builder.build().await }); + let leaked_client = client_rx + .recv() + .await + .expect("captured install client") + .upgrade() + .expect("client under construction"); + entered_rx.recv().await.expect("client service startup"); + + assert!(leaked_client.plugin::().is_none()); + assert!(leaked_client.plugin_manifests().is_empty()); + assert!(leaked_client.plugin_stats().is_none()); + assert!(leaked_client.plugin_event_router().is_none()); + + release.wait(); + let client = build + .await + .expect("builder task") + .expect("successful build") + .into_client(); + assert_eq!( + client + .plugin::() + .as_deref() + .map(String::as_str), + Some("published-api") + ); + assert_eq!(client.plugin_manifests().len(), 1); + assert!(client.plugin_stats().is_some()); + assert!(client.plugin_event_router().is_some()); + client.disconnect().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn rejected_construction_never_publishes_staged_plugin_surfaces() { + let (client_tx, client_rx) = async_channel::bounded(1); + let (entered_tx, entered_rx) = async_channel::bounded(1); + let release = Arc::new(Barrier::new(2)); + let builder = complete_builder() + .await + .with_runtime(BlockingFirstSpawnRuntime { + blocked: AtomicBool::new(false), + entered: entered_tx, + release: Arc::clone(&release), + }) + .with_lifecycle(CaptureInstallClient { client: client_tx }) + .with_plugin(PublicationProbePlugin); + + let build = tokio::spawn(async move { builder.build().await }); + let leaked_client = client_rx + .recv() + .await + .expect("captured install client") + .upgrade() + .expect("client under construction"); + entered_rx.recv().await.expect("client service startup"); + leaked_client.signal_shutdown_sync(); + + assert!(leaked_client.plugin::().is_none()); + assert!(leaked_client.plugin_manifests().is_empty()); + assert!(leaked_client.plugin_stats().is_none()); + assert!(leaked_client.plugin_event_router().is_none()); + + release.wait(); + assert!(matches!( + build.await.expect("builder task"), + Err(ClientBuilderError::PluginInstall(_)) + )); + assert!(leaked_client.plugin::().is_none()); + assert!(leaked_client.plugin_manifests().is_empty()); + assert!(leaked_client.plugin_stats().is_none()); + assert!(leaked_client.plugin_event_router().is_none()); + } + #[tokio::test] async fn shutdown_cancels_an_inflight_plugin_install_and_closes_its_resources() { let (client_tx, client_rx) = async_channel::bounded(1); @@ -3405,7 +3591,8 @@ mod tests { let client = build.into_client(); wait_for_flag(&install_started).await; let host = client.plugin_host.as_ref().expect("plugin host").clone(); - let resources = Arc::clone(&host.installed.get().expect("installed plugins")[0].resources); + let resources = + Arc::clone(&host.installed.get().expect("installed plugins").plugins[0].resources); let stats = client.plugin_stats().expect("plugin stats"); assert_eq!(stats.health, PluginHealth::Healthy); assert_eq!(stats.plugins[0].state, PluginState::Active); From 719011e16bb819c67b24ff78754375c9994ba493 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 22 Jul 2026 10:02:31 -0300 Subject: [PATCH 39/46] fix(plugins): close lifecycle review races --- src/client.rs | 1 + src/client/builder.rs | 50 +++----- src/client/extension_lifecycle.rs | 196 +++++++++++++++++++++++++----- src/client/lifecycle.rs | 4 + src/client/node_io.rs | 2 + src/plugins/mod.rs | 132 ++++++++++++++------ 6 files changed, 284 insertions(+), 101 deletions(-) diff --git a/src/client.rs b/src/client.rs index ed7898597..5829835a1 100644 --- a/src/client.rs +++ b/src/client.rs @@ -704,6 +704,7 @@ pub struct Client { pub(crate) media_conn: Arc>>, pub(crate) is_logged_in: Arc, + #[cfg(feature = "client-lifecycle")] pub(crate) login_transition: std::sync::Mutex<()>, pub(crate) is_connecting: Arc, pub(crate) is_running: Arc, diff --git a/src/client/builder.rs b/src/client/builder.rs index f560d2e48..2c97c7e76 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -490,41 +490,29 @@ impl ClientBuilder { ); let _ = build.client.saver_handle.set(saver_handle); } - #[cfg(feature = "plugins")] - if let Some(plugin_host) = &client.plugin_host - && !plugin_host.activate() - { - client.signal_shutdown_sync(); - client.shutdown_lifecycle().await; - return Err(ClientBuilderError::PluginInstall(anyhow::anyhow!( - "client shutdown began before plugin activation" - ))); - } #[cfg(feature = "client-lifecycle")] - if let Some(lifecycle) = &client.lifecycle - && !lifecycle.activate() - { - client.signal_shutdown_sync(); - client.shutdown_lifecycle().await; + if let Some(lifecycle) = &client.lifecycle { #[cfg(feature = "plugins")] - if client.plugin_host.is_some() { - return Err(ClientBuilderError::PluginInstall(anyhow::anyhow!( - "client shutdown raced plugin activation" + let activated = if let Some(plugin_host) = &client.plugin_host { + lifecycle.activate_with(|| plugin_host.commit()) + } else { + lifecycle.activate() + }; + #[cfg(not(feature = "plugins"))] + let activated = lifecycle.activate(); + if !activated { + client.signal_shutdown_sync(); + client.shutdown_lifecycle().await; + #[cfg(feature = "plugins")] + if client.plugin_host.is_some() { + return Err(ClientBuilderError::PluginInstall(anyhow::anyhow!( + "client shutdown raced plugin publication" + ))); + } + return Err(ClientBuilderError::LifecycleInstall(anyhow::anyhow!( + "client shutdown raced lifecycle activation" ))); } - return Err(ClientBuilderError::LifecycleInstall(anyhow::anyhow!( - "client shutdown raced lifecycle activation" - ))); - } - #[cfg(feature = "plugins")] - if let Some(plugin_host) = &client.plugin_host - && !plugin_host.publish_apis() - { - client.signal_shutdown_sync(); - client.shutdown_lifecycle().await; - return Err(ClientBuilderError::PluginInstall(anyhow::anyhow!( - "plugin APIs could not be published" - ))); } #[cfg(feature = "client-lifecycle")] construction.disarm(); diff --git a/src/client/extension_lifecycle.rs b/src/client/extension_lifecycle.rs index 6ff902d0e..59f573431 100644 --- a/src/client/extension_lifecycle.rs +++ b/src/client/extension_lifecycle.rs @@ -169,6 +169,7 @@ pub(super) struct LifecycleRegistration { shutdown_notifier: ShutdownNotifier, callback_timeout: Duration, terminal: AtomicBool, + construction_transition: std::sync::Mutex<()>, construction_state: AtomicU8, construction_notifier: ShutdownNotifier, } @@ -198,31 +199,32 @@ struct CallbackQueue { } impl CallbackQueue { - fn push_with_pressure_policy( - &mut self, - callback: LifecycleCallback, - ) -> Option { - if self.pending.len() < CALLBACK_QUEUE_TARGET_CAPACITY { + fn push_with_pressure_policy(&mut self, callback: LifecycleCallback) -> Vec { + if self.pending.len() < CALLBACK_QUEUE_TARGET_CAPACITY && !self.overflowed { self.pending.push_back(callback); - return None; - } - - self.overflowed = true; - let ready_position = self - .pending - .iter() - .position(|pending| matches!(pending, LifecycleCallback::Ready { .. })); - match (callback, ready_position) { - (callback @ LifecycleCallback::Ready { .. }, None) => Some(callback), - (callback, Some(position)) => { - let dropped = self.pending.remove(position); + return Vec::new(); + } + + self.overflowed |= self.pending.len() >= CALLBACK_QUEUE_TARGET_CAPACITY; + match callback { + callback @ LifecycleCallback::Ready { .. } => { + let mut dropped = Vec::new(); + let mut retained = VecDeque::with_capacity(self.pending.len()); + for pending in self.pending.drain(..) { + if matches!(pending, LifecycleCallback::Ready { .. }) { + dropped.push(pending); + } else { + retained.push_back(pending); + } + } + self.pending = retained; self.pending.push_back(callback); dropped } - (callback, None) => { - // Every closed scope must reach the extension even when a callback stalls. + callback => { + // Closures are lossless, so the target remains soft under backlog. self.pending.push_back(callback); - None + Vec::new() } } } @@ -307,6 +309,7 @@ impl LifecycleRegistration { shutdown_notifier: ShutdownNotifier::new(), callback_timeout, terminal: AtomicBool::new(false), + construction_transition: std::sync::Mutex::new(()), construction_state: AtomicU8::new(CONSTRUCTION_INSTALLING), construction_notifier: ShutdownNotifier::new(), } @@ -359,10 +362,25 @@ impl LifecycleRegistration { } pub(super) fn activate(&self) -> bool { + self.activate_with(|| true) + } + + pub(super) fn activate_with(&self, commit: impl FnOnce() -> bool) -> bool { + let _transition = self + .construction_transition + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); if self.terminal.load(Ordering::Acquire) { self.reject_construction(); return false; } + if self.construction_state.load(Ordering::Acquire) == CONSTRUCTION_ACTIVE { + return true; + } + if !commit() { + self.reject_construction(); + return false; + } match self.construction_state.compare_exchange( CONSTRUCTION_INSTALLING, CONSTRUCTION_ACTIVE, @@ -373,7 +391,7 @@ impl LifecycleRegistration { Err(CONSTRUCTION_ACTIVE) => {} Err(_) => return false, } - !self.terminal.load(Ordering::Acquire) + true } pub(super) async fn wait_until_active(&self) -> bool { @@ -558,13 +576,20 @@ impl LifecycleRegistration { } pub(super) fn signal_shutdown_sync(&self) { - self.reject_construction(); - let first_signal = if ready_publication_active(self) { - self.mark_terminal_and_cancel_scopes() + let first_signal = { + let _transition = self + .construction_transition + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + self.reject_construction(); + !self.terminal.swap(true, Ordering::AcqRel) + }; + if ready_publication_active(self) { + self.cancel_all_scopes() } else { let _publication = self.ready_publication(); - self.mark_terminal_and_cancel_scopes() - }; + self.cancel_all_scopes() + } if first_signal && std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { self.handler.signal_shutdown(); @@ -579,7 +604,7 @@ impl LifecycleRegistration { let (should_spawn, dropped) = { let mut queue = self.callback_queue(); if queue.shutdown_requested || self.terminal.load(Ordering::Acquire) { - (false, Some(callback)) + (false, vec![callback]) } else { let dropped = queue.push_with_pressure_policy(callback); let should_spawn = !queue.drain_scheduled; @@ -587,7 +612,7 @@ impl LifecycleRegistration { (should_spawn, dropped) } }; - warn_dropped_callbacks(dropped.into_iter().collect()); + warn_dropped_callbacks(dropped); self.spawn_callback_driver(should_spawn); } @@ -753,8 +778,7 @@ impl LifecycleRegistration { } } - fn mark_terminal_and_cancel_scopes(&self) -> bool { - let first_signal = !self.terminal.swap(true, Ordering::AcqRel); + fn cancel_all_scopes(&self) { let scopes = self.scopes(); if let Some(scope) = &scopes.active { scope.cancel(); @@ -762,7 +786,6 @@ impl LifecycleRegistration { for scope in &scopes.retired { scope.cancel(); } - first_signal } fn scope_for(&self, generation: u64) -> Option { @@ -1253,6 +1276,58 @@ mod tests { assert!(!published.load(Ordering::Acquire)); } + #[test] + fn terminal_signal_waits_for_construction_commit() { + let registration = Arc::new(LifecycleRegistration::new( + Arc::new(RecordingLifecycle::default()), + Arc::new(TokioRuntime), + )); + let (commit_started_tx, commit_started_rx) = std::sync::mpsc::sync_channel(1); + let (release_commit_tx, release_commit_rx) = std::sync::mpsc::sync_channel(1); + let (activation_tx, activation_rx) = std::sync::mpsc::sync_channel(1); + let activation_registration = registration.clone(); + let activation = std::thread::spawn(move || { + let activated = activation_registration.activate_with(|| { + commit_started_tx.send(()).expect("publish commit start"); + release_commit_rx.recv().expect("release publish commit"); + true + }); + activation_tx.send(activated).expect("activation result"); + }); + commit_started_rx + .recv_timeout(Duration::from_secs(2)) + .expect("construction commit started"); + + let (shutdown_tx, shutdown_rx) = std::sync::mpsc::sync_channel(1); + let shutdown_registration = registration.clone(); + let shutdown = std::thread::spawn(move || { + shutdown_registration.signal_shutdown_sync(); + shutdown_tx.send(()).expect("shutdown result"); + }); + assert!( + shutdown_rx + .recv_timeout(Duration::from_millis(100)) + .is_err() + ); + + release_commit_tx.send(()).expect("finish publish commit"); + assert!( + activation_rx + .recv_timeout(Duration::from_secs(2)) + .expect("construction activated") + ); + shutdown_rx + .recv_timeout(Duration::from_secs(2)) + .expect("terminal signal completed"); + activation.join().expect("activation thread"); + shutdown.join().expect("shutdown thread"); + assert_eq!( + registration.construction_state.load(Ordering::Acquire), + CONSTRUCTION_ACTIVE + ); + assert!(registration.terminal.load(Ordering::Acquire)); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn terminal_cancellation_waits_for_ready_publication() { let registration = Arc::new(LifecycleRegistration::new( @@ -1731,6 +1806,63 @@ mod tests { ); } + #[test] + fn callback_queue_retains_latest_ready_with_lossless_close_backlog() { + let mut queue = CallbackQueue::default(); + for generation in 1..=CALLBACK_QUEUE_TARGET_CAPACITY as u64 { + let scope = ConnectionScope::new(generation); + scope.close(); + assert!( + queue + .push_with_pressure_policy(LifecycleCallback::Closed(scope)) + .is_empty() + ); + } + + let (first_done, _first_completion) = async_channel::bounded(1); + assert!( + queue + .push_with_pressure_policy(LifecycleCallback::Ready { + scope: ConnectionScope::new(100), + done: first_done, + }) + .is_empty() + ); + assert_eq!(queue.pending.len(), CALLBACK_QUEUE_TARGET_CAPACITY + 1); + + let (latest_done, _latest_completion) = async_channel::bounded(1); + let mut dropped = queue.push_with_pressure_policy(LifecycleCallback::Ready { + scope: ConnectionScope::new(101), + done: latest_done, + }); + assert_eq!(dropped.len(), 1); + let dropped = dropped.pop().expect("older ready callback is replaceable"); + assert!(matches!( + dropped, + LifecycleCallback::Ready { scope, .. } if scope.generation() == 100 + )); + + let extra_closed = ConnectionScope::new(102); + extra_closed.close(); + assert!( + queue + .push_with_pressure_policy(LifecycleCallback::Closed(extra_closed)) + .is_empty() + ); + assert_eq!(queue.pending.len(), CALLBACK_QUEUE_TARGET_CAPACITY + 2); + assert_eq!( + queue + .pending + .iter() + .filter(|callback| matches!(callback, LifecycleCallback::Ready { .. })) + .count(), + 1 + ); + assert!(queue.pending.iter().any(|callback| { + matches!(callback, LifecycleCallback::Ready { scope, .. } if scope.generation() == 101) + })); + } + #[tokio::test] async fn callback_queue_preserves_every_scope_closure_before_shutdown() { let (ready_started_tx, ready_started_rx) = async_channel::bounded(1); @@ -1773,7 +1905,7 @@ mod tests { } assert_eq!( registration.callback_queue().pending.len(), - usize::try_from(closed_callbacks).expect("closure count fits usize") + usize::try_from(closed_callbacks + 1).expect("callback count fits usize") ); let shutdown_registration = registration.clone(); diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index a2d9beee8..78ab08d77 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -114,6 +114,7 @@ impl Client { } } + #[cfg(feature = "client-lifecycle")] let _login_transition = self .login_transition .lock() @@ -228,6 +229,7 @@ impl Client { persistence_manager: persistence_manager.clone(), media_conn: Arc::new(RwLock::new(None)), is_logged_in: Arc::new(AtomicBool::new(false)), + #[cfg(feature = "client-lifecycle")] login_transition: std::sync::Mutex::new(()), is_connecting: Arc::new(AtomicBool::new(false)), is_running: Arc::new(AtomicBool::new(false)), @@ -940,6 +942,7 @@ impl Client { } async fn cleanup_connection_state_inner(&self) { + #[cfg(feature = "client-lifecycle")] let login_transition = self .login_transition .lock() @@ -984,6 +987,7 @@ impl Client { // outgoing stanzas, which are transport-scoped. self.clear_sent_node_waiters(); self.is_logged_in.store(false, Ordering::Relaxed); + #[cfg(feature = "client-lifecycle")] drop(login_transition); self.is_ready.store(false, Ordering::Relaxed); // Publish the disconnected state BEFORE draining VoIP calls (it used to be cleared only after diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 49bf1b4ab..9d8765c14 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -679,6 +679,7 @@ impl Client { tracing::instrument(name = "wa.conn.success", level = "debug", skip_all) )] pub(crate) async fn handle_success(self: &Arc, node: &wacore_binary::NodeRef<'_>) { + #[cfg(feature = "client-lifecycle")] let login_transition = self .login_transition .lock() @@ -713,6 +714,7 @@ impl Client { return; } } + #[cfg(feature = "client-lifecycle")] drop(login_transition); info!( diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index ecd40ad32..58049c8dd 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -278,6 +278,8 @@ pub struct PluginStats { pub struct PluginHostStats { pub terminal: bool, pub health: PluginHealth, + pub upstream_callback_failures: u64, + pub upstream_callback_timeouts: u64, pub plugins: Vec, pub event_router: Option, } @@ -407,12 +409,23 @@ impl PluginResources { }) } + #[cfg(test)] fn activate(&self) { + self.prepare_activation(); + self.publish_activation(); + } + + fn prepare_activation(&self) { if self.closed.load(Ordering::Acquire) { return; } self.active.store(true, Ordering::Release); - self.activation.notify(); + } + + fn publish_activation(&self) { + if self.active.load(Ordering::Acquire) && !self.closed.load(Ordering::Acquire) { + self.activation.notify(); + } } fn ensure_active(&self) -> Result<(), PluginResourceError> { @@ -809,6 +822,12 @@ struct PluginCoreEventHandler { impl EventHandler for PluginCoreEventHandler { fn handle_event(&self, event: Arc) { + let Some(resources) = self.resources.upgrade() else { + return; + }; + if resources.ensure_active().is_err() { + return; + } let Some(inner) = &self.inner else { return; }; @@ -1570,6 +1589,8 @@ pub(crate) struct PluginHost { terminal: AtomicBool, terminal_notifier: ShutdownNotifier, installing_resources: Mutex>>, + upstream_callback_failures: AtomicU64, + upstream_callback_timeouts: AtomicU64, } impl PluginHost { @@ -1614,6 +1635,8 @@ impl PluginHost { terminal: AtomicBool::new(false), terminal_notifier: ShutdownNotifier::new(), installing_resources: Mutex::new(Vec::new()), + upstream_callback_failures: AtomicU64::new(0), + upstream_callback_timeouts: AtomicU64::new(0), }) } @@ -1638,6 +1661,8 @@ impl PluginHost { pub(crate) fn stats(&self) -> PluginHostStats { let terminal = self.terminal.load(Ordering::Acquire); + let upstream_callback_failures = self.upstream_callback_failures.load(Ordering::Relaxed); + let upstream_callback_timeouts = self.upstream_callback_timeouts.load(Ordering::Relaxed); let plugins = self .manifests .iter() @@ -1650,9 +1675,11 @@ impl PluginHost { diagnostics.snapshot(&manifest.id, terminal, events) }) .collect::>(); - let health = if plugins - .iter() - .any(|plugin| plugin.health == PluginHealth::Degraded) + let health = if upstream_callback_failures > 0 + || upstream_callback_timeouts > 0 + || plugins + .iter() + .any(|plugin| plugin.health == PluginHealth::Degraded) { PluginHealth::Degraded } else { @@ -1661,6 +1688,8 @@ impl PluginHost { PluginHostStats { terminal, health, + upstream_callback_failures, + upstream_callback_timeouts, plugins, event_router: self.event_router.as_ref().map(PluginEventRouter::stats), } @@ -1933,23 +1962,11 @@ impl PluginHost { } } - pub(crate) fn activate(&self) -> bool { + pub(crate) fn commit(&self) -> bool { if self.terminal.load(Ordering::Acquire) { self.close_installed_resources(); return false; } - for plugin in self.installed_plugins() { - plugin.resources.activate(); - } - if self.terminal.load(Ordering::Acquire) { - self.close_installed_resources(); - false - } else { - true - } - } - - pub(crate) fn publish_apis(&self) -> bool { let Some(installed) = self.installed.get() else { return false; }; @@ -1957,16 +1974,22 @@ impl PluginHost { .staged_apis .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let Some(apis) = staged.take() else { - return self.apis.get().is_some(); - }; - match self.apis.set(apis) { - Ok(()) => true, - Err(apis) => { - *staged = Some(apis); - false - } + if let Some(apis) = staged.take() + && let Err(apis) = self.apis.set(apis) + { + *staged = Some(apis); + return false; + } + if self.apis.get().is_none() { + return false; + } + for plugin in &installed.plugins { + plugin.resources.prepare_activation(); + } + for plugin in &installed.plugins { + plugin.resources.publish_activation(); } + true } fn close_installed_resources(&self) { @@ -1984,6 +2007,26 @@ impl PluginHost { })?; bounded_plugin_callback(&**runtime, self.callback_timeout, make_future).await } + + fn record_upstream_callback(&self, result: &Result<(), PluginCallbackError>) { + match result { + Ok(()) => {} + Err(PluginCallbackError::Timeout { .. }) => { + self.upstream_callback_timeouts + .fetch_add(1, Ordering::Relaxed); + } + Err(PluginCallbackError::TimeoutCancellationPanic { .. }) => { + self.upstream_callback_timeouts + .fetch_add(1, Ordering::Relaxed); + self.upstream_callback_failures + .fetch_add(1, Ordering::Relaxed); + } + Err(PluginCallbackError::Callback(_)) => { + self.upstream_callback_failures + .fetch_add(1, Ordering::Relaxed); + } + } + } } impl ClientLifecycle for PluginHost { @@ -1994,10 +2037,12 @@ impl ClientLifecycle for PluginHost { fn on_ready(&self, scope: ConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { Box::pin(async move { let mut failures = Vec::new(); - if let Some(upstream) = &self.upstream - && let Err(error) = self.run_callback(|| upstream.on_ready(scope.clone())).await - { - failures.push(format!("upstream: {error:#}")); + if let Some(upstream) = &self.upstream { + let result = self.run_callback(|| upstream.on_ready(scope.clone())).await; + self.record_upstream_callback(&result); + if let Err(error) = result { + failures.push(format!("upstream: {error:#}")); + } } for plugin in self.installed_plugins() { let task_tracker = plugin @@ -2062,10 +2107,12 @@ impl ClientLifecycle for PluginHost { failures.push(format!("{}: {error:#}", plugin.manifest.id)); } } - if let Some(upstream) = &self.upstream - && let Err(error) = self.run_callback(|| upstream.on_closed(scope)).await - { - failures.push(format!("upstream: {error:#}")); + if let Some(upstream) = &self.upstream { + let result = self.run_callback(|| upstream.on_closed(scope)).await; + self.record_upstream_callback(&result); + if let Err(error) = result { + failures.push(format!("upstream: {error:#}")); + } } finish_callbacks("closed", failures) }) @@ -2082,6 +2129,8 @@ impl ClientLifecycle for PluginHost { if let Some(upstream) = &self.upstream && std::panic::catch_unwind(AssertUnwindSafe(|| upstream.signal_shutdown())).is_err() { + self.upstream_callback_failures + .fetch_add(1, Ordering::Relaxed); log::warn!("Upstream lifecycle synchronous shutdown signal panicked"); } } @@ -2105,10 +2154,12 @@ impl ClientLifecycle for PluginHost { failures.push(format!("{}: {error:#}", plugin.manifest.id)); } } - if let Some(upstream) = &self.upstream - && let Err(error) = self.run_callback(|| upstream.shutdown()).await - { - failures.push(format!("upstream: {error:#}")); + if let Some(upstream) = &self.upstream { + let result = self.run_callback(|| upstream.shutdown()).await; + self.record_upstream_callback(&result); + if let Err(error) = result { + failures.push(format!("upstream: {error:#}")); + } } finish_callbacks("shutdown", failures) }) @@ -3938,6 +3989,11 @@ mod tests { *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), vec!["ready:ready-probe"] ); + let stats = client.plugin_stats().expect("plugin host stats"); + assert_eq!(stats.health, PluginHealth::Degraded); + assert_eq!(stats.upstream_callback_failures, 1); + assert_eq!(stats.upstream_callback_timeouts, 0); + assert_eq!(stats.plugins[0].health, PluginHealth::Healthy); client.disconnect().await; } From 29481d9e7a9cf3dd1639dd7399b65285b8bad882 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:37:44 -0300 Subject: [PATCH 40/46] feat(plugins): add manifest-keyed untyped adapters --- agent_docs/plugin_architecture.md | 12 +- src/bot.rs | 19 ++- src/client/builder.rs | 21 +++- src/lib.rs | 3 +- src/plugins/mod.rs | 187 +++++++++++++++++++++++++++--- 5 files changed, 224 insertions(+), 18 deletions(-) diff --git a/agent_docs/plugin_architecture.md b/agent_docs/plugin_architecture.md index b8f36d081..913661a16 100644 --- a/agent_docs/plugin_architecture.md +++ b/agent_docs/plugin_architecture.md @@ -125,6 +125,13 @@ plugins may therefore expose the same API type without colliding. is selected at runtime by the builder. Encoding that set in `Client` generics would make the client type viral and substantially increase monomorphization. +An adapter that represents runtime-defined plugins implements +`UntypedClientPlugin` and registers each instance with +`with_untyped_plugin(...)`. Those instances are keyed only by manifest ID, may +share one concrete Rust adapter type, and do not appear in +`Client::plugin::

()`. Native plugins keep the typed path above; a future +bridge multiplexes its language-specific handles behind the untyped adapter. + During installation, `PluginContext::plugin::

()` exposes only directly declared dependencies. The context keeps a weak dependency view so an API that retains its context cannot create a registry ownership cycle. APIs should keep @@ -251,7 +258,10 @@ phone numbers, or message bodies. See `observability.md` for accounting rules. ## Future foreign-language adapter seam A future bridge should be a Rust adapter at the host boundary, not a second -client lifecycle. It can map a foreign endpoint onto the existing semantics: +client lifecycle. Each runtime-defined instance can use +`UntypedClientPlugin`, so one adapter type can host multiple manifest IDs +without colliding in the native `TypeId` API registry. It can map a foreign +endpoint onto the existing semantics: - build-time registration and stable install-scoped handles across reconnects; - explicit capability grants checked for every foreign command; diff --git a/src/bot.rs b/src/bot.rs index c772d40a4..b7d6b8628 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -2,7 +2,7 @@ use crate::cache_config::CacheConfig; use crate::client::{Client, ClientBuilderError}; use crate::pair_code::PairCodeOptions; #[cfg(feature = "plugins")] -use crate::plugins::{ClientPlugin, PluginRegistration}; +use crate::plugins::{ClientPlugin, PluginRegistration, UntypedClientPlugin}; use crate::store::commands::DeviceCommand; use crate::store::error::StoreError; use crate::store::persistence_manager::PersistenceManager; @@ -829,6 +829,23 @@ impl BotBuilder { self } + /// Register a manifest-ID-keyed plugin that exposes no Rust typed API. + #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] + pub fn with_untyped_plugin(mut self, plugin: P) -> Self { + self.plugins.push(PluginRegistration::new_untyped(plugin)); + self + } + + /// Register an already-shared manifest-ID-keyed plugin. + #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] + pub fn with_untyped_plugin_arc(mut self, plugin: Arc

) -> Self { + self.plugins + .push(PluginRegistration::new_untyped_arc(plugin)); + self + } + // ── Event handler registration (additive; order of registration is kept, // but handlers run on their own tasks, so cross-event ordering is not // guaranteed) ────────────────────────────────────────────────────── diff --git a/src/client/builder.rs b/src/client/builder.rs index b085069ba..81ec9982f 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -10,7 +10,9 @@ use super::{ClientLifecycle, LifecycleRegistration}; use crate::cache_config::CacheConfig; use crate::http::HttpClient; #[cfg(feature = "plugins")] -use crate::plugins::{ClientPlugin, PluginHost, PluginPlan, PluginPlanError, PluginRegistration}; +use crate::plugins::{ + ClientPlugin, PluginHost, PluginPlan, PluginPlanError, PluginRegistration, UntypedClientPlugin, +}; use crate::store::error::StoreError; use crate::store::persistence_manager::PersistenceManager; use crate::sync_task::MajorSyncTask; @@ -320,6 +322,23 @@ impl ClientBuilder { self } + /// Register a manifest-ID-keyed plugin that exposes no Rust typed API. + #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] + pub fn with_untyped_plugin(mut self, plugin: P) -> Self { + self.plugins.push(PluginRegistration::new_untyped(plugin)); + self + } + + /// Register an already-shared manifest-ID-keyed plugin. + #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] + pub fn with_untyped_plugin_arc(mut self, plugin: Arc

) -> Self { + self.plugins + .push(PluginRegistration::new_untyped_arc(plugin)); + self + } + #[cfg(feature = "plugins")] pub(crate) fn with_plugin_registrations( mut self, diff --git a/src/lib.rs b/src/lib.rs index 3f2c5de13..cf213d9cb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -127,6 +127,7 @@ pub use plugins::{ PluginEventTryReceiveError, PluginEvents, PluginFuture, PluginHealth, PluginHostStats, PluginIq, PluginIqError, PluginManifest, PluginMessaging, PluginMessagingError, PluginPlanError, PluginResourceError, PluginState, PluginStats, PluginTasks, + UntypedClientPlugin, }; pub mod request; pub(crate) mod signal_flush; @@ -207,7 +208,7 @@ pub mod prelude { ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, PluginEventEndpointConfig, PluginEventOverflow, PluginEventPayloadEncoding, PluginEventRouter, PluginEventSelector, PluginEventSubscription, PluginEventTopic, - PluginEvents, PluginFuture, PluginManifest, + PluginEvents, PluginFuture, PluginManifest, UntypedClientPlugin, }; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 58049c8dd..8abc68677 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -166,6 +166,29 @@ pub trait ClientPlugin: MaybeSendSync + 'static { } } +/// A trusted plugin instance identified only by its manifest ID. +/// +/// Unlike [`ClientPlugin`], this trait publishes no Rust type-indexed API, so +/// multiple instances of the same adapter type may be registered. It is the +/// intended host seam for runtime-defined or foreign-language plugins. +pub trait UntypedClientPlugin: MaybeSendSync + 'static { + fn manifest(&self) -> PluginManifest; + + fn install(&self, context: PluginContext) -> PluginFuture<'_, anyhow::Result<()>>; + + fn on_ready(&self, _scope: PluginConnectionScope) -> PluginFuture<'_, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } + + fn on_closed(&self, _scope: PluginConnectionScope) -> PluginFuture<'_, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } + + fn shutdown(&self) -> PluginFuture<'_, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } +} + /// Manifest validation or dependency-ordering failure. #[derive(Debug, Error)] #[non_exhaustive] @@ -1248,10 +1271,10 @@ fn downcast_api(api: &ErasedApi) -> Option> { } trait ErasedClientPlugin: MaybeSendSync { - fn marker_type_id(&self) -> TypeId; + fn marker_type_id(&self) -> Option; fn marker_type_name(&self) -> &'static str; fn manifest(&self) -> PluginManifest; - fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>; + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>>; fn on_ready(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>>; fn on_closed(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>>; fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>>; @@ -1260,8 +1283,8 @@ trait ErasedClientPlugin: MaybeSendSync { struct PluginAdapter

(Arc

); impl ErasedClientPlugin for PluginAdapter

{ - fn marker_type_id(&self) -> TypeId { - TypeId::of::

() + fn marker_type_id(&self) -> Option { + Some(TypeId::of::

()) } fn marker_type_name(&self) -> &'static str { @@ -1272,10 +1295,45 @@ impl ErasedClientPlugin for PluginAdapter

{ self.0.manifest() } - fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result> { + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { Box::pin(async move { let api = self.0.install(context).await?; - Ok(Arc::new(TypedApi(api)) as ErasedApi) + Ok(Some(Arc::new(TypedApi(api)) as ErasedApi)) + }) + } + + fn on_ready(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + self.0.on_ready(scope) + } + + fn on_closed(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + self.0.on_closed(scope) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + self.0.shutdown() + } +} + +struct UntypedPluginAdapter

(Arc

); + +impl ErasedClientPlugin for UntypedPluginAdapter

{ + fn marker_type_id(&self) -> Option { + None + } + + fn marker_type_name(&self) -> &'static str { + std::any::type_name::

() + } + + fn manifest(&self) -> PluginManifest { + self.0.manifest() + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + Box::pin(async move { + self.0.install(context).await?; + Ok(None) }) } @@ -1306,6 +1364,16 @@ impl PluginRegistration { plugin: Arc::new(PluginAdapter(plugin)), } } + + pub(crate) fn new_untyped(plugin: P) -> Self { + Self::new_untyped_arc(Arc::new(plugin)) + } + + pub(crate) fn new_untyped_arc(plugin: Arc

) -> Self { + Self { + plugin: Arc::new(UntypedPluginAdapter(plugin)), + } + } } struct PlannedPlugin { @@ -1332,8 +1400,9 @@ impl PluginPlan { for registration in registrations { let plugin = registration.plugin; - let marker = plugin.marker_type_id(); - if !marker_types.insert(marker) { + if let Some(marker) = plugin.marker_type_id() + && !marker_types.insert(marker) + { return Err(PluginPlanError::DuplicateType { plugin_type: plugin.marker_type_name(), }); @@ -1376,8 +1445,9 @@ impl PluginPlan { }; indegree[plugin_index] += 1; dependents[dependency_index].push(plugin_index); - dependency_markers[plugin_index] - .push(plugins[dependency_index].plugin.marker_type_id()); + if let Some(marker) = plugins[dependency_index].plugin.marker_type_id() { + dependency_markers[plugin_index].push(marker); + } } } for (planned, markers) in plugins.iter_mut().zip(dependency_markers) { @@ -1895,7 +1965,17 @@ impl PluginHost { ); } }; - staging.insert(planned.plugin.marker_type_id(), api); + match (planned.plugin.marker_type_id(), api) { + (Some(marker), Some(api)) => staging.insert(marker, api), + (None, None) => {} + _ => { + rollback.rollback().await; + anyhow::bail!( + "plugin `{}` returned an API inconsistent with its registration", + planned.manifest.id + ); + } + } self.abort_install_if_terminal(&mut rollback).await?; let Some(installed) = rollback.current.take() else { rollback.rollback().await; @@ -2320,9 +2400,9 @@ async fn plugin_callback<'a>( result? } -async fn plugin_install<'a>( - make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result>, -) -> anyhow::Result { +async fn plugin_install<'a, T>( + make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result>, +) -> anyhow::Result { let mut future = std::panic::catch_unwind(AssertUnwindSafe(make_future)) .map_err(|_| anyhow::anyhow!("install panicked before returning a future"))?; let result = AssertUnwindSafe(std::future::poll_fn(|context| { @@ -2421,6 +2501,40 @@ mod tests { log: Log, } + struct RuntimePluginAdapter { + id: &'static str, + dependency: Option<&'static str>, + log: Log, + } + + impl UntypedClientPlugin for RuntimePluginAdapter { + fn manifest(&self) -> PluginManifest { + let manifest = PluginManifest::new(self.id, "0.1.0"); + match self.dependency { + Some(dependency) => manifest.with_dependency(dependency), + None => manifest, + } + } + + fn install(&self, _context: PluginContext) -> BoxFuture<'_, anyhow::Result<()>> { + let id = self.id; + let log = Arc::clone(&self.log); + Box::pin(async move { + record(&log, format!("install:{id}")); + Ok(()) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let id = self.id; + let log = Arc::clone(&self.log); + Box::pin(async move { + record(&log, format!("shutdown:{id}")); + Ok(()) + }) + } + } + struct ShutdownDuringPluginInstall; struct CaptureInstallClient { @@ -2566,6 +2680,51 @@ mod tests { } } + #[tokio::test] + async fn untyped_instances_share_an_adapter_type_and_remain_manifest_keyed() { + let log = Arc::new(Mutex::new(Vec::new())); + let client = complete_builder() + .await + .with_untyped_plugin(RuntimePluginAdapter { + id: "runtime-dependent", + dependency: Some("runtime-foundation"), + log: Arc::clone(&log), + }) + .with_untyped_plugin(RuntimePluginAdapter { + id: "runtime-foundation", + dependency: None, + log: Arc::clone(&log), + }) + .build() + .await + .expect("untyped plugin plan") + .into_client(); + + assert_eq!( + client + .plugin_manifests() + .iter() + .map(PluginManifest::id) + .collect::>(), + vec!["runtime-foundation", "runtime-dependent"] + ); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["install:runtime-foundation", "install:runtime-dependent"] + ); + + client.disconnect().await; + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec![ + "install:runtime-foundation", + "install:runtime-dependent", + "shutdown:runtime-dependent", + "shutdown:runtime-foundation" + ] + ); + } + #[tokio::test] async fn shutdown_during_upstream_install_prevents_plugin_installation() { let log = Arc::new(Mutex::new(Vec::new())); From 03c4b8bb5d3a34756e3504307e9ea4a436252642 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:44:24 -0300 Subject: [PATCH 41/46] feat(plugins): own core event subscriptions --- agent_docs/plugin_architecture.md | 7 +- plugins/metrics/src/lib.rs | 9 +- src/lib.rs | 25 +-- src/plugins/mod.rs | 275 ++++++++++++++++++++++++++---- 4 files changed, 261 insertions(+), 55 deletions(-) diff --git a/agent_docs/plugin_architecture.md b/agent_docs/plugin_architecture.md index 913661a16..a5d86644c 100644 --- a/agent_docs/plugin_architecture.md +++ b/agent_docs/plugin_architecture.md @@ -215,9 +215,10 @@ Core events remain the sealed `wacore::types::events::Event` contract. Subscriptions use explicit `EventInterest`; interest changes go through the retained `Subscription`, and the aggregate 128-bit mask provides the producer fast path. Plugin core handlers run inline, must not block, and should hand work -to a task capability. `PluginCoreEvents` retains its subscriptions for the -plugin lifetime; early removal is not part of the initial plugin API. Requesting -`RawNode` retains a forwarding lease for the same lifetime. +to a task capability. `PluginCoreEvents::subscribe` returns an owned token; +dropping or explicitly unsubscribing it removes the handler immediately, while +host shutdown invalidates tokens retained by plugin APIs. Updating its interest +also acquires or releases the `RawNode` forwarding lease in the same operation. Custom events never enter the core enum or consume an `EventInterest` bit. `PluginEventRouter` routes exact `(plugin_id, topic)` selectors and gives each diff --git a/plugins/metrics/src/lib.rs b/plugins/metrics/src/lib.rs index 7eb85186c..3f11f32ad 100644 --- a/plugins/metrics/src/lib.rs +++ b/plugins/metrics/src/lib.rs @@ -27,8 +27,9 @@ use portable_atomic::{AtomicBool, AtomicU64, Ordering}; use serde::{Deserialize, Serialize}; use whatsapp_rust::wacore::types::events::{Event, EventHandler, EventInterest, EventKind}; use whatsapp_rust::{ - ClientPlugin, PluginCapability, PluginConnectionScope, PluginEventPayloadEncoding, - PluginEventSelector, PluginEventTopic, PluginEvents, PluginFuture, PluginManifest, PluginTasks, + ClientPlugin, PluginCapability, PluginConnectionScope, PluginCoreEventSubscription, + PluginEventPayloadEncoding, PluginEventSelector, PluginEventTopic, PluginEvents, PluginFuture, + PluginManifest, PluginTasks, }; pub const METRICS_PLUGIN_ID: &str = "wa.metrics"; @@ -159,6 +160,7 @@ impl EventHandler for MetricsEventHandler { pub struct MetricsApi { state: Arc, tick_selector: PluginEventSelector, + _core_events: PluginCoreEventSubscription, } impl MetricsApi { @@ -233,7 +235,7 @@ impl ClientPlugin for MetricsPlugin { .set(state.clone()) .map_err(|_| anyhow::anyhow!("metrics plugin was installed more than once"))?; - core_events.subscribe( + let core_events = core_events.subscribe( EventInterest::of(&[ EventKind::Messages, EventKind::Receipt, @@ -246,6 +248,7 @@ impl ClientPlugin for MetricsPlugin { let api = Arc::new(MetricsApi { state: state.clone(), tick_selector: plugin_events.selector(&tick), + _core_events: core_events, }); spawn_install_ticker(tasks, plugin_events, tick, state, self.interval)?; Ok(api) diff --git a/src/lib.rs b/src/lib.rs index cf213d9cb..84111d3b0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -119,15 +119,15 @@ pub mod plugins; #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] pub use plugins::{ ClientPlugin, PluginCapabilities, PluginCapability, PluginConnectionScope, - PluginConnectionTasks, PluginContext, PluginCoreEvents, PluginEventEndpointConfig, - PluginEventEndpointStats, PluginEventEnvelope, PluginEventOverflow, PluginEventPayloadEncoding, - PluginEventPublishError, PluginEventPublishReport, PluginEventPublisherStats, - PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, PluginEventRouterStats, - PluginEventSelector, PluginEventSubscribeError, PluginEventSubscription, PluginEventTopic, - PluginEventTryReceiveError, PluginEvents, PluginFuture, PluginHealth, PluginHostStats, - PluginIq, PluginIqError, PluginManifest, PluginMessaging, PluginMessagingError, - PluginPlanError, PluginResourceError, PluginState, PluginStats, PluginTasks, - UntypedClientPlugin, + PluginConnectionTasks, PluginContext, PluginCoreEventSubscription, PluginCoreEvents, + PluginEventEndpointConfig, PluginEventEndpointStats, PluginEventEnvelope, PluginEventOverflow, + PluginEventPayloadEncoding, PluginEventPublishError, PluginEventPublishReport, + PluginEventPublisherStats, PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, + PluginEventRouterStats, PluginEventSelector, PluginEventSubscribeError, + PluginEventSubscription, PluginEventTopic, PluginEventTryReceiveError, PluginEvents, + PluginFuture, PluginHealth, PluginHostStats, PluginIq, PluginIqError, PluginManifest, + PluginMessaging, PluginMessagingError, PluginPlanError, PluginResourceError, PluginState, + PluginStats, PluginTasks, UntypedClientPlugin, }; pub mod request; pub(crate) mod signal_flush; @@ -206,9 +206,10 @@ pub mod prelude { #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] pub use crate::plugins::{ ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, - PluginEventEndpointConfig, PluginEventOverflow, PluginEventPayloadEncoding, - PluginEventRouter, PluginEventSelector, PluginEventSubscription, PluginEventTopic, - PluginEvents, PluginFuture, PluginManifest, UntypedClientPlugin, + PluginCoreEventSubscription, PluginEventEndpointConfig, PluginEventOverflow, + PluginEventPayloadEncoding, PluginEventRouter, PluginEventSelector, + PluginEventSubscription, PluginEventTopic, PluginEvents, PluginFuture, PluginManifest, + UntypedClientPlugin, }; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 8abc68677..c6ef5c4bc 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -314,7 +314,7 @@ struct PluginResources { shutdown: ShutdownNotifier, install_tasks: Arc, connection_tasks: Mutex, - subscriptions: Mutex>, + subscriptions: Mutex>>, teardown_panics: AtomicU64, } @@ -413,9 +413,134 @@ impl Drop for TaskLease { } } -struct PluginCoreEventSubscription { - _subscription: Subscription, - _raw_node_lease: Option, +struct PluginCoreEventSubscriptionState { + subscription: Option, + raw_node_lease: Option, + interest: EventInterest, +} + +struct PluginCoreEventSubscriptionInner { + client: Weak, + resources: Weak, + plugin_id: Arc, + state: Mutex, +} + +impl PluginCoreEventSubscriptionInner { + fn is_active(&self) -> bool { + self.state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .subscription + .is_some() + } + + fn update_interest(&self, interest: EventInterest) -> Result { + let resources = self + .resources + .upgrade() + .ok_or(PluginResourceError::ShuttingDown)?; + if resources.closed.load(Ordering::Acquire) { + return Err(PluginResourceError::ShuttingDown); + } + + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let wants_raw_node = interest.wants(EventKind::RawNode); + let acquired_raw_node_lease = if wants_raw_node && state.raw_node_lease.is_none() { + Some( + self.client + .upgrade() + .ok_or(PluginResourceError::ClientUnavailable)? + .acquire_raw_node_forwarding(), + ) + } else { + None + }; + let Some(subscription) = state.subscription.as_ref() else { + return Ok(false); + }; + if !subscription.update_interest(interest) { + let registration = (state.subscription.take(), state.raw_node_lease.take()); + drop(state); + drop(acquired_raw_node_lease); + drop(registration); + return Ok(false); + } + + state.interest = interest; + if let Some(lease) = acquired_raw_node_lease { + state.raw_node_lease = Some(lease); + } + let retired_raw_node_lease = (!wants_raw_node) + .then(|| state.raw_node_lease.take()) + .flatten(); + drop(state); + drop(retired_raw_node_lease); + Ok(true) + } + + fn close(&self) -> bool { + let registration = { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + (state.subscription.take(), state.raw_node_lease.take()) + }; + let active = registration.0.is_some(); + if std::panic::catch_unwind(AssertUnwindSafe(|| drop(registration))).is_err() { + if let Some(resources) = self.resources.upgrade() { + resources.teardown_panics.fetch_add(1, Ordering::Relaxed); + } + log::warn!( + "Plugin `{}` core-event subscription panicked while closing", + self.plugin_id + ); + } + active + } +} + +/// Ownership token for one plugin core-event subscription. +/// +/// Dropping the token unsubscribes immediately. Host shutdown also invalidates +/// a retained token, so keeping it in a plugin API cannot extend client work. +#[must_use = "dropping the token immediately unregisters the plugin event handler"] +pub struct PluginCoreEventSubscription { + inner: Arc, +} + +impl PluginCoreEventSubscription { + pub fn interest(&self) -> EventInterest { + self.inner + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .interest + } + + /// Replace the filter while preserving the handler registration. + pub fn update_interest(&self, interest: EventInterest) -> Result { + self.inner.update_interest(interest) + } + + pub fn is_active(&self) -> bool { + self.inner.is_active() + } + + /// Remove the handler now instead of waiting for `Drop`. + pub fn unsubscribe(&self) -> bool { + self.inner.close() + } +} + +impl Drop for PluginCoreEventSubscription { + fn drop(&mut self) { + self.inner.close(); + } } impl PluginResources { @@ -463,30 +588,43 @@ impl PluginResources { fn retain_subscription( &self, + client: Weak, + resources: Weak, + plugin_id: Arc, + interest: EventInterest, subscription: Subscription, raw_node_lease: Option, - ) -> Result<(), PluginResourceError> { - let registration = PluginCoreEventSubscription { - _subscription: subscription, - _raw_node_lease: raw_node_lease, - }; + ) -> Result { + let registration = Arc::new(PluginCoreEventSubscriptionInner { + client, + resources, + plugin_id, + state: Mutex::new(PluginCoreEventSubscriptionState { + subscription: Some(subscription), + raw_node_lease, + interest, + }), + }); let rejected = { let mut subscriptions = self .subscriptions .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); if self.closed.load(Ordering::Acquire) { - Some(registration) + true } else { - subscriptions.push(registration); - None + subscriptions.retain(|subscription| subscription.is_active()); + subscriptions.push(Arc::clone(®istration)); + false } }; - if let Some(rejected) = rejected { - drop(rejected); + if rejected { + registration.close(); Err(PluginResourceError::ShuttingDown) } else { - Ok(()) + Ok(PluginCoreEventSubscription { + inner: registration, + }) } } @@ -594,10 +732,7 @@ impl PluginResources { std::mem::take(&mut *subscriptions) }; for subscription in subscriptions { - if std::panic::catch_unwind(AssertUnwindSafe(|| drop(subscription))).is_err() { - self.teardown_panics.fetch_add(1, Ordering::Relaxed); - log::warn!("Plugin core-event subscription panicked while being dropped"); - } + subscription.close(); } } } @@ -634,7 +769,9 @@ impl PluginResources { .subscriptions .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) - .len(), + .iter() + .filter(|subscription| subscription.is_active()) + .count(), teardown_panics: self.teardown_panics.load(Ordering::Relaxed), } } @@ -888,7 +1025,7 @@ impl PluginCoreEvents { &self, interest: EventInterest, handler: Arc, - ) -> Result<(), PluginResourceError> { + ) -> Result { let client = self .client .upgrade() @@ -903,8 +1040,14 @@ impl PluginCoreEvents { diagnostics: Arc::clone(&self.diagnostics), }); let subscription = client.subscribe(interest, handler); - self.resources - .retain_subscription(subscription, raw_node_lease) + self.resources.retain_subscription( + self.client.clone(), + Arc::downgrade(&self.resources), + Arc::clone(&self.plugin_id), + interest, + subscription, + raw_node_lease, + ) } } @@ -4317,7 +4460,7 @@ mod tests { struct PanickingCoreEventPlugin; impl ClientPlugin for PanickingCoreEventPlugin { - type Api = (); + type Api = PluginCoreEventSubscription; fn manifest(&self) -> PluginManifest { PluginManifest::new("panicking-core-event", "0.1.0") @@ -4326,14 +4469,14 @@ mod tests { fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { Box::pin(async move { - context + let subscription = context .core_events() .ok_or_else(|| anyhow::anyhow!("core events capability missing"))? .subscribe( EventInterest::of(&[EventKind::Connected]), Arc::new(PanickingCoreEventHandler), )?; - Ok(Arc::new(())) + Ok(Arc::new(subscription)) }) } } @@ -4341,7 +4484,7 @@ mod tests { struct PanickingSubscriptionPlugin; impl ClientPlugin for PanickingSubscriptionPlugin { - type Api = (); + type Api = PluginCoreEventSubscription; fn manifest(&self) -> PluginManifest { PluginManifest::new("panicking-subscription", "0.1.0") @@ -4350,14 +4493,14 @@ mod tests { fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { Box::pin(async move { - context + let subscription = context .core_events() .ok_or_else(|| anyhow::anyhow!("core events capability missing"))? .subscribe( EventInterest::of(&[EventKind::Connected]), Arc::new(PanickingDropEventHandler), )?; - Ok(Arc::new(())) + Ok(Arc::new(subscription)) }) } } @@ -4391,7 +4534,7 @@ mod tests { } impl ClientPlugin for EventSubscriptionPlugin { - type Api = (); + type Api = PluginCoreEventSubscription; fn manifest(&self) -> PluginManifest { PluginManifest::new("event-subscription", "0.1.0") @@ -4400,14 +4543,14 @@ mod tests { fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { Box::pin(async move { - context + let subscription = context .core_events() .ok_or_else(|| anyhow::anyhow!("core events capability missing"))? .subscribe( EventInterest::of(&[EventKind::Connected, EventKind::RawNode]), Arc::new(NoopEventHandler), )?; - Ok(Arc::new(())) + Ok(Arc::new(subscription)) }) } } @@ -4431,8 +4574,13 @@ mod tests { struct ReentrantSubscriptionPlugin; + struct ReentrantSubscriptionApi { + events: PluginCoreEvents, + _subscription: PluginCoreEventSubscription, + } + impl ClientPlugin for ReentrantSubscriptionPlugin { - type Api = PluginCoreEvents; + type Api = ReentrantSubscriptionApi; fn manifest(&self) -> PluginManifest { PluginManifest::new("reentrant-subscription", "0.1.0") @@ -4445,13 +4593,16 @@ mod tests { .core_events() .cloned() .ok_or_else(|| anyhow::anyhow!("core events capability missing"))?; - events.subscribe( + let subscription = events.subscribe( EventInterest::of(&[EventKind::Connected]), Arc::new(ReentrantSubscriptionHandler { events: events.clone(), }), )?; - Ok(Arc::new(events)) + Ok(Arc::new(ReentrantSubscriptionApi { + events, + _subscription: subscription, + })) }) } } @@ -4483,6 +4634,56 @@ mod tests { assert!(!client.raw_node_forwarding_enabled()); } + #[tokio::test] + async fn plugin_subscription_updates_interest_and_can_unsubscribe_early() { + let client = complete_builder() + .await + .with_plugin(EventSubscriptionPlugin) + .build() + .await + .expect("event subscription plugin") + .into_client(); + let subscription = client + .plugin::() + .expect("subscription API"); + + assert!(subscription.is_active()); + assert!(subscription.interest().wants(EventKind::RawNode)); + assert!(client.raw_node_forwarding_enabled()); + assert!( + subscription + .update_interest(EventInterest::of(&[EventKind::Connected])) + .expect("interest update") + ); + assert!(!client.raw_node_forwarding_enabled()); + assert!(!subscription.interest().wants(EventKind::RawNode)); + assert!(client.core.event_bus.has_handler_for(EventKind::Connected)); + assert!( + subscription + .update_interest(EventInterest::of(&[EventKind::RawNode])) + .expect("raw-node interest update") + ); + assert!(client.raw_node_forwarding_enabled()); + assert!(!client.core.event_bus.has_handler_for(EventKind::Connected)); + + assert!(subscription.unsubscribe()); + assert!(!subscription.is_active()); + assert!(!subscription.unsubscribe()); + assert!(!client.raw_node_forwarding_enabled()); + assert!(!client.core.event_bus.has_handler_for(EventKind::Connected)); + assert_eq!( + client + .plugin_stats() + .expect("plugin stats") + .plugins + .first() + .expect("plugin stats entry") + .core_event_subscriptions, + 0 + ); + client.disconnect().await; + } + #[tokio::test] async fn panicking_core_event_handler_is_isolated_and_degrades_only_its_plugin() { let client = complete_builder() @@ -4615,10 +4816,10 @@ mod tests { let (completed_tx, completed_rx) = std::sync::mpsc::sync_channel(1); let subscribe_events = events.clone(); let subscribe = std::thread::spawn(move || { - let result = subscribe_events.subscribe( + let result = subscribe_events.events.subscribe( EventInterest::of(&[EventKind::Connected]), Arc::new(ReentrantSubscriptionHandler { - events: (*subscribe_events).clone(), + events: subscribe_events.events.clone(), }), ); let _ = completed_tx.send(result); From 1444153ba994146b16193235ff840f0a2e4c391d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:53:44 -0300 Subject: [PATCH 42/46] feat(plugins): support cooperative task draining --- agent_docs/plugin_architecture.md | 18 +- src/bot.rs | 20 +- src/client/builder.rs | 34 ++- src/lib.rs | 10 +- src/plugins/mod.rs | 408 ++++++++++++++++++++++++++++-- 5 files changed, 453 insertions(+), 37 deletions(-) diff --git a/agent_docs/plugin_architecture.md b/agent_docs/plugin_architecture.md index a5d86644c..c9fc48404 100644 --- a/agent_docs/plugin_architecture.md +++ b/agent_docs/plugin_architecture.md @@ -191,13 +191,17 @@ install once readied first and closed/shut down last. `PluginTasks` survives reconnects and receives cancellation only on rollback or -terminal shutdown. `PluginConnectionTasks` is tied to one generation: -cancellation is signalled synchronously, while actual future destruction is -cooperative at executor poll boundaries. Plugin tasks must not block an executor -thread or detach untracked work. A non-cooperative task may outlive the bounded -drain, in which case teardown proceeds and diagnostics remain degraded. - -Callbacks are serialized, bounded by timeouts, and isolated from panics, +terminal shutdown. `PluginConnectionTasks` is tied to one generation. Their +default `spawn` drops the future when cancellation wins. `spawn_cooperative` +instead keeps polling accepted work after signalling shutdown; that future must +observe `shutdown_signal()` or `cancellation_signal()` and finish itself. The +host waits only through the configured task-drain deadline, then proceeds and +marks the plugin degraded if work remains. Plugin tasks must not block an +executor thread or detach untracked work. + +`PluginHostConfig` independently configures per-callback and per-task-drain +deadlines; both default to five seconds and reject zero. Callbacks are +serialized, bounded by those timeouts, and isolated from panics, including panics while constructing, polling, cancelling, or destroying their futures. One faulty plugin must not suppress later callbacks. Stale `Ready` work is bounded under reconnect pressure; every accepted `Closed` callback is diff --git a/src/bot.rs b/src/bot.rs index b7d6b8628..0fc419378 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -2,7 +2,7 @@ use crate::cache_config::CacheConfig; use crate::client::{Client, ClientBuilderError}; use crate::pair_code::PairCodeOptions; #[cfg(feature = "plugins")] -use crate::plugins::{ClientPlugin, PluginRegistration, UntypedClientPlugin}; +use crate::plugins::{ClientPlugin, PluginHostConfig, PluginRegistration, UntypedClientPlugin}; use crate::store::commands::DeviceCommand; use crate::store::error::StoreError; use crate::store::persistence_manager::PersistenceManager; @@ -636,6 +636,8 @@ pub struct BotBuilder< alloc_meter: Option>, #[cfg(feature = "plugins")] plugins: Vec, + #[cfg(feature = "plugins")] + plugin_host_config: PluginHostConfig, _marker: PhantomData<(B, T, H, R)>, } @@ -663,6 +665,8 @@ impl BotBuilder BotBuilder { alloc_meter: self.alloc_meter, #[cfg(feature = "plugins")] plugins: self.plugins, + #[cfg(feature = "plugins")] + plugin_host_config: self.plugin_host_config, _marker: PhantomData, } } @@ -846,6 +852,14 @@ impl BotBuilder { self } + /// Configure plugin lifecycle and tracked-task deadlines. + #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] + pub fn with_plugin_host_config(mut self, config: PluginHostConfig) -> Self { + self.plugin_host_config = config; + self + } + // ── Event handler registration (additive; order of registration is kept, // but handlers run on their own tasks, so cross-event ordering is not // guaranteed) ────────────────────────────────────────────────────── @@ -1293,7 +1307,9 @@ impl BotBuilder { .with_skip_history_sync(self.skip_history_sync) .with_background_saver_interval(std::time::Duration::from_secs(30)); #[cfg(feature = "plugins")] - let client_builder = client_builder.with_plugin_registrations(self.plugins); + let client_builder = client_builder + .with_plugin_registrations(self.plugins) + .with_plugin_host_config(self.plugin_host_config); let mut client_builder = client_builder; if let Some(version) = self.override_version { diff --git a/src/client/builder.rs b/src/client/builder.rs index 81ec9982f..4ee2816dc 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -11,7 +11,8 @@ use crate::cache_config::CacheConfig; use crate::http::HttpClient; #[cfg(feature = "plugins")] use crate::plugins::{ - ClientPlugin, PluginHost, PluginPlan, PluginPlanError, PluginRegistration, UntypedClientPlugin, + ClientPlugin, PluginHost, PluginHostConfig, PluginPlan, PluginPlanError, PluginRegistration, + UntypedClientPlugin, }; use crate::store::error::StoreError; use crate::store::persistence_manager::PersistenceManager; @@ -69,6 +70,14 @@ pub enum ClientBuilderError { MissingHttpClient, #[error("background saver interval must be greater than zero")] InvalidBackgroundSaverInterval, + #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] + #[error("plugin callback timeout must be greater than zero")] + InvalidPluginCallbackTimeout, + #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] + #[error("plugin task-drain timeout must be greater than zero")] + InvalidPluginTaskDrainTimeout, #[error("the configured backend does not support the inbound durability hook: {0}")] UnsupportedDurabilityBackend(String), #[cfg(feature = "client-lifecycle")] @@ -109,6 +118,8 @@ pub struct ClientBuilder { lifecycle: Option>, #[cfg(feature = "plugins")] plugins: Vec, + #[cfg(feature = "plugins")] + plugin_host_config: PluginHostConfig, } impl Default for ClientBuilder { @@ -139,6 +150,8 @@ impl ClientBuilder { lifecycle: None, #[cfg(feature = "plugins")] plugins: Vec::new(), + #[cfg(feature = "plugins")] + plugin_host_config: PluginHostConfig::default(), } } @@ -339,6 +352,14 @@ impl ClientBuilder { self } + /// Configure plugin lifecycle and tracked-task deadlines. + #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] + pub fn with_plugin_host_config(mut self, config: PluginHostConfig) -> Self { + self.plugin_host_config = config; + self + } + #[cfg(feature = "plugins")] pub(crate) fn with_plugin_registrations( mut self, @@ -383,6 +404,15 @@ impl ClientBuilder { return Err(ClientBuilderError::InvalidBackgroundSaverInterval); } + #[cfg(feature = "plugins")] + if self.plugin_host_config.callback_timeout() == Duration::ZERO { + return Err(ClientBuilderError::InvalidPluginCallbackTimeout); + } + #[cfg(feature = "plugins")] + if self.plugin_host_config.task_drain_timeout() == Duration::ZERO { + return Err(ClientBuilderError::InvalidPluginTaskDrainTimeout); + } + if self.inbound_durability_hook.is_some() { probe_durability_backend(&persistence_manager.backend()).await?; } @@ -435,7 +465,7 @@ impl ClientBuilder { let (lifecycle_handler, plugin_host) = { let mut lifecycle_handler = lifecycle_handler; let plugin_host = plugin_plan.map(|plan| { - let host = PluginHost::new(plan, lifecycle_handler.take()); + let host = PluginHost::new(plan, lifecycle_handler.take(), self.plugin_host_config); lifecycle_handler = Some(host.clone()); host }); diff --git a/src/lib.rs b/src/lib.rs index 84111d3b0..968c0f6b6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -125,9 +125,9 @@ pub use plugins::{ PluginEventPublisherStats, PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, PluginEventRouterStats, PluginEventSelector, PluginEventSubscribeError, PluginEventSubscription, PluginEventTopic, PluginEventTryReceiveError, PluginEvents, - PluginFuture, PluginHealth, PluginHostStats, PluginIq, PluginIqError, PluginManifest, - PluginMessaging, PluginMessagingError, PluginPlanError, PluginResourceError, PluginState, - PluginStats, PluginTasks, UntypedClientPlugin, + PluginFuture, PluginHealth, PluginHostConfig, PluginHostStats, PluginIq, PluginIqError, + PluginManifest, PluginMessaging, PluginMessagingError, PluginPlanError, PluginResourceError, + PluginState, PluginStats, PluginTasks, UntypedClientPlugin, }; pub mod request; pub(crate) mod signal_flush; @@ -208,8 +208,8 @@ pub mod prelude { ClientPlugin, PluginCapability, PluginConnectionScope, PluginContext, PluginCoreEventSubscription, PluginEventEndpointConfig, PluginEventOverflow, PluginEventPayloadEncoding, PluginEventRouter, PluginEventSelector, - PluginEventSubscription, PluginEventTopic, PluginEvents, PluginFuture, PluginManifest, - UntypedClientPlugin, + PluginEventSubscription, PluginEventTopic, PluginEvents, PluginFuture, PluginHostConfig, + PluginManifest, UntypedClientPlugin, }; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index c6ef5c4bc..bdf956a9a 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -41,7 +41,8 @@ const CAP_TASKS: u8 = 1 << 1; const CAP_MESSAGING: u8 = 1 << 2; const CAP_IQ: u8 = 1 << 3; const CAP_PLUGIN_EVENTS: u8 = 1 << 4; -const PLUGIN_CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); +const DEFAULT_PLUGIN_CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); +const DEFAULT_PLUGIN_TASK_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); /// A capability a plugin asks the host to expose during installation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -92,6 +93,48 @@ impl PluginCapabilities { } } +/// Deadlines applied by the native plugin host. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PluginHostConfig { + callback_timeout: Duration, + task_drain_timeout: Duration, +} + +impl PluginHostConfig { + pub const fn new() -> Self { + Self { + callback_timeout: DEFAULT_PLUGIN_CALLBACK_TIMEOUT, + task_drain_timeout: DEFAULT_PLUGIN_TASK_DRAIN_TIMEOUT, + } + } + + /// Bound each `on_ready`, `on_closed`, and `shutdown` callback. + pub const fn with_callback_timeout(mut self, timeout: Duration) -> Self { + self.callback_timeout = timeout; + self + } + + /// Bound each install- or connection-scoped task drain. + pub const fn with_task_drain_timeout(mut self, timeout: Duration) -> Self { + self.task_drain_timeout = timeout; + self + } + + pub const fn callback_timeout(self) -> Duration { + self.callback_timeout + } + + pub const fn task_drain_timeout(self) -> Duration { + self.task_drain_timeout + } +} + +impl Default for PluginHostConfig { + fn default() -> Self { + Self::new() + } +} + /// Build-time declaration used for validation, ordering, and future foreign adapters. #[derive(Debug, Clone, PartialEq, Eq)] #[non_exhaustive] @@ -929,6 +972,26 @@ pub struct PluginTasks { impl PluginTasks { pub fn spawn(&self, future: F) -> Result<(), PluginResourceError> + where + F: Future + Spawnable, + { + self.spawn_with_mode(future, PluginTaskShutdown::Abort) + } + + /// Track work that must observe [`shutdown_signal`](Self::shutdown_signal) + /// and finish itself after shutdown is signalled. + pub fn spawn_cooperative(&self, future: F) -> Result<(), PluginResourceError> + where + F: Future + Spawnable, + { + self.spawn_with_mode(future, PluginTaskShutdown::Cooperative) + } + + fn spawn_with_mode( + &self, + future: F, + shutdown: PluginTaskShutdown, + ) -> Result<(), PluginResourceError> where F: Future + Spawnable, { @@ -943,6 +1006,7 @@ impl PluginTasks { Arc::clone(&self.plugin_id), lease, future, + shutdown, ); Ok(()) } @@ -1192,6 +1256,26 @@ pub struct PluginConnectionTasks { impl PluginConnectionTasks { pub fn spawn(&self, future: F) -> Result<(), PluginResourceError> + where + F: Future + Spawnable, + { + self.spawn_with_mode(future, PluginTaskShutdown::Abort) + } + + /// Track work that must observe [`cancellation_signal`](Self::cancellation_signal) + /// and finish itself after this generation is cancelled. + pub fn spawn_cooperative(&self, future: F) -> Result<(), PluginResourceError> + where + F: Future + Spawnable, + { + self.spawn_with_mode(future, PluginTaskShutdown::Cooperative) + } + + fn spawn_with_mode( + &self, + future: F, + shutdown: PluginTaskShutdown, + ) -> Result<(), PluginResourceError> where F: Future + Spawnable, { @@ -1206,10 +1290,15 @@ impl PluginConnectionTasks { Arc::clone(&self.plugin_id), lease, future, + shutdown, ); Ok(()) } + pub fn cancellation_signal(&self) -> ShutdownSignal { + self.scope.cancellation_signal() + } + /// Sleep through the configured runtime, returning promptly when this generation retires. pub async fn sleep(&self, duration: Duration) -> Result<(), PluginResourceError> { if self.scope.is_cancelled() { @@ -1305,6 +1394,12 @@ impl> Drop for GuardedPluginTask { } } +#[derive(Clone, Copy)] +enum PluginTaskShutdown { + Abort, + Cooperative, +} + fn spawn_until_cancelled( runtime: &Arc, cancellation: ShutdownSignal, @@ -1312,6 +1407,7 @@ fn spawn_until_cancelled( plugin_id: Arc, lease: TaskLease, future: F, + shutdown: PluginTaskShutdown, ) where F: Future + Spawnable, { @@ -1319,9 +1415,14 @@ fn spawn_until_cancelled( runtime .spawn(Box::pin(async move { let _lease = lease; - let cancelled = Box::pin(wait_for_shutdown(&cancellation)); let work = Box::pin(work); - let _ = futures::future::select(cancelled, work).await; + match shutdown { + PluginTaskShutdown::Abort => { + let cancelled = Box::pin(wait_for_shutdown(&cancellation)); + let _ = futures::future::select(cancelled, work).await; + } + PluginTaskShutdown::Cooperative => work.await, + } })) .detach(); } @@ -1333,6 +1434,7 @@ fn spawn_after_activation( plugin_id: Arc, lease: TaskLease, future: F, + shutdown: PluginTaskShutdown, ) where F: Future + Spawnable, { @@ -1353,9 +1455,14 @@ fn spawn_after_activation( if resources.closed.load(Ordering::Acquire) { return; } - let cancelled = Box::pin(wait_for_shutdown(&cancellation)); let work = Box::pin(work); - let _ = futures::future::select(cancelled, work).await; + match shutdown { + PluginTaskShutdown::Abort => { + let cancelled = Box::pin(wait_for_shutdown(&cancellation)); + let _ = futures::future::select(cancelled, work).await; + } + PluginTaskShutdown::Cooperative => work.await, + } })) .detach(); } @@ -1678,6 +1785,7 @@ struct InstalledPlugin { struct PluginInstallRollback { runtime: Arc, + config: PluginHostConfig, installed: Vec, current: Option, upstream: Option>, @@ -1686,9 +1794,10 @@ struct PluginInstallRollback { } impl PluginInstallRollback { - fn new(runtime: Arc, capacity: usize) -> Self { + fn new(runtime: Arc, capacity: usize, config: PluginHostConfig) -> Self { Self { runtime, + config, installed: Vec::with_capacity(capacity), current: None, upstream: None, @@ -1729,10 +1838,12 @@ impl PluginInstallRollback { let completion = completed.subscribe(); let runtime = self.runtime.clone(); let cleanup_runtime = runtime.clone(); + let config = self.config; runtime .spawn(Box::pin(async move { let result = AssertUnwindSafe(shutdown_staged_plugins( cleanup_runtime, + config, current, installed, upstream, @@ -1798,7 +1909,7 @@ pub(crate) struct PluginHost { apis: OnceLock>, runtime: OnceLock>, event_router: Option, - callback_timeout: Duration, + config: PluginHostConfig, terminal: AtomicBool, terminal_notifier: ShutdownNotifier, installing_resources: Mutex>>, @@ -1807,14 +1918,31 @@ pub(crate) struct PluginHost { } impl PluginHost { - pub(crate) fn new(plan: PluginPlan, upstream: Option>) -> Arc { - Self::new_with_callback_timeout(plan, upstream, PLUGIN_CALLBACK_TIMEOUT) + pub(crate) fn new( + plan: PluginPlan, + upstream: Option>, + config: PluginHostConfig, + ) -> Arc { + Self::new_with_config(plan, upstream, config) } + #[cfg(test)] fn new_with_callback_timeout( plan: PluginPlan, upstream: Option>, callback_timeout: Duration, + ) -> Arc { + Self::new_with_config( + plan, + upstream, + PluginHostConfig::new().with_callback_timeout(callback_timeout), + ) + } + + fn new_with_config( + plan: PluginPlan, + upstream: Option>, + config: PluginHostConfig, ) -> Arc { let manifests = plan .ordered @@ -1844,7 +1972,7 @@ impl PluginHost { apis: OnceLock::new(), runtime: OnceLock::new(), event_router, - callback_timeout, + config, terminal: AtomicBool::new(false), terminal_notifier: ShutdownNotifier::new(), installing_resources: Mutex::new(Vec::new()), @@ -1920,8 +2048,14 @@ impl PluginHost { .contains(PluginCapability::Tasks) }) .count(); - self.callback_timeout - .saturating_mul((callback_count + task_barrier_count) as u32) + self.config + .callback_timeout() + .saturating_mul(callback_count as u32) + .saturating_add( + self.config + .task_drain_timeout() + .saturating_mul(task_barrier_count as u32), + ) .saturating_add(Duration::from_secs(1)) } @@ -2023,7 +2157,12 @@ impl PluginHost { .runtime .get() .ok_or(PluginTaskDrainError::RuntimeUnavailable)?; - wait_for_plugin_tasks(&**runtime, self.callback_timeout, completion_signals).await + wait_for_plugin_tasks( + &**runtime, + self.config.task_drain_timeout(), + completion_signals, + ) + .await } async fn install_all(&self, client: Weak) -> anyhow::Result<()> { @@ -2044,7 +2183,8 @@ impl PluginHost { .unwrap_or_else(|poisoned| poisoned.into_inner()) .clear(); }); - let mut rollback = PluginInstallRollback::new(runtime.clone(), self.ordered.len()); + let mut rollback = + PluginInstallRollback::new(runtime.clone(), self.ordered.len(), self.config); self.abort_install_if_terminal(&mut rollback).await?; if let Some(upstream) = &self.upstream { if let Err(error) = plugin_callback(|| upstream.install(client.clone())).await { @@ -2228,7 +2368,7 @@ impl PluginHost { let runtime = self.runtime.get().ok_or_else(|| { PluginCallbackError::Callback(anyhow::anyhow!("plugin runtime is unavailable")) })?; - bounded_plugin_callback(&**runtime, self.callback_timeout, make_future).await + bounded_plugin_callback(&**runtime, self.config.callback_timeout(), make_future).await } fn record_upstream_callback(&self, result: &Result<(), PluginCallbackError>) { @@ -2417,6 +2557,7 @@ enum PluginTaskDrainError { async fn shutdown_staged_plugins( runtime: Arc, + config: PluginHostConfig, current: Option, mut installed: Vec, upstream: Option>, @@ -2425,7 +2566,7 @@ async fn shutdown_staged_plugins( if let Some(plugin) = current { let task_result = wait_for_plugin_tasks( &*runtime, - PLUGIN_CALLBACK_TIMEOUT, + config.task_drain_timeout(), plugin.resources.task_completion_signals(), ) .await; @@ -2436,7 +2577,7 @@ async fn shutdown_staged_plugins( plugin.manifest.id ); } - let callback_result = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { + let callback_result = bounded_plugin_callback(&*runtime, config.callback_timeout(), || { plugin.plugin.shutdown() }) .await; @@ -2452,7 +2593,7 @@ async fn shutdown_staged_plugins( while let Some(plugin) = installed.pop() { let task_result = wait_for_plugin_tasks( &*runtime, - PLUGIN_CALLBACK_TIMEOUT, + config.task_drain_timeout(), plugin.resources.task_completion_signals(), ) .await; @@ -2463,7 +2604,7 @@ async fn shutdown_staged_plugins( plugin.manifest.id ); } - let callback_result = bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || { + let callback_result = bounded_plugin_callback(&*runtime, config.callback_timeout(), || { plugin.plugin.shutdown() }) .await; @@ -2475,7 +2616,7 @@ async fn shutdown_staged_plugins( } if let Some(upstream) = upstream && let Err(error) = - bounded_plugin_callback(&*runtime, PLUGIN_CALLBACK_TIMEOUT, || upstream.shutdown()) + bounded_plugin_callback(&*runtime, config.callback_timeout(), || upstream.shutdown()) .await { log::warn!("Upstream lifecycle rollback failed: {error:#}"); @@ -2640,6 +2781,31 @@ mod tests { .with_http_client(MockHttpClient) } + #[tokio::test] + async fn rejects_zero_plugin_host_deadlines() { + let callback = complete_builder() + .await + .with_plugin_host_config(PluginHostConfig::new().with_callback_timeout(Duration::ZERO)) + .build() + .await; + assert!(matches!( + callback, + Err(ClientBuilderError::InvalidPluginCallbackTimeout) + )); + + let task_drain = complete_builder() + .await + .with_plugin_host_config( + PluginHostConfig::new().with_task_drain_timeout(Duration::ZERO), + ) + .build() + .await; + assert!(matches!( + task_drain, + Err(ClientBuilderError::InvalidPluginTaskDrainTimeout) + )); + } + struct FoundationPlugin { log: Log, } @@ -3358,7 +3524,8 @@ mod tests { let api: ErasedApi = Arc::new(TypedApi(Arc::new(PanickingDropApi))); registry.insert(TypeId::of::(), api); - let mut rollback = PluginInstallRollback::new(Arc::new(TokioRuntime), 1); + let mut rollback = + PluginInstallRollback::new(Arc::new(TokioRuntime), 1, PluginHostConfig::default()); rollback.installed.push(InstalledPlugin { plugin: erased_plugin, manifest, @@ -3849,6 +4016,123 @@ mod tests { shutdown_after_task: Arc, } + struct CooperativeTaskPlugin { + install_started: Arc, + install_finished: Arc, + connection_started: Arc, + connection_finished: Arc, + closed_after_task: Arc, + shutdown_after_task: Arc, + } + + impl ClientPlugin for CooperativeTaskPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("cooperative-tasks", "0.1.0") + .with_capability(PluginCapability::Tasks) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + let tasks = context + .tasks() + .cloned() + .ok_or_else(|| anyhow::anyhow!("tasks capability missing")); + let started = Arc::clone(&self.install_started); + let finished = Arc::clone(&self.install_finished); + Box::pin(async move { + let tasks = tasks?; + let shutdown = tasks.shutdown_signal(); + tasks.spawn_cooperative(async move { + started.store(true, Ordering::Release); + wait_for_shutdown(&shutdown).await; + tokio::task::yield_now().await; + finished.store(true, Ordering::Release); + })?; + Ok(Arc::new(())) + }) + } + + fn on_ready(&self, scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + let tasks = scope + .tasks() + .cloned() + .ok_or_else(|| anyhow::anyhow!("connection tasks capability missing")); + let started = Arc::clone(&self.connection_started); + let finished = Arc::clone(&self.connection_finished); + Box::pin(async move { + let tasks = tasks?; + let cancelled = tasks.cancellation_signal(); + tasks.spawn_cooperative(async move { + started.store(true, Ordering::Release); + wait_for_shutdown(&cancelled).await; + tokio::task::yield_now().await; + finished.store(true, Ordering::Release); + })?; + Ok(()) + }) + } + + fn on_closed(&self, _scope: PluginConnectionScope) -> BoxFuture<'_, anyhow::Result<()>> { + let finished = self.connection_finished.load(Ordering::Acquire); + let observed = Arc::clone(&self.closed_after_task); + Box::pin(async move { + observed.store(finished, Ordering::Release); + Ok(()) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let finished = self.install_finished.load(Ordering::Acquire); + let observed = Arc::clone(&self.shutdown_after_task); + Box::pin(async move { + observed.store(finished, Ordering::Release); + Ok(()) + }) + } + } + + struct TimedDrainPlugin { + started: Arc, + finished: Arc, + release_tx: async_channel::Sender<()>, + release_rx: async_channel::Receiver<()>, + } + + impl ClientPlugin for TimedDrainPlugin { + type Api = (); + + fn manifest(&self) -> PluginManifest { + PluginManifest::new("timed-drain", "0.1.0").with_capability(PluginCapability::Tasks) + } + + fn install(&self, context: PluginContext) -> BoxFuture<'_, anyhow::Result>> { + let tasks = context + .tasks() + .cloned() + .ok_or_else(|| anyhow::anyhow!("tasks capability missing")); + let started = Arc::clone(&self.started); + let finished = Arc::clone(&self.finished); + let release = self.release_rx.clone(); + Box::pin(async move { + tasks?.spawn_cooperative(async move { + started.store(true, Ordering::Release); + let _ = release.recv().await; + finished.store(true, Ordering::Release); + })?; + Ok(Arc::new(())) + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let release = self.release_tx.clone(); + Box::pin(async move { + let _ = release.try_send(()); + Ok(()) + }) + } + } + impl ClientPlugin for ScopedTaskPlugin { type Api = (); @@ -3998,6 +4282,88 @@ mod tests { assert_eq!(stats.plugins[0].callbacks_completed, 3); } + #[tokio::test] + async fn cooperative_tasks_drain_before_lifecycle_callbacks() { + let install_started = Arc::new(AtomicBool::new(false)); + let install_finished = Arc::new(AtomicBool::new(false)); + let connection_started = Arc::new(AtomicBool::new(false)); + let connection_finished = Arc::new(AtomicBool::new(false)); + let closed_after_task = Arc::new(AtomicBool::new(false)); + let shutdown_after_task = Arc::new(AtomicBool::new(false)); + let config = PluginHostConfig::new() + .with_callback_timeout(Duration::from_secs(1)) + .with_task_drain_timeout(Duration::from_secs(1)); + let client = complete_builder() + .await + .with_plugin_host_config(config) + .with_plugin(CooperativeTaskPlugin { + install_started: Arc::clone(&install_started), + install_finished: Arc::clone(&install_finished), + connection_started: Arc::clone(&connection_started), + connection_finished: Arc::clone(&connection_finished), + closed_after_task: Arc::clone(&closed_after_task), + shutdown_after_task: Arc::clone(&shutdown_after_task), + }) + .build() + .await + .expect("cooperative task plugin") + .into_client(); + let host = client.plugin_host.as_ref().expect("plugin host").clone(); + assert_eq!(host.config, config); + wait_for_flag(&install_started).await; + + let scope = ConnectionScope::new(212); + host.on_ready(scope.clone()) + .await + .expect("cooperative ready callback"); + wait_for_flag(&connection_started).await; + scope.cancel(); + host.on_closed(scope) + .await + .expect("cooperative closed callback"); + assert!(connection_finished.load(Ordering::Acquire)); + assert!(closed_after_task.load(Ordering::Acquire)); + + client.disconnect().await; + assert!(install_finished.load(Ordering::Acquire)); + assert!(shutdown_after_task.load(Ordering::Acquire)); + let stats = client.plugin_stats().expect("cooperative plugin stats"); + assert_eq!(stats.plugins[0].task_drain_timeouts, 0); + assert_eq!(stats.plugins[0].health, PluginHealth::Healthy); + } + + #[tokio::test] + async fn configured_task_drain_timeout_degrades_and_continues_shutdown() { + let started = Arc::new(AtomicBool::new(false)); + let finished = Arc::new(AtomicBool::new(false)); + let (release_tx, release_rx) = async_channel::bounded(1); + let client = complete_builder() + .await + .with_plugin_host_config( + PluginHostConfig::new() + .with_callback_timeout(Duration::from_secs(1)) + .with_task_drain_timeout(Duration::from_millis(10)), + ) + .with_plugin(TimedDrainPlugin { + started: Arc::clone(&started), + finished: Arc::clone(&finished), + release_tx, + release_rx, + }) + .build() + .await + .expect("timed drain plugin") + .into_client(); + wait_for_flag(&started).await; + + client.disconnect().await; + wait_for_flag(&finished).await; + let stats = client.plugin_stats().expect("timed drain stats"); + assert_eq!(stats.plugins[0].task_drain_timeouts, 1); + assert_eq!(stats.plugins[0].health, PluginHealth::Degraded); + assert_eq!(stats.plugins[0].state, PluginState::Stopped); + } + #[tokio::test] async fn connection_scoped_tasks_stop_when_the_generation_is_cancelled() { let scope = ConnectionScope::new(77); From 60fda0bce4648766ddcb6f09306224151d027dff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 22 Jul 2026 12:17:00 -0300 Subject: [PATCH 43/46] fix(plugins): bound partial installation cleanup --- agent_docs/plugin_architecture.md | 20 ++-- src/bot.rs | 5 +- src/client/builder.rs | 15 ++- src/plugins/mod.rs | 148 +++++++++++++++++++++++++++--- 4 files changed, 165 insertions(+), 23 deletions(-) diff --git a/agent_docs/plugin_architecture.md b/agent_docs/plugin_architecture.md index c9fc48404..5b7eb80a7 100644 --- a/agent_docs/plugin_architecture.md +++ b/agent_docs/plugin_architecture.md @@ -131,6 +131,8 @@ An adapter that represents runtime-defined plugins implements share one concrete Rust adapter type, and do not appear in `Client::plugin::

()`. Native plugins keep the typed path above; a future bridge multiplexes its language-specific handles behind the untyped adapter. +`with_untyped_plugin_arc(...)` also accepts a trait object when a host needs to +erase multiple adapter implementations before configuring the client. During installation, `PluginContext::plugin::

()` exposes only directly declared dependencies. The context keeps a weak dependency view so an API that @@ -199,14 +201,16 @@ host waits only through the configured task-drain deadline, then proceeds and marks the plugin degraded if work remains. Plugin tasks must not block an executor thread or detach untracked work. -`PluginHostConfig` independently configures per-callback and per-task-drain -deadlines; both default to five seconds and reject zero. Callbacks are -serialized, bounded by those timeouts, and isolated from panics, -including panics while constructing, polling, cancelling, or destroying their -futures. One faulty plugin must not suppress later callbacks. Stale `Ready` -work is bounded under reconnect pressure; every accepted `Closed` callback is -lossless and precedes terminal `Shutdown`, so the queue may temporarily exceed -its target to preserve cleanup. +`PluginHostConfig` independently configures installation, per-callback, and +per-task-drain deadlines. Installation defaults to thirty seconds; callbacks +and drains default to five seconds; all reject zero. A timed-out partial +installation is cancelled and follows the same LIFO rollback as an explicit +failure. Callbacks are serialized, bounded by their timeout, and isolated from +panics, including panics while constructing, polling, cancelling, or destroying +their futures. One faulty plugin must not suppress later callbacks. Stale +`Ready` work is bounded under reconnect pressure; every accepted `Closed` +callback is lossless and precedes terminal `Shutdown`, so the queue may +temporarily exceed its target to preserve cleanup. `signal_shutdown_sync()` closes tasks, subscriptions, event routes, and capability handles promptly. `disconnect().await` remains required for async diff --git a/src/bot.rs b/src/bot.rs index 0fc419378..1fa37cedf 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -846,7 +846,10 @@ impl BotBuilder { /// Register an already-shared manifest-ID-keyed plugin. #[cfg(feature = "plugins")] #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] - pub fn with_untyped_plugin_arc(mut self, plugin: Arc

) -> Self { + pub fn with_untyped_plugin_arc( + mut self, + plugin: Arc

, + ) -> Self { self.plugins .push(PluginRegistration::new_untyped_arc(plugin)); self diff --git a/src/client/builder.rs b/src/client/builder.rs index 4ee2816dc..fd63f0541 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -72,6 +72,10 @@ pub enum ClientBuilderError { InvalidBackgroundSaverInterval, #[cfg(feature = "plugins")] #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] + #[error("plugin install timeout must be greater than zero")] + InvalidPluginInstallTimeout, + #[cfg(feature = "plugins")] + #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] #[error("plugin callback timeout must be greater than zero")] InvalidPluginCallbackTimeout, #[cfg(feature = "plugins")] @@ -96,7 +100,7 @@ pub enum ClientBuilderError { /// Runtime-validated, low-level builder for [`Client`]. /// -/// Unlike [`crate::BotBuilder`], this builder deliberately does not use +/// Unlike [`crate::bot::BotBuilder`], this builder deliberately does not use /// typestate. FFI and embedded hosts can populate dependencies dynamically and /// receive a typed error without encoding Rust generic state in their wrapper. pub struct ClientBuilder { @@ -346,7 +350,10 @@ impl ClientBuilder { /// Register an already-shared manifest-ID-keyed plugin. #[cfg(feature = "plugins")] #[cfg_attr(docsrs, doc(cfg(feature = "plugins")))] - pub fn with_untyped_plugin_arc(mut self, plugin: Arc

) -> Self { + pub fn with_untyped_plugin_arc( + mut self, + plugin: Arc

, + ) -> Self { self.plugins .push(PluginRegistration::new_untyped_arc(plugin)); self @@ -404,6 +411,10 @@ impl ClientBuilder { return Err(ClientBuilderError::InvalidBackgroundSaverInterval); } + #[cfg(feature = "plugins")] + if self.plugin_host_config.install_timeout() == Duration::ZERO { + return Err(ClientBuilderError::InvalidPluginInstallTimeout); + } #[cfg(feature = "plugins")] if self.plugin_host_config.callback_timeout() == Duration::ZERO { return Err(ClientBuilderError::InvalidPluginCallbackTimeout); diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index bdf956a9a..ddc059c9d 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -41,6 +41,7 @@ const CAP_TASKS: u8 = 1 << 1; const CAP_MESSAGING: u8 = 1 << 2; const CAP_IQ: u8 = 1 << 3; const CAP_PLUGIN_EVENTS: u8 = 1 << 4; +const DEFAULT_PLUGIN_INSTALL_TIMEOUT: Duration = Duration::from_secs(30); const DEFAULT_PLUGIN_CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); const DEFAULT_PLUGIN_TASK_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); @@ -96,6 +97,7 @@ impl PluginCapabilities { /// Deadlines applied by the native plugin host. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct PluginHostConfig { + install_timeout: Duration, callback_timeout: Duration, task_drain_timeout: Duration, } @@ -103,11 +105,18 @@ pub struct PluginHostConfig { impl PluginHostConfig { pub const fn new() -> Self { Self { + install_timeout: DEFAULT_PLUGIN_INSTALL_TIMEOUT, callback_timeout: DEFAULT_PLUGIN_CALLBACK_TIMEOUT, task_drain_timeout: DEFAULT_PLUGIN_TASK_DRAIN_TIMEOUT, } } + /// Bound each plugin and upstream lifecycle installation. + pub const fn with_install_timeout(mut self, timeout: Duration) -> Self { + self.install_timeout = timeout; + self + } + /// Bound each `on_ready`, `on_closed`, and `shutdown` callback. pub const fn with_callback_timeout(mut self, timeout: Duration) -> Self { self.callback_timeout = timeout; @@ -120,6 +129,10 @@ impl PluginHostConfig { self } + pub const fn install_timeout(self) -> Duration { + self.install_timeout + } + pub const fn callback_timeout(self) -> Duration { self.callback_timeout } @@ -1565,9 +1578,9 @@ impl ErasedClientPlugin for PluginAdapter

{ } } -struct UntypedPluginAdapter

(Arc

); +struct UntypedPluginAdapter(Arc

); -impl ErasedClientPlugin for UntypedPluginAdapter

{ +impl ErasedClientPlugin for UntypedPluginAdapter

{ fn marker_type_id(&self) -> Option { None } @@ -1619,7 +1632,7 @@ impl PluginRegistration { Self::new_untyped_arc(Arc::new(plugin)) } - pub(crate) fn new_untyped_arc(plugin: Arc

) -> Self { + pub(crate) fn new_untyped_arc(plugin: Arc

) -> Self { Self { plugin: Arc::new(UntypedPluginAdapter(plugin)), } @@ -2187,11 +2200,16 @@ impl PluginHost { PluginInstallRollback::new(runtime.clone(), self.ordered.len(), self.config); self.abort_install_if_terminal(&mut rollback).await?; if let Some(upstream) = &self.upstream { - if let Err(error) = plugin_callback(|| upstream.install(client.clone())).await { - rollback.disarm(); + rollback.upstream = Some(upstream.clone()); + if let Err(error) = + bounded_plugin_install(&*runtime, self.config.install_timeout(), || { + upstream.install(client.clone()) + }) + .await + { + rollback.rollback().await; return Err(error); } - rollback.upstream = Some(upstream.clone()); self.abort_install_if_terminal(&mut rollback).await?; } @@ -2224,7 +2242,11 @@ impl PluginHost { } let terminal = self.terminal_notifier.subscribe(); let cancelled = Box::pin(wait_for_shutdown(&terminal)); - let install = Box::pin(plugin_install(|| planned.plugin.install(context))); + let install = Box::pin(bounded_plugin_install( + &*runtime, + self.config.install_timeout(), + || planned.plugin.install(context), + )); let install_result = match futures::future::select(cancelled, install).await { futures::future::Either::Left((_, install)) => { if std::panic::catch_unwind(AssertUnwindSafe(|| drop(install))).is_err() { @@ -2702,6 +2724,29 @@ async fn plugin_install<'a, T>( result? } +async fn bounded_plugin_install<'a, T>( + runtime: &dyn Runtime, + timeout: Duration, + make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result>, +) -> anyhow::Result { + let install = Box::pin(plugin_install(make_future)); + match futures::future::select(install, runtime.sleep(timeout)).await { + futures::future::Either::Left((result, _)) => result, + futures::future::Either::Right(((), install)) => { + if std::panic::catch_unwind(AssertUnwindSafe(|| drop(install))).is_err() { + anyhow::bail!( + "install timed out after {:.3} seconds and panicked while being cancelled", + timeout.as_secs_f64() + ); + } + anyhow::bail!( + "install timed out after {:.3} seconds", + timeout.as_secs_f64() + ) + } + } +} + fn finish_callbacks(stage: &str, failures: Vec) -> anyhow::Result<()> { if failures.is_empty() { Ok(()) @@ -2783,6 +2828,16 @@ mod tests { #[tokio::test] async fn rejects_zero_plugin_host_deadlines() { + let install = complete_builder() + .await + .with_plugin_host_config(PluginHostConfig::new().with_install_timeout(Duration::ZERO)) + .build() + .await; + assert!(matches!( + install, + Err(ClientBuilderError::InvalidPluginInstallTimeout) + )); + let callback = complete_builder() .await .with_plugin_host_config(PluginHostConfig::new().with_callback_timeout(Duration::ZERO)) @@ -2846,6 +2901,10 @@ mod tests { struct ShutdownDuringPluginInstall; + struct FailingInstallLifecycle { + log: Log, + } + struct CaptureInstallClient { client: async_channel::Sender>, } @@ -2962,6 +3021,24 @@ mod tests { } } + impl ClientLifecycle for FailingInstallLifecycle { + fn install(&self, _client: Weak) -> BoxFuture<'_, anyhow::Result<()>> { + let log = Arc::clone(&self.log); + Box::pin(async move { + record(&log, "install:failing-upstream"); + anyhow::bail!("injected upstream install failure") + }) + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + let log = Arc::clone(&self.log); + Box::pin(async move { + record(&log, "shutdown:failing-upstream"); + Ok(()) + }) + } + } + impl ClientPlugin for FoundationPlugin { type Api = String; @@ -2992,6 +3069,11 @@ mod tests { #[tokio::test] async fn untyped_instances_share_an_adapter_type_and_remain_manifest_keyed() { let log = Arc::new(Mutex::new(Vec::new())); + let foundation: Arc = Arc::new(RuntimePluginAdapter { + id: "runtime-foundation", + dependency: None, + log: Arc::clone(&log), + }); let client = complete_builder() .await .with_untyped_plugin(RuntimePluginAdapter { @@ -2999,11 +3081,7 @@ mod tests { dependency: Some("runtime-foundation"), log: Arc::clone(&log), }) - .with_untyped_plugin(RuntimePluginAdapter { - id: "runtime-foundation", - dependency: None, - log: Arc::clone(&log), - }) + .with_untyped_plugin_arc(foundation) .build() .await .expect("untyped plugin plan") @@ -3052,6 +3130,27 @@ mod tests { ); } + #[tokio::test] + async fn upstream_install_failure_runs_partial_rollback() { + let log = Arc::new(Mutex::new(Vec::new())); + let result = complete_builder() + .await + .with_lifecycle(FailingInstallLifecycle { + log: Arc::clone(&log), + }) + .with_plugin(FoundationPlugin { + log: Arc::clone(&log), + }) + .build() + .await; + + assert!(matches!(result, Err(ClientBuilderError::PluginInstall(_)))); + assert_eq!( + *log.lock().unwrap_or_else(|poisoned| poisoned.into_inner()), + vec!["install:failing-upstream", "shutdown:failing-upstream"] + ); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn plugin_surfaces_publish_only_after_final_activation() { let (client_tx, client_rx) = async_channel::bounded(1); @@ -3178,6 +3277,31 @@ mod tests { drop(client); } + #[tokio::test] + async fn install_timeout_rolls_back_the_partial_plugin() { + let (started_tx, started_rx) = async_channel::unbounded(); + let install_dropped = Arc::new(AtomicBool::new(false)); + let shutdown_called = Arc::new(AtomicBool::new(false)); + let result = complete_builder() + .await + .with_plugin_host_config( + PluginHostConfig::new().with_install_timeout(Duration::from_millis(10)), + ) + .with_plugin(TerminalBlockingInstallPlugin { + started: started_tx, + install_dropped: install_dropped.clone(), + shutdown_called: shutdown_called.clone(), + }) + .build() + .await; + + let resource_shutdown = started_rx.recv().await.expect("plugin install started"); + assert!(matches!(result, Err(ClientBuilderError::PluginInstall(_)))); + assert!(resource_shutdown.is_fired()); + assert!(install_dropped.load(Ordering::Acquire)); + assert!(shutdown_called.load(Ordering::Acquire)); + } + struct DependentPlugin { log: Log, } From bfdb2699a110bbff967492f53480a26f4e060f07 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 22 Jul 2026 12:44:53 -0300 Subject: [PATCH 44/46] fix(plugins): release closed subscription registry entries --- src/plugins/mod.rs | 94 ++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 87 insertions(+), 7 deletions(-) diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index ddc059c9d..4815ab873 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -370,7 +370,7 @@ struct PluginResources { shutdown: ShutdownNotifier, install_tasks: Arc, connection_tasks: Mutex, - subscriptions: Mutex>>, + subscriptions: Mutex>>, teardown_panics: AtomicU64, } @@ -519,10 +519,9 @@ impl PluginCoreEventSubscriptionInner { return Ok(false); }; if !subscription.update_interest(interest) { - let registration = (state.subscription.take(), state.raw_node_lease.take()); drop(state); drop(acquired_raw_node_lease); - drop(registration); + self.close(); return Ok(false); } @@ -539,6 +538,7 @@ impl PluginCoreEventSubscriptionInner { } fn close(&self) -> bool { + let resources = self.resources.upgrade(); let registration = { let mut state = self .state @@ -548,7 +548,7 @@ impl PluginCoreEventSubscriptionInner { }; let active = registration.0.is_some(); if std::panic::catch_unwind(AssertUnwindSafe(|| drop(registration))).is_err() { - if let Some(resources) = self.resources.upgrade() { + if let Some(resources) = &resources { resources.teardown_panics.fetch_add(1, Ordering::Relaxed); } log::warn!( @@ -556,6 +556,9 @@ impl PluginCoreEventSubscriptionInner { self.plugin_id ); } + if let Some(resources) = resources { + resources.forget_subscription(self); + } active } } @@ -669,8 +672,12 @@ impl PluginResources { if self.closed.load(Ordering::Acquire) { true } else { - subscriptions.retain(|subscription| subscription.is_active()); - subscriptions.push(Arc::clone(®istration)); + subscriptions.retain(|subscription| { + subscription + .upgrade() + .is_some_and(|subscription| subscription.is_active()) + }); + subscriptions.push(Arc::downgrade(®istration)); false } }; @@ -684,6 +691,16 @@ impl PluginResources { } } + fn forget_subscription(&self, subscription: &PluginCoreEventSubscriptionInner) { + let subscription_ptr = std::ptr::from_ref(subscription); + self.subscriptions + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .retain(|candidate| { + candidate.strong_count() != 0 && !std::ptr::eq(candidate.as_ptr(), subscription_ptr) + }); + } + fn connection_task_tracker(&self, generation: u64) -> (Arc, bool) { let mut registry = self .connection_tasks @@ -788,7 +805,9 @@ impl PluginResources { std::mem::take(&mut *subscriptions) }; for subscription in subscriptions { - subscription.close(); + if let Some(subscription) = subscription.upgrade() { + subscription.close(); + } } } } @@ -826,6 +845,7 @@ impl PluginResources { .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) .iter() + .filter_map(Weak::upgrade) .filter(|subscription| subscription.is_active()) .count(), teardown_panics: self.teardown_panics.load(Ordering::Relaxed), @@ -5174,6 +5194,66 @@ mod tests { client.disconnect().await; } + #[tokio::test] + async fn dropped_plugin_subscriptions_leave_no_retained_registry_entries() { + let client = complete_builder() + .await + .build() + .await + .expect("client") + .into_client(); + let resources = PluginResources::new(); + resources.activate(); + let diagnostics = PluginDiagnostics::new(); + diagnostics.attach_resources(&resources); + let events = PluginCoreEvents { + client: Arc::downgrade(&client), + resources: Arc::clone(&resources), + plugin_id: Arc::from("subscription-churn"), + diagnostics, + }; + + let subscriptions = (0..128) + .map(|_| { + events + .subscribe( + EventInterest::of(&[EventKind::Connected]), + Arc::new(NoopEventHandler), + ) + .expect("subscription") + }) + .collect::>(); + let registrations = subscriptions + .iter() + .map(|subscription| Arc::downgrade(&subscription.inner)) + .collect::>(); + assert_eq!( + resources + .subscriptions + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .len(), + subscriptions.len() + ); + + drop(subscriptions); + + assert!( + resources + .subscriptions + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .is_empty() + ); + assert!( + registrations + .into_iter() + .all(|registration| registration.upgrade().is_none()) + ); + assert_eq!(resources.stats().core_event_subscriptions, 0); + client.disconnect().await; + } + #[tokio::test] async fn panicking_core_event_handler_is_isolated_and_degrades_only_its_plugin() { let client = complete_builder() From 21b430c293396e3254f31dc575ec1489422e835b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 22 Jul 2026 12:53:31 -0300 Subject: [PATCH 45/46] docs(plugins): define subscription ownership --- agent_docs/plugin_architecture.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/agent_docs/plugin_architecture.md b/agent_docs/plugin_architecture.md index 5b7eb80a7..68324e12e 100644 --- a/agent_docs/plugin_architecture.md +++ b/agent_docs/plugin_architecture.md @@ -166,6 +166,12 @@ Capability handles keep `Weak` internally and reject calls before activation or after shutdown. This avoids `Client -> plugin API -> Client` cycles and gives terminal resource invalidation a synchronous boundary. +`PluginCoreEvents::subscribe` returns the ownership token for its registration. +Dropping or explicitly unsubscribing that token removes the handler, any +`RawNodeLease`, and its host registry entry immediately. The host indexes live +tokens weakly so terminal shutdown can invalidate retained tokens without +extending the lifetime of registrations that plugins already released. + ## Lifecycle and task ownership The host maps the client's existing `connection_generation` to From 30d06ff6d256621d79d1ea665440ea9925ac6b57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 22 Jul 2026 13:00:26 -0300 Subject: [PATCH 46/46] refactor(plugins): leave capability expansion room --- src/plugins/mod.rs | 35 ++++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 4815ab873..0adb4f696 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -36,11 +36,11 @@ use crate::client::{ClientLifecycle, ConnectionScope, ConnectionScopeState, RawN use crate::request::IqError; use crate::send::{SendError, SendResult}; -const CAP_CORE_EVENTS: u8 = 1 << 0; -const CAP_TASKS: u8 = 1 << 1; -const CAP_MESSAGING: u8 = 1 << 2; -const CAP_IQ: u8 = 1 << 3; -const CAP_PLUGIN_EVENTS: u8 = 1 << 4; +const CAP_CORE_EVENTS: u64 = 1 << 0; +const CAP_TASKS: u64 = 1 << 1; +const CAP_MESSAGING: u64 = 1 << 2; +const CAP_IQ: u64 = 1 << 3; +const CAP_PLUGIN_EVENTS: u64 = 1 << 4; const DEFAULT_PLUGIN_INSTALL_TIMEOUT: Duration = Duration::from_secs(30); const DEFAULT_PLUGIN_CALLBACK_TIMEOUT: Duration = Duration::from_secs(5); const DEFAULT_PLUGIN_TASK_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); @@ -67,7 +67,7 @@ impl PluginCapability { } } - const fn bit(self) -> u8 { + const fn bit(self) -> u64 { match self { Self::CoreEvents => CAP_CORE_EVENTS, Self::Tasks => CAP_TASKS, @@ -80,7 +80,7 @@ impl PluginCapability { /// Compact set of capabilities requested by one plugin. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct PluginCapabilities(u8); +pub struct PluginCapabilities(u64); impl PluginCapabilities { pub const NONE: Self = Self(0); @@ -2846,6 +2846,27 @@ mod tests { .with_http_client(MockHttpClient) } + #[test] + fn capability_bits_are_distinct_and_composable() { + let capabilities = [ + PluginCapability::CoreEvents, + PluginCapability::Tasks, + PluginCapability::Messaging, + PluginCapability::Iq, + PluginCapability::PluginEvents, + ]; + let combined = capabilities + .into_iter() + .fold(PluginCapabilities::NONE, PluginCapabilities::with); + + assert!( + capabilities + .into_iter() + .all(|capability| combined.contains(capability)) + ); + assert_eq!(combined.0.count_ones(), capabilities.len() as u32); + } + #[tokio::test] async fn rejects_zero_plugin_host_deadlines() { let install = complete_builder()