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/Cargo.lock b/Cargo.lock index 66dfc9507..79ad66dc2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4313,6 +4313,7 @@ dependencies = [ "async-lock", "async-trait", "base64", + "bon", "buffa", "bytes", "cbc 0.2.1", @@ -4377,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 102b0cdc4..be907c3f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,11 +9,16 @@ 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 = [ ".", "examples/voip-cli", "http_clients/ureq-client", + "plugins/metrics", "storages/chat-store", "storages/sqlite-storage", "tests/bench-integration", @@ -49,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" @@ -112,6 +120,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 = ["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. @@ -170,6 +184,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/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 987e23a97..f0030c034 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, 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 new file mode 100644 index 000000000..68324e12e --- /dev/null +++ b/agent_docs/plugin_architecture.md @@ -0,0 +1,318 @@ +# 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. + +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 +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. 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 +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. + +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. +`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 +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. + +`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 +`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 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 receives cancellation only on rollback or +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 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 +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::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 +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 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. + +## Future foreign-language adapter seam + +A future bridge should be a Rust adapter at the host boundary, not a second +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; +- 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. 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/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..3f11f32ad --- /dev/null +++ b/plugins/metrics/src/lib.rs @@ -0,0 +1,494 @@ +//! 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, PluginCoreEventSubscription, + 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, + _core_events: PluginCoreEventSubscription, +} + +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"))?; + + let core_events = 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), + _core_events: core_events, + }); + 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, + PluginHealth, 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); + 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); + assert!(matches!( + router.subscribe( + [selector], + PluginEventEndpointConfig::new(1, PluginEventOverflow::DropNewest), + ), + Err(PluginEventSubscribeError::Closed) + )); + } +} diff --git a/src/bot.rs b/src/bot.rs index 29709ceae..1fa37cedf 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -1,6 +1,8 @@ use crate::cache_config::CacheConfig; -use crate::client::Client; +use crate::client::{Client, ClientBuilderError}; use crate::pair_code::PairCodeOptions; +#[cfg(feature = "plugins")] +use crate::plugins::{ClientPlugin, PluginHostConfig, PluginRegistration, UntypedClientPlugin}; use crate::store::commands::DeviceCommand; use crate::store::error::StoreError; use crate::store::persistence_manager::PersistenceManager; @@ -87,49 +89,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 @@ -585,60 +546,22 @@ 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() { 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 @@ -711,6 +634,10 @@ pub struct BotBuilder< resend_rate_limit: Option<(u32, u32)>, task_instrument: Option>, alloc_meter: Option>, + #[cfg(feature = "plugins")] + plugins: Vec, + #[cfg(feature = "plugins")] + plugin_host_config: PluginHostConfig, _marker: PhantomData<(B, T, H, R)>, } @@ -736,6 +663,10 @@ 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, + #[cfg(feature = "plugins")] + plugin_host_config: self.plugin_host_config, _marker: PhantomData, } } @@ -884,6 +819,50 @@ impl BotBuilder { self } + /// 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 + } + + /// 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 + } + + /// 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 + } + + /// 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) ────────────────────────────────────────────────────── @@ -1278,18 +1257,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 +1300,42 @@ 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 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)); + #[cfg(feature = "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 { + 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, @@ -1455,6 +1409,44 @@ mod tests { .client() } + #[cfg(feature = "plugins")] + struct BotBuilderPlugin; + + #[cfg(feature = "plugins")] + 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")) }) + } + } + + #[cfg(feature = "plugins")] + #[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 93778357b..33584c98e 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,9 +1,12 @@ mod accessors; mod adapters; mod app_state; +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; mod lifecycle; @@ -13,6 +16,13 @@ pub(crate) mod offline_resume; mod sender_keys; 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")] +#[cfg_attr(docsrs, doc(cfg(feature = "client-lifecycle")))] +pub use extension_lifecycle::{ClientLifecycle, ConnectionScope, ConnectionScopeState}; pub use voip::{CallError, Voip}; use crate::cache::Cache; @@ -51,6 +61,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. @@ -267,15 +296,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), @@ -294,7 +339,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 } } @@ -376,6 +424,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)?; @@ -634,6 +704,8 @@ 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, /// Whether the noise socket is established (connected to WhatsApp servers). @@ -662,6 +734,12 @@ 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. + #[cfg(feature = "client-lifecycle")] + 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 /// watchdog reads its activity timestamps. Snapshot via [`Client::stats`]. @@ -1012,9 +1090,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 46426ac13..ec168b83f 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -27,17 +27,41 @@ 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) + } + + /// 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), + } } - /// 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); + 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. @@ -169,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, @@ -200,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()), } @@ -516,6 +596,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/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 new file mode 100644 index 000000000..fd63f0541 --- /dev/null +++ b/src/client/builder.rs @@ -0,0 +1,1269 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use thiserror::Error; + +use super::Client; +#[cfg(feature = "client-lifecycle")] +use super::{ClientLifecycle, LifecycleRegistration}; +use crate::cache_config::CacheConfig; +use crate::http::HttpClient; +#[cfg(feature = "plugins")] +use crate::plugins::{ + ClientPlugin, PluginHost, PluginHostConfig, PluginPlan, PluginPlanError, PluginRegistration, + UntypedClientPlugin, +}; +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`]. +/// +/// 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, +} + +impl ClientBuild { + pub(crate) fn new( + client: Arc, + sync_task_receiver: async_channel::Receiver, + ) -> Self { + Self { + client, + sync_task_receiver, + } + } + + /// 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 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) + } +} + +/// A validated client-construction failure. +#[derive(Debug, Error)] +#[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, + #[error("background saver interval must be greater than zero")] + 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")] + #[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")] + #[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), +} + +/// Runtime-validated, low-level builder for [`Client`]. +/// +/// 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 { + runtime: Option>, + persistence_manager: Option>, + transport_factory: Option>, + 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, + #[cfg(feature = "client-lifecycle")] + lifecycle: Option>, + #[cfg(feature = "plugins")] + plugins: Vec, + #[cfg(feature = "plugins")] + plugin_host_config: PluginHostConfig, +} + +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(), + 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, + #[cfg(feature = "client-lifecycle")] + lifecycle: None, + #[cfg(feature = "plugins")] + plugins: Vec::new(), + #[cfg(feature = "plugins")] + plugin_host_config: PluginHostConfig::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 + } + + /// 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 + } + + /// 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, + { + self.lifecycle = Some(Arc::new(lifecycle)); + self + } + + /// 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 + } + + /// 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 + } + + /// 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 + } + + /// 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 + } + + /// 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, + 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 + } + + #[inline(never)] + fn build_boxed( + self, + ) -> wacore::runtime::BoxFuture<'static, Result> { + Box::pin(async move { + 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)?; + + if self.background_saver_interval == Some(Duration::ZERO) { + 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); + } + #[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?; + } + + self.finish(runtime, persistence_manager, transport_factory, http_client) + .await + }) + } + + pub(crate) async fn build_required( + runtime: Arc, + persistence_manager: Arc, + transport_factory: Arc, + http_client: Arc, + override_version: Option<(u32, u32, u32)>, + cache_config: CacheConfig, + ) -> ClientBuild { + 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}"), + } + } + + async fn finish( + self, + runtime: Arc, + persistence_manager: Arc, + 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) => { + Arc::new(wacore::stats::InstrumentedRuntime::new(runtime, instrument)) + } + None => runtime, + }; + + #[cfg(feature = "client-lifecycle")] + 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(), self.plugin_host_config); + lifecycle_handler = Some(host.clone()); + host + }); + (lifecycle_handler, plugin_host) + }; + #[cfg(feature = "client-lifecycle")] + 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), + transport_factory, + http_client, + 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() { + 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); + } + #[cfg(feature = "client-lifecycle")] + if let Some(lifecycle) = &client.lifecycle + && let Err(error) = lifecycle.install(Arc::downgrade(&client)).await + { + #[cfg(feature = "plugins")] + if client.plugin_host.is_some() { + return Err(ClientBuilderError::PluginInstall(error)); + } + return Err(ClientBuilderError::LifecycleInstall(error)); + } + + 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); + } + #[cfg(feature = "client-lifecycle")] + if let Some(lifecycle) = &client.lifecycle { + #[cfg(feature = "plugins")] + 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" + ))); + } + } + #[cfg(feature = "client-lifecycle")] + construction.disarm(); + Ok(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 { + client: Arc, + 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 { + client, + armed: true, + } + } + + fn disarm(&mut self) { + self.armed = false; + } +} + +#[cfg(feature = "client-lifecycle")] +impl Drop for ClientConstructionGuard { + fn drop(&mut self) { + if self.armed { + self.client.signal_shutdown_sync(); + } + } +} + +#[derive(Default)] +pub(super) struct ClientExtensions { + #[cfg(feature = "client-lifecycle")] + pub(super) lifecycle: Option>, + #[cfg(feature = "plugins")] + pub(super) plugin_host: Option>, +} + +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) + } + + fn client(&self) -> Arc { + Arc::clone(&self.client) + } +} + +#[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; + + #[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")] + 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>( + &'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(()) + }) + } + } + + #[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>( + &'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, + } + + #[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() + } + } + + 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!( + 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(), + ClientExtensions::default(), + ); + assert_eq!(spawns.load(Ordering::SeqCst), 0); + + let build = assembly.start(); + assert_eq!(spawns.load(Ordering::SeqCst), 1); + build.into_client().signal_shutdown_sync(); + } + + #[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); + 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] + #[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() { + 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(); + + assert!(matches!( + 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()) + .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( + 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.into_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(); + } + + #[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) + )); + } + + #[cfg(feature = "client-lifecycle")] + struct PanickingInstallLifecycle { + when_polled: bool, + } + + #[cfg(feature = "client-lifecycle")] + 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] + #[cfg(feature = "client-lifecycle")] + 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] + #[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) + .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..59f573431 --- /dev/null +++ b/src/client/extension_lifecycle.rs @@ -0,0 +1,2379 @@ +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 futures::FutureExt; +use wacore::runtime::{BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, wait_for_shutdown}; + +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_TARGET_CAPACITY: usize = 64; + +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. +#[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 { + pub(crate) 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, + } + } + + /// 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 + } + + fn mark_ready(&self) -> bool { + self.inner + .state + .compare_exchange(SCOPE_OPEN, SCOPE_READY, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + } + + 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( + 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. Connection callbacks are +/// 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 { + 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(()) }) + } + + /// 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(()) }) + } +} + +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, + shutdown_notifier: ShutdownNotifier, + callback_timeout: Duration, + terminal: AtomicBool, + construction_transition: std::sync::Mutex<()>, + construction_state: AtomicU8, + construction_notifier: ShutdownNotifier, +} + +#[derive(Default)] +struct ScopeRegistry { + active: Option, + 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, + overflowed: bool, +} + +impl CallbackQueue { + 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 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 => { + // Closures are lossless, so the target remains soft under backlog. + self.pending.push_back(callback); + Vec::new() + } + } + } + + fn compact_for_shutdown(&mut self) -> Vec { + if !self.overflowed { + return Vec::new(); + } + + 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 + } +} + +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); + 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)) +} + +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) + } + + pub(super) fn new_with_timeout( + handler: Arc, + runtime: Arc, + callback_timeout: Duration, + ) -> Self { + 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), + 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(), + } + } + + 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 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 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 + } + 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 + { + return Err(anyhow::anyhow!( + "client shutdown began during lifecycle installation" + )); + } + result + } + + 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, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => self.construction_notifier.notify(), + Err(CONSTRUCTION_ACTIVE) => {} + Err(_) => return false, + } + true + } + + 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( + &self, + generation: u64, + is_current: impl FnOnce() -> bool, + ) -> bool { + if self.terminal.load(Ordering::Acquire) { + return false; + } + + let scope = ConnectionScope::new(generation); + 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 { + log::warn!( + "Replacing unclosed connection scope for generation {}", + replaced.generation() + ); + replaced.cancel(); + scopes.retired.push(replaced); + } + true + } + + pub(super) async fn ready(self: &Arc, generation: u64) -> bool { + if self.terminal.load(Ordering::Acquire) { + 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 + }; + + done_rx.recv().await.unwrap_or(false) && !scope.is_cancelled() + } + + 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_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); + } + } + + 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, || {}); + } + + /// `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(); + let scope = if scopes + .active + .as_ref() + .is_some_and(|scope| scope.generation() == generation) + { + scopes.active.take() + } else { + scopes + .retired + .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(); + let mut dropped = Vec::new(); + if queue.shutdown_requested { + 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_with_pressure_policy(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); + } + + pub(super) async fn shutdown(self: &Arc) { + self.request_shutdown(); + + 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; + } + + 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 = { + 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.cancel_all_scopes() + } + 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, dropped) = { + let mut queue = self.callback_queue(); + if queue.shutdown_requested || self.terminal.load(Ordering::Acquire) { + (false, vec![callback]) + } else { + let dropped = queue.push_with_pressure_policy(callback); + let should_spawn = !queue.drain_scheduled; + queue.drain_scheduled = true; + (should_spawn, dropped) + } + }; + warn_dropped_callbacks(dropped); + 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() + }; + let (should_spawn, dropped) = { + let mut queue = self.callback_queue(); + if !queue.shutdown_requested || queue.shutdown_enqueued { + return; + } + let dropped = queue.compact_for_shutdown(); + 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); + } + + 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 callback = { + let mut queue = self.callback_queue(); + match queue.pending.pop_front() { + Some(callback) => callback, + None => { + queue.drain_scheduled = false; + queue.overflowed = false; + return; + } + } + }; + + match callback { + LifecycleCallback::Ready { scope, done } => { + 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", move |handler| handler.on_closed(scope)) + .await; + } + LifecycleCallback::Shutdown => { + self.run_callback("shutdown", |handler| handler.shutdown()) + .await; + self.shutdown_complete.store(true, Ordering::Release); + self.shutdown_notifier.notify(); + } + } + } + } + + async fn run_callback<'a>( + &'a self, + name: &'static str, + 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 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); + drop(callback); + })) + .is_err(); + if drop_panicked { + log::warn!("Client lifecycle {name} panicked while being dropped"); + return; + } + match result { + 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"), + } + } + + 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()) + } + + 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 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(); + } + } + + fn cancel_all_scopes(&self) { + let scopes = self.scopes(); + if let Some(scope) = &scopes.active { + scope.cancel(); + } + for scope in &scopes.retired { + scope.cancel(); + } + } + + fn scope_for(&self, generation: u64) -> Option { + let scopes = self.scopes(); + scopes + .active + .as_ref() + .filter(|scope| scope.generation() == generation) + .or_else(|| { + scopes + .retired + .iter() + .find(|scope| scope.generation() == generation) + }) + .cloned() + } +} + +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; + + 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 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<()>, + scope: std::sync::Mutex>, + events: std::sync::Mutex>, + } + + #[derive(Default)] + struct ReentrantDisconnectLifecycle { + client: std::sync::Mutex>>, + 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 { + *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, + } + + 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 { + self.calls.fetch_add(1, Ordering::SeqCst); + let _ = self.started.try_send(()); + let _ = self.release.recv().await; + self.completed.store(true, Ordering::Release); + Ok(()) + }) + } + } + + #[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, + closed_calls: AtomicUsize, + 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); + 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"); + } + } + + 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, + } + + 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 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)); + } + + #[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( + 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( + 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.into_client(); + const GENERATION: u64 = 9; + 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(GENERATION).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); + 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 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 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( + 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") + .into_client(); + const GENERATION: u64 = 13; + client + .connection_generation + .store(GENERATION, Ordering::SeqCst); + 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 { + ready_client.dispatch_connected(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"); + + 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"); + 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"); + client.shutdown_lifecycle().await; + 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 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") + .into_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(GENERATION), + ) + .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 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(GENERATION), + ) + .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); + 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"] + ); + } + + #[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); + 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"); + + 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(), + usize::try_from(closed_callbacks + 1).expect("callback count fits usize") + ); + + 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(), + usize::try_from(closed_callbacks + 1) + .expect("terminal callback count fits usize") + ); + 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), + usize::try_from(closed_callbacks).expect("closure count fits usize") + ); + 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); + 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 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] + 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()); + 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") + .into_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(GENERATION).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 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 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 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(GENERATION).await; + let scope = registration + .scope_for(GENERATION) + .expect("connection scope"); + *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()); + 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(); + } + + #[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 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( + 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()); + let registration = Arc::new(LifecycleRegistration::new( + lifecycle.clone(), + Arc::new(TokioRuntime), + )); + + assert!(registration.begin_scope_if_current(31, || true)); + assert!(registration.ready(31).await); + let first = registration.scope_for(31).expect("first scope"); + + 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); + assert_eq!(first.state(), ConnectionScopeState::Closed); + assert!(registration.scope_for(31).is_none()); + assert!(registration.scope_for(32).is_some()); + + registration.close_scope(32); + 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()); + let registration = Arc::new(LifecycleRegistration::new( + lifecycle.clone(), + Arc::new(TokioRuntime), + )); + const GENERATION: u64 = 21; + assert!(registration.begin_scope_if_current(GENERATION, || true)); + 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); + shutdown.await.expect("shutdown did not panic"); + assert_eq!( + lifecycle.events(), + vec!["closed:21".to_string(), "shutdown".to_string()] + ); + + 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; + 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") + .into_client(); + client + .subscribe_handler(Arc::new(LogoutOrderHandler { + lifecycle: lifecycle.clone(), + })) + .detach(); + + 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 424652858..6bdbb826b 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -6,11 +6,22 @@ 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. 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() } @@ -25,6 +36,10 @@ 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(); + } self.notify_connection_shutdown(); } @@ -66,8 +81,54 @@ impl Client { self.is_connected() && self.is_logged_in() && self.is_ready.load(Ordering::Relaxed) } - /// Dispatch the Connected event and notify waiters. - pub(crate) 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")] + { + if let Some(lifecycle) = &self.lifecycle { + if !lifecycle.ready(expected_generation).await { + debug!( + "Skipping Connected dispatch for retired generation {expected_generation}" + ); + return; + } + + // 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(expected_generation, || self.publish_connected()) { + debug!("Skipping Connected dispatch after lifecycle cancellation"); + } + return; + } + } + + #[cfg(feature = "client-lifecycle")] + 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(); + } + + fn publish_connected(&self) { self.is_ready.store(true, Ordering::Relaxed); wacore::telemetry::set_connected(true); self.core.event_bus.dispatch(Event::Connected( @@ -76,6 +137,20 @@ 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(); + } + } + /// Create a new `Client` with default cache configuration. /// /// This is the standard constructor. Use [`Client::new_with_cache_config`] @@ -87,7 +162,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, @@ -96,6 +171,7 @@ impl Client { CacheConfig::default(), ) .await + .into_parts() } /// Create a new `Client` with a custom [`CacheConfig`]. @@ -107,6 +183,33 @@ 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, + ) + .await + .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, + extensions: ClientExtensions, + ) -> ClientAssembly { + let ClientExtensions { + #[cfg(feature = "client-lifecycle")] + lifecycle, + #[cfg(feature = "plugins")] + plugin_host, + } = extensions; let mut unique_id_bytes = [0u8; 2]; rand::make_rng::().fill_bytes(&mut unique_id_bytes); @@ -126,6 +229,8 @@ 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)), is_connected: Arc::new(AtomicBool::new(false)), @@ -133,6 +238,10 @@ 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, stats: Arc::new(wacore::stats::SessionStats::new()), transport: Arc::new(Mutex::new(None)), @@ -287,7 +396,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")] @@ -303,17 +412,18 @@ 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}"); } })) .detach(); - - (arc, rx) } // Deliberately NOT instrumented: this span would live for the entire client @@ -321,10 +431,26 @@ 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 + { + 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 @@ -436,6 +562,8 @@ impl Client { ); self.runtime.sleep(delay).await; } + #[cfg(feature = "client-lifecycle")] + self.shutdown_lifecycle().await; info!("Client run loop has shut down."); } @@ -443,6 +571,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 } @@ -598,8 +732,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) @@ -607,6 +739,8 @@ impl Client { .build(), )); + self.disconnect().await; + Ok(()) } @@ -620,6 +754,8 @@ 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 // closing it, so a disconnect mid-offline-sync still acks the @@ -666,6 +802,8 @@ 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; } /// Backoff step used by [`reconnect()`] to create an offline window. @@ -694,6 +832,10 @@ 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(); + } wacore::telemetry::reconnect(); self.intentional_reconnect.store(true, Ordering::Relaxed); self.auto_reconnect_errors @@ -730,6 +872,10 @@ 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(); + } self.expected_disconnect.store(true, Ordering::Relaxed); // Same durable-before-receipts gate as disconnect(). @@ -754,7 +900,46 @@ impl Client { feature = "tracing", tracing::instrument(name = "wa.conn.cleanup", level = "debug", skip_all) )] - pub(crate) async fn cleanup_connection_state(&self) { + #[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; + return; + } + + // Scope closure must survive a caller dropping its cleanup waiter. + let (completed, completion) = futures::channel::oneshot::channel(); + let client = Arc::clone(self); + self.runtime + .spawn(Box::pin(async move { + let result = std::panic::AssertUnwindSafe(client.cleanup_connection_state_inner()) + .catch_unwind() + .await; + let _ = completed.send(result); + })) + .detach(); + 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) { + #[cfg(feature = "client-lifecycle")] + 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 @@ -763,7 +948,28 @@ 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")] + 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); + } + 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. @@ -774,6 +980,8 @@ 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 // the socket teardown below): a concurrent accept()/call() setup that finishes its async work @@ -789,11 +997,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()` @@ -923,6 +1126,8 @@ impl Client { if let Some(proc) = self.app_state_processor.lock().await.as_ref() { proc.clear_key_cache().await; } + #[cfg(feature = "client-lifecycle")] + drop(scope_close); } /// Waits for the noise socket to be established. diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 4db844432..c2dbb3bfa 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))); @@ -707,6 +707,11 @@ 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() + .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. @@ -725,6 +730,20 @@ 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 + && !self.expected_disconnect.load(Ordering::Acquire) + }); + if !opened { + self.is_logged_in.store(false, Ordering::SeqCst); + debug!("Ignoring stanza retired during lifecycle publication"); + return; + } + } + #[cfg(feature = "client-lifecycle")] + drop(login_transition); info!( "Successfully authenticated with WhatsApp servers! (gen={})", @@ -1132,7 +1151,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(task_generation).await; } Err(e) => { client_clone.log_sync_error("critical app state sync", &e); @@ -1194,7 +1213,7 @@ impl Client { // for an outdated connection that was replaced mid-await. check_generation!(); - client_clone.dispatch_connected(); + client_clone.dispatch_connected(task_generation).await; } })).detach(); } diff --git a/src/client/tests.rs b/src/client/tests.rs index 6266bc0bc..d10555c2d 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/ib.rs b/src/handlers/ib.rs index 7bd366fcd..a1eaee330 100644 --- a/src/handlers/ib.rs +++ b/src/handlers/ib.rs @@ -1,6 +1,6 @@ use super::traits::StanzaHandler; use crate::client::Client; -use crate::types::events::{DirtyState, Event, OfflineSyncPreview}; +use crate::types::events::{DirtyState, Event, EventKind, OfflineSyncPreview}; use async_trait::async_trait; use futures::FutureExt; use log::{debug, info, warn}; @@ -68,12 +68,14 @@ async fn handle_ib_impl(client: Arc, node: &wacore_binary::NodeRef<'_>) ); let needs_resync = bit.dirty_type == DirtyType::SyncdAppState; - client.core.event_bus.dispatch(Event::DirtyState( - DirtyState::builder() - .dirty_type(bit.dirty_type.clone()) - .maybe_timestamp(bit.timestamp) - .build(), - )); + if client.core.event_bus.has_handler_for(EventKind::DirtyState) { + client.core.event_bus.dispatch(Event::DirtyState( + DirtyState::builder() + .dirty_type(bit.dirty_type.clone()) + .maybe_timestamp(bit.timestamp) + .build(), + )); + } debug!( "Received dirty state notification for type: '{dirty_type_str}'. Sending clean IQ." @@ -234,7 +236,7 @@ mod tests { async fn valid_dirty_marker_dispatches_typed_event() { let client = create_test_client().await; let collector = Arc::new(TestEventCollector::default()); - client.register_handler(collector.clone()); + let _subscription = client.subscribe_handler(collector.clone()); let node = NodeBuilder::new("ib") .children([NodeBuilder::new("dirty") .attr("type", "account_sync") diff --git a/src/handlers/notification/mod.rs b/src/handlers/notification/mod.rs index 6aa3f9fbd..579352763 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 { @@ -939,7 +939,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 @@ -968,7 +968,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") @@ -988,7 +988,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; @@ -1015,7 +1015,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 @@ -1037,7 +1037,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; @@ -1148,7 +1148,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). @@ -1186,7 +1186,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(); @@ -1274,7 +1274,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"; @@ -1332,7 +1332,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 3e733ccb6..968c0f6b6 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). @@ -84,7 +85,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 +94,10 @@ pub use client::{ StatsSnapshot, StorageResourceReport, TransportResourceReport, }; 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; pub mod download; @@ -108,6 +112,23 @@ pub(crate) mod msg_secret_buffer; 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, PluginCoreEventSubscription, PluginCoreEvents, + PluginEventEndpointConfig, PluginEventEndpointStats, PluginEventEnvelope, PluginEventOverflow, + PluginEventPayloadEncoding, PluginEventPublishError, PluginEventPublishReport, + PluginEventPublisherStats, PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, + PluginEventRouterStats, PluginEventSelector, PluginEventSubscribeError, + PluginEventSubscription, PluginEventTopic, PluginEventTryReceiveError, PluginEvents, + PluginFuture, PluginHealth, PluginHostConfig, PluginHostStats, PluginIq, PluginIqError, + PluginManifest, PluginMessaging, PluginMessagingError, PluginPlanError, PluginResourceError, + PluginState, PluginStats, PluginTasks, UntypedClientPlugin, +}; pub mod request; pub(crate) mod signal_flush; pub use request::IqError; @@ -177,7 +198,19 @@ 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, 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, + PluginCoreEventSubscription, PluginEventEndpointConfig, PluginEventOverflow, + PluginEventPayloadEncoding, PluginEventRouter, PluginEventSelector, + PluginEventSubscription, PluginEventTopic, PluginEvents, PluginFuture, PluginHostConfig, + PluginManifest, UntypedClientPlugin, + }; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] pub use crate::runtime_impl::TokioRuntime; @@ -186,7 +219,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 f1000bb3e..1d70f77df 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() @@ -5016,7 +5016,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"; @@ -9533,7 +9533,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 @@ -9597,7 +9597,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"; @@ -9690,7 +9690,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"; @@ -9770,7 +9770,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"; @@ -9898,7 +9898,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"; @@ -9974,7 +9974,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"; @@ -10090,7 +10090,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"; @@ -10239,7 +10239,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"; @@ -10311,7 +10311,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"; @@ -10397,7 +10397,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"; @@ -11077,7 +11077,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"; @@ -11164,7 +11164,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"; @@ -11269,7 +11269,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"; @@ -11391,7 +11391,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"; @@ -11479,7 +11479,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"; @@ -11677,7 +11677,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/plugins/events.rs b/src/plugins/events.rs new file mode 100644 index 000000000..f70408dad --- /dev/null +++ b/src/plugins/events.rs @@ -0,0 +1,1261 @@ +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. +/// +/// `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 { + 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, +} + +/// 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. +/// `dropped` follows the discarded envelope, including cross-namespace `DropOldest` eviction. +#[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")] +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 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.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>, + 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(dropped)) => { + self.dropped.fetch_add(1, Ordering::Relaxed); + dropped.publisher.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 let Some(evicted) = evicted { + self.dropped.fetch_add(1, Ordering::Relaxed); + evicted.publisher.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 { + publishers: HashMap, Arc>, + queue_memory: Arc, + 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 { + let publishers = plugin_ids + .into_iter() + .map(|plugin_id| (Arc::from(plugin_id), Arc::new(PublisherCounters::default()))) + .collect(); + Self { + inner: Arc::new(PluginEventRouterInner { + publishers, + queue_memory: Arc::new(QueueMemory::default()), + 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()) + } + + /// 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, + 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.publishers.contains_key(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, + publisher: Arc, + plugin_id: &Arc, + topic: &PluginEventTopic, + publication: PluginEventPublication, + ) -> Result { + 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) { + 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 envelope = Arc::new( + PluginEventEnvelope::builder() + .plugin_id(plugin_id.clone()) + .topic(topic.clone()) + .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), + ..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); + event.publisher.delivered.fetch_add(1, Ordering::Relaxed); + Ok(event.envelope.clone()) + } + + 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); + event.publisher.delivered.fetch_add(1, Ordering::Relaxed); + Ok(event.envelope.clone()) + } +} + +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, + stats: Arc, + 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 stats(&self) -> PluginEventPublisherStats { + self.stats.snapshot() + } + + pub fn publish( + &self, + topic: &PluginEventTopic, + schema_version: u32, + payload_encoding: PluginEventPayloadEncoding, + payload: impl Into, + ) -> Result { + 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, + PluginEventPublication { + schema_version, + payload_encoding, + payload: payload.into(), + connection_generation: self.connection_generation.load(Ordering::Acquire), + }, + ) + } +} + +pub(super) fn publisher( + plugin_id: &str, + router: PluginEventRouter, + resources: Arc, + connection_generation: Arc, +) -> 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 { + 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 { + let publisher = router + .inner + .publishers + .get(plugin_id) + .cloned() + .expect("registered publisher"); + router + .publish( + publisher, + &Arc::from(plugin_id), + topic, + 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()]); + 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, + } + ); + assert_eq!( + router.publisher_stats("metrics"), + Some(PluginEventPublisherStats { + published: 4, + matched: 4, + enqueued: 2, + delivered: 2, + dropped: 2, + ..PluginEventPublisherStats::default() + }) + ); + } + + #[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); + assert_eq!(router.stats().dropped, 2); + 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()]); + 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 new file mode 100644 index 000000000..0adb4f696 --- /dev/null +++ b/src/plugins/mod.rs @@ -0,0 +1,5649 @@ +//! Build-time client plugins and their capability-scoped host. + +mod events; + +pub use events::{ + PluginEventEndpointConfig, PluginEventEndpointStats, PluginEventEnvelope, PluginEventOverflow, + PluginEventPayloadEncoding, PluginEventPublishError, PluginEventPublishReport, + PluginEventPublisherStats, PluginEventReceiveError, PluginEventRouteError, PluginEventRouter, + PluginEventRouterStats, PluginEventSelector, PluginEventSubscribeError, + PluginEventSubscription, PluginEventTopic, PluginEventTryReceiveError, PluginEvents, +}; + +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 std::time::Duration; + +use futures::FutureExt; +use portable_atomic::AtomicU64; +use thiserror::Error; +use wacore::iq::spec::IqSpec; +use wacore::runtime::{ + BoxFuture, Runtime, ShutdownNotifier, ShutdownSignal, Spawnable, timeout as runtime_timeout, + wait_for_shutdown, +}; +use wacore::sync_marker::MaybeSendSync; +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, RawNodeLease}; +use crate::request::IqError; +use crate::send::{SendError, SendResult}; + +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); + +/// 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, + PluginEvents, +} + +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", + Self::PluginEvents => "events.plugin.publish", + } + } + + const fn bit(self) -> u64 { + match self { + Self::CoreEvents => CAP_CORE_EVENTS, + Self::Tasks => CAP_TASKS, + Self::Messaging => CAP_MESSAGING, + Self::Iq => CAP_IQ, + Self::PluginEvents => CAP_PLUGIN_EVENTS, + } + } +} + +/// Compact set of capabilities requested by one plugin. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct PluginCapabilities(u64); + +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 + } +} + +/// 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, +} + +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; + 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 install_timeout(self) -> Duration { + self.install_timeout + } + + 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] +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 + } +} + +/// 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) -> 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(()) }) + } + + /// Release plugin-owned state. This may run after `install` began but returned an error. + fn shutdown(&self) -> PluginFuture<'_, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } +} + +/// 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] +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, + #[error("the plugin task capacity is exhausted")] + TaskCapacityExceeded, +} + +#[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), +} + +/// 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, + /// 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. + 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 upstream_callback_failures: u64, + pub upstream_callback_timeouts: u64, + pub plugins: Vec, + pub event_router: Option, +} + +struct PluginResources { + active: AtomicBool, + closed: AtomicBool, + activation: ShutdownNotifier, + shutdown: ShutdownNotifier, + install_tasks: Arc, + connection_tasks: Mutex, + subscriptions: Mutex>>, + teardown_panics: AtomicU64, +} + +#[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() + } + + fn active(&self) -> usize { + self.state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .active + } +} + +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 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) { + drop(state); + drop(acquired_raw_node_lease); + self.close(); + 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 resources = self.resources.upgrade(); + 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) = &resources { + resources.teardown_panics.fetch_add(1, Ordering::Relaxed); + } + log::warn!( + "Plugin `{}` core-event subscription panicked while closing", + self.plugin_id + ); + } + if let Some(resources) = resources { + resources.forget_subscription(self); + } + 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 { + fn new() -> Arc { + Arc::new(Self { + active: AtomicBool::new(false), + 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()), + teardown_panics: AtomicU64::new(0), + }) + } + + #[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); + } + + fn publish_activation(&self) { + if self.active.load(Ordering::Acquire) && !self.closed.load(Ordering::Acquire) { + 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, + client: Weak, + resources: Weak, + plugin_id: Arc, + interest: EventInterest, + subscription: Subscription, + raw_node_lease: Option, + ) -> 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) { + true + } else { + subscriptions.retain(|subscription| { + subscription + .upgrade() + .is_some_and(|subscription| subscription.is_active()) + }); + subscriptions.push(Arc::downgrade(®istration)); + false + } + }; + if rejected { + registration.close(); + Err(PluginResourceError::ShuttingDown) + } else { + Ok(PluginCoreEventSubscription { + inner: registration, + }) + } + } + + 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 + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if registry.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) + } + } + } + + 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 = self + .connection_tasks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .trackers + .get(&generation) + .cloned() + .unwrap_or_else(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 + .subscriptions + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + std::mem::take(&mut *subscriptions) + }; + for subscription in subscriptions { + if let Some(subscription) = subscription.upgrade() { + subscription.close(); + } + } + } +} + +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()) + .iter() + .filter_map(Weak::upgrade) + .filter(|subscription| subscription.is_active()) + .count(), + 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, + task_panics: 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), + task_panics: 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 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) { + 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 + || task_panics > 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, + task_panics, + 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)] +pub struct PluginTasks { + runtime: Arc, + resources: Arc, + diagnostics: Arc, + plugin_id: Arc, +} + +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, + { + if self.resources.closed.load(Ordering::Acquire) { + return Err(PluginResourceError::ShuttingDown); + } + let lease = self.resources.install_tasks.register()?; + spawn_after_activation( + &self.runtime, + Arc::clone(&self.resources), + Arc::clone(&self.diagnostics), + Arc::clone(&self.plugin_id), + lease, + future, + shutdown, + ); + Ok(()) + } + + 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. +/// Handlers run inline and must hand slow work to a task capability. +#[derive(Clone)] +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(resources) = self.resources.upgrade() else { + return; + }; + if resources.ensure_active().is_err() { + return; + } + 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 { + pub fn subscribe( + &self, + interest: EventInterest, + handler: Arc, + ) -> Result { + let client = self + .client + .upgrade() + .ok_or(PluginResourceError::ClientUnavailable)?; + 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( + self.client.clone(), + Arc::downgrade(&self.resources), + Arc::clone(&self.plugin_id), + interest, + subscription, + raw_node_lease, + ) + } +} + +/// 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, + dependencies: HashMap, + core_events: Option, + tasks: Option, + messaging: Option, + iq: Option, + plugin_events: Option, +} + +impl PluginContext { + pub fn plugin_id(&self) -> &str { + &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> { + let api = self.dependencies.get(&TypeId::of::

())?.upgrade()?; + downcast_api::(&api) + } + + 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() + } + + pub fn plugin_events(&self) -> Option<&PluginEvents> { + self.plugin_events.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 cancellation is signalled synchronously when its generation retires. +#[derive(Clone)] +pub struct PluginConnectionTasks { + runtime: Arc, + scope: ConnectionScope, + tracker: Arc, + diagnostics: Arc, + plugin_id: Arc, +} + +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, + { + if self.scope.is_cancelled() { + return Err(PluginResourceError::ShuttingDown); + } + let lease = self.tracker.register()?; + spawn_until_cancelled( + &self.runtime, + self.scope.cancellation_signal(), + Arc::clone(&self.diagnostics), + 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() { + 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(()), + } + } +} + +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"); + } + } +} + +#[derive(Clone, Copy)] +enum PluginTaskShutdown { + Abort, + Cooperative, +} + +fn spawn_until_cancelled( + runtime: &Arc, + cancellation: ShutdownSignal, + diagnostics: Arc, + plugin_id: Arc, + lease: TaskLease, + future: F, + shutdown: PluginTaskShutdown, +) where + F: Future + Spawnable, +{ + let work = GuardedPluginTask::new(future, diagnostics, plugin_id); + runtime + .spawn(Box::pin(async move { + let _lease = lease; + let work = Box::pin(work); + match shutdown { + PluginTaskShutdown::Abort => { + let cancelled = Box::pin(wait_for_shutdown(&cancellation)); + let _ = futures::future::select(cancelled, work).await; + } + PluginTaskShutdown::Cooperative => work.await, + } + })) + .detach(); +} + +fn spawn_after_activation( + runtime: &Arc, + resources: Arc, + diagnostics: Arc, + plugin_id: Arc, + lease: TaskLease, + future: F, + shutdown: PluginTaskShutdown, +) where + F: Future + Spawnable, +{ + 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; + 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 work = Box::pin(work); + match shutdown { + PluginTaskShutdown::Abort => { + let cancelled = Box::pin(wait_for_shutdown(&cancellation)); + let _ = futures::future::select(cancelled, work).await; + } + PluginTaskShutdown::Cooperative => 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; +type WeakErasedApi = Weak; + +#[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 dependency_view(&self, markers: &[TypeId]) -> HashMap { + let values = self + .values + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + markers + .iter() + .filter_map(|marker| values.get(marker).map(|api| (*marker, Arc::downgrade(api)))) + .collect() + } +} + +fn downcast_api(api: &ErasedApi) -> Option> { + api.as_any() + .downcast_ref::>() + .map(|typed| typed.0.clone()) +} + +trait ErasedClientPlugin: MaybeSendSync { + 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 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) -> Option { + Some(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(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) + }) + } + + 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)), + } + } + + 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 { + 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; + if let Some(marker) = plugin.marker_type_id() + && !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); + 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) { + 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, + diagnostics: Arc, +} + +struct PluginInstallRollback { + runtime: Arc, + config: PluginHostConfig, + installed: Vec, + current: Option, + upstream: Option>, + staged_apis: Option>, + armed: bool, +} + +impl PluginInstallRollback { + fn new(runtime: Arc, capacity: usize, config: PluginHostConfig) -> Self { + Self { + runtime, + config, + installed: Vec::with_capacity(capacity), + current: None, + upstream: None, + staged_apis: None, + armed: true, + } + } + + fn close_resources(&self) { + if let Some(current) = &self.current { + close_plugin_resources(¤t.manifest.id, ¤t.resources); + } + for plugin in self.installed.iter().rev() { + close_plugin_resources(&plugin.manifest.id, &plugin.resources); + } + } + + 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(); + let staged_apis = self.staged_apis.take(); + self.armed = false; + 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(); + let config = self.config; + runtime + .spawn(Box::pin(async move { + let result = AssertUnwindSafe(shutdown_staged_plugins( + cleanup_runtime, + config, + current, + installed, + upstream, + staged_apis, + )) + .catch_unwind() + .await; + completed.notify(); + if result.is_err() { + log::warn!("Plugin installation rollback panicked"); + } + })) + .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 { + 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; + self.staged_apis = None; + } +} + +impl Drop for PluginInstallRollback { + fn drop(&mut self) { + let _ = self.schedule_rollback(); + } +} + +struct PluginContextParts { + resources: Arc, + apis: Arc, + runtime: Arc, + connection_generation: Arc, + diagnostics: Arc, +} + +struct InstalledPlugins { + plugins: Vec, + staged_apis: Mutex>>, +} + +pub(crate) struct PluginHost { + ordered: Vec, + manifests: Vec, + diagnostics: Vec>, + upstream: Option>, + installed: OnceLock, + apis: OnceLock>, + runtime: OnceLock>, + event_router: Option, + config: PluginHostConfig, + terminal: AtomicBool, + terminal_notifier: ShutdownNotifier, + installing_resources: Mutex>>, + upstream_callback_failures: AtomicU64, + upstream_callback_timeouts: AtomicU64, +} + +impl PluginHost { + 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 + .iter() + .map(|plugin| plugin.manifest.clone()) + .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)); + let diagnostics = (0..manifests.len()) + .map(|_| PluginDiagnostics::new()) + .collect(); + Arc::new(Self { + ordered: plan.ordered, + manifests, + diagnostics, + upstream, + installed: OnceLock::new(), + apis: OnceLock::new(), + runtime: OnceLock::new(), + event_router, + config, + 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), + }) + } + + pub(crate) fn plugin(&self) -> Option> { + 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 + } + + 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() + .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 upstream_callback_failures > 0 + || upstream_callback_timeouts > 0 + || plugins + .iter() + .any(|plugin| plugin.health == PluginHealth::Degraded) + { + PluginHealth::Degraded + } else { + PluginHealth::Healthy + }; + PluginHostStats { + terminal, + health, + upstream_callback_failures, + upstream_callback_timeouts, + 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 + .ordered + .iter() + .filter(|plugin| { + plugin + .manifest + .capabilities + .contains(PluginCapability::Tasks) + }) + .count(); + 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)) + } + + fn context( + &self, + client: &Weak, + planned: &PlannedPlugin, + parts: PluginContextParts, + ) -> PluginContext { + let PluginContextParts { + resources, + apis, + runtime, + connection_generation, + diagnostics, + } = 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), + core_events: capabilities + .contains(PluginCapability::CoreEvents) + .then(|| PluginCoreEvents { + client: client.clone(), + resources: Arc::clone(&resources), + plugin_id: Arc::clone(&plugin_id), + diagnostics: Arc::clone(&diagnostics), + }), + tasks: capabilities + .contains(PluginCapability::Tasks) + .then(|| PluginTasks { + runtime: Arc::clone(&runtime), + resources: Arc::clone(&resources), + diagnostics: Arc::clone(&diagnostics), + plugin_id, + }), + 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), + }), + plugin_events: self + .event_router + .as_ref() + .filter(|_| capabilities.contains(PluginCapability::PluginEvents)) + .and_then(|router| { + events::publisher( + &manifest.id, + router.clone(), + Arc::clone(&resources), + connection_generation, + ) + }), + } + } + + fn connection_scope( + &self, + scope: ConnectionScope, + plugin: &InstalledPlugin, + task_tracker: Option>, + ) -> PluginConnectionScope { + let tasks = if plugin + .manifest + .capabilities + .contains(PluginCapability::Tasks) + { + self.runtime + .get() + .cloned() + .zip(task_tracker) + .map(|(runtime, tracker)| PluginConnectionTasks { + runtime, + scope: scope.clone(), + tracker, + diagnostics: Arc::clone(&plugin.diagnostics), + plugin_id: Arc::from(plugin.manifest.id.as_str()), + }) + } else { + None + }; + PluginConnectionScope { scope, tasks } + } + + async fn wait_for_tasks( + &self, + completion_signals: Vec, + ) -> Result<(), PluginTaskDrainError> { + let runtime = self + .runtime + .get() + .ok_or(PluginTaskDrainError::RuntimeUnavailable)?; + wait_for_plugin_tasks( + &**runtime, + self.config.task_drain_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"); + }; + let runtime = strong_client.runtime.clone(); + let connection_generation = strong_client.connection_generation.clone(); + drop(strong_client); + self.runtime + .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.config); + self.abort_install_if_terminal(&mut rollback).await?; + if let Some(upstream) = &self.upstream { + 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); + } + self.abort_install_if_terminal(&mut rollback).await?; + } + + let staging = Arc::new(ApiRegistry::default()); + rollback.staged_apis = Some(Arc::clone(&staging)); + 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, + 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; + 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(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() { + 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; + anyhow::bail!( + "plugin `{}` installation failed: {error:#}", + planned.manifest.id + ); + } + }; + 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; + anyhow::bail!("plugin installation rollback state was lost"); + }; + rollback.installed.push(installed); + } + self.abort_install_if_terminal(&mut rollback).await?; + + 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.plugins); + anyhow::bail!("plugins were installed more than once"); + } + rollback.disarm(); + Ok(()) + } + + 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") + } + + 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); + 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; + } + 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 { + 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"); + } + } + } + + pub(crate) fn commit(&self) -> bool { + if self.terminal.load(Ordering::Acquire) { + self.close_installed_resources(); + return false; + } + let Some(installed) = self.installed.get() else { + return false; + }; + let mut staged = installed + .staged_apis + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + 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) { + for plugin in self.installed_plugins().iter().rev() { + close_plugin_resources(&plugin.manifest.id, &plugin.resources); + } + } + + async fn run_callback<'a>( + &'a self, + make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result<()>>, + ) -> Result<(), PluginCallbackError> { + let runtime = self.runtime.get().ok_or_else(|| { + PluginCallbackError::Callback(anyhow::anyhow!("plugin runtime is unavailable")) + })?; + bounded_plugin_callback(&**runtime, self.config.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 { + 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 { + let mut failures = Vec::new(); + 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 + .manifest + .capabilities + .contains(PluginCapability::Tasks) + .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, task_tracker); + let result = self + .run_callback(|| plugin.plugin.on_ready(plugin_scope)) + .await; + plugin.diagnostics.record_callback(&result); + if let Err(error) = result { + 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_plugins().iter().rev() { + let task_tracker = plugin + .manifest + .capabilities + .contains(PluginCapability::Tasks) + .then(|| plugin.resources.close_connection_tasks(scope.generation())); + if let Some(task_tracker) = &task_tracker { + let result = self + .wait_for_tasks(vec![task_tracker.completion_signal()]) + .await; + plugin.diagnostics.record_task_drain(&result); + match result { + 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, task_tracker); + let result = self + .run_callback(|| plugin.plugin.on_closed(plugin_scope)) + .await; + plugin.diagnostics.record_callback(&result); + if let Err(error) = result { + failures.push(format!("{}: {error:#}", plugin.manifest.id)); + } + } + 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) + }) + } + + 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(); + } + self.close_installed_resources(); + 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"); + } + } + + fn shutdown(&self) -> BoxFuture<'_, anyhow::Result<()>> { + Box::pin(async move { + let mut failures = Vec::new(); + self.signal_shutdown(); + for plugin in self.installed_plugins().iter().rev() { + let task_result = self + .wait_for_tasks(plugin.resources.task_completion_signals()) + .await; + plugin.diagnostics.record_task_drain(&task_result); + if let Err(error) = task_result { + failures.push(format!("{} tasks: {error:#}", plugin.manifest.id)); + } + 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)); + } + } + 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) + }) + } +} + +impl Drop for PluginHost { + fn drop(&mut self) { + self.signal_shutdown(); + } +} + +#[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, + config: PluginHostConfig, + current: Option, + mut installed: Vec, + upstream: Option>, + staged_apis: Option>, +) { + if let Some(plugin) = current { + let task_result = wait_for_plugin_tasks( + &*runtime, + config.task_drain_timeout(), + plugin.resources.task_completion_signals(), + ) + .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 + ); + } + let callback_result = bounded_plugin_callback(&*runtime, config.callback_timeout(), || { + plugin.plugin.shutdown() + }) + .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 + ); + } + } + while let Some(plugin) = installed.pop() { + let task_result = wait_for_plugin_tasks( + &*runtime, + config.task_drain_timeout(), + plugin.resources.task_completion_signals(), + ) + .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 + ); + } + let callback_result = bounded_plugin_callback(&*runtime, config.callback_timeout(), || { + plugin.plugin.shutdown() + }) + .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); + } + } + if let Some(upstream) = upstream + && let Err(error) = + bounded_plugin_callback(&*runtime, config.callback_timeout(), || upstream.shutdown()) + .await + { + log::warn!("Upstream lifecycle rollback failed: {error:#}"); + } + 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( + runtime: &dyn Runtime, + timeout: Duration, + completion_signals: Vec, +) -> Result<(), PluginTaskDrainError> { + 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(|_| PluginTaskDrainError::Timeout { + timeout_seconds: timeout.as_secs_f64(), + }) +} + +async fn bounded_plugin_callback<'a>( + runtime: &dyn Runtime, + timeout: Duration, + make_future: impl FnOnce() -> BoxFuture<'a, 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.map_err(PluginCallbackError::Callback), + futures::future::Either::Right(((), callback)) => { + let cancellation_panicked = + std::panic::catch_unwind(AssertUnwindSafe(|| drop(callback))).is_err(); + if cancellation_panicked { + return Err(PluginCallbackError::TimeoutCancellationPanic { + timeout_seconds: timeout.as_secs_f64(), + }); + } + Err(PluginCallbackError::Timeout { + timeout_seconds: timeout.as_secs_f64(), + }) + } + } +} + +async fn plugin_callback<'a>( + make_future: impl FnOnce() -> BoxFuture<'a, anyhow::Result<()>>, +) -> anyhow::Result<()> { + let mut future = std::panic::catch_unwind(AssertUnwindSafe(make_future)) + .map_err(|_| anyhow::anyhow!("callback panicked before returning a future"))?; + 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, 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| { + 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? +} + +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(()) + } 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() + .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() + .filter(|host| host.is_published()) + .map(|host| host.stats()) + } + + /// 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() + .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; + + use bytes::Bytes; + + 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) + } + + #[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() + .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)) + .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, + } + + 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 FailingInstallLifecycle { + log: Log, + } + + struct CaptureInstallClient { + 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(); + Box::pin(async move { + sender.send(client).await?; + Ok(()) + }) + } + } + + 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, + 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 { + client + .upgrade() + .ok_or_else(|| anyhow::anyhow!("client unavailable during install"))? + .signal_shutdown_sync(); + Ok(()) + }) + } + } + + 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; + + 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(()) + }) + } + } + + #[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 { + id: "runtime-dependent", + dependency: Some("runtime-foundation"), + log: Arc::clone(&log), + }) + .with_untyped_plugin_arc(foundation) + .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())); + 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() + ); + } + + #[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); + 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); + 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); + } + + #[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, + } + + 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.into_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.into_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 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 { + 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, PluginHostConfig::default()); + rollback.installed.push(InstalledPlugin { + plugin: erased_plugin, + manifest, + resources, + diagnostics: PluginDiagnostics::new(), + }); + 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, + } + + 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, + api_dropped: Arc, + } + + struct RollbackApi { + _drop_flag: DropFlag, + } + + impl ClientPlugin for RollbackPlugin { + type Api = RollbackApi; + + 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(); + let api_dropped = self.api_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(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(()) + }) + } + } + + 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 api_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(), + api_dropped: api_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"); + assert!(api_dropped.load(Ordering::Acquire)); + } + + 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 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); + 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(), + api_dropped: api_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) + && api_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<()>, + 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 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() + .await + .with_plugin(BlockingInstallPlugin { + log: log.clone(), + started: started_tx, + release: release_rx, + }) + .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 }); + 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) + && api_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, + } + + 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)), + api_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, + closed_after_task: Arc, + 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 = (); + + 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(()) + }) + } + + 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) { + 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 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 { + install_started: install_started.clone(), + 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 + .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").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 { + 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)); + 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] + 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); + let task_dropped = Arc::new(AtomicBool::new(false)); + let tasks = PluginConnectionTasks { + 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 + .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) + )); + } + + #[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(); + resources.activate(); + 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 }); + 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(), + 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 }); + 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, + } + + 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(()) + }) + } + } + + 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(()) + }) + } + } + + 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())); + 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"] + ); + 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; + } + + #[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"] + ); + 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; + } + + #[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"] + ); + 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; + } + + #[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.into_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; + + 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 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 = PluginCoreEventSubscription; + + 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 { + 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(subscription)) + }) + } + } + + struct PanickingSubscriptionPlugin; + + impl ClientPlugin for PanickingSubscriptionPlugin { + type Api = PluginCoreEventSubscription; + + 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 { + 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(subscription)) + }) + } + } + + 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 = PluginCoreEventSubscription; + + 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 { + 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(subscription)) + }) + } + } + + 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; + + struct ReentrantSubscriptionApi { + events: PluginCoreEvents, + _subscription: PluginCoreEventSubscription, + } + + impl ClientPlugin for ReentrantSubscriptionPlugin { + type Api = ReentrantSubscriptionApi; + + 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"))?; + let subscription = events.subscribe( + EventInterest::of(&[EventKind::Connected]), + Arc::new(ReentrantSubscriptionHandler { + events: events.clone(), + }), + )?; + Ok(Arc::new(ReentrantSubscriptionApi { + events, + _subscription: subscription, + })) + }) + } + } + + #[tokio::test] + 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.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!( + !client + .core + .event_bus + .has_handler_for(wacore::types::events::EventKind::Connected) + ); + 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 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() + .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)); + 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)); + 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; + } + + #[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 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.events.subscribe( + EventInterest::of(&[EventKind::Connected]), + Arc::new(ReentrantSubscriptionHandler { + events: subscribe_events.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)); + let client = complete_builder() + .await + .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() + .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; + + impl ClientPlugin for CapabilityProbe { + type Api = [bool; 5]; + + 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(), + context.plugin_events().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.into_client(); + assert_eq!( + client.plugin::().as_deref(), + 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)); + const TICK_PAYLOAD: &[u8] = br#"{"messages":1}"#; + let generation = client.connection_generation.load(Ordering::Acquire); + assert_eq!( + publisher + .publish( + &tick, + 2, + PluginEventPayloadEncoding::Json, + Bytes::from_static(TICK_PAYLOAD), + ) + .expect("publish tick"), + PluginEventPublishReport { + matched: 1, + enqueued: 1, + dropped: 0, + 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(TICK_PAYLOAD)); + 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) + )); + 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); + } +} diff --git a/src/receipt.rs b/src/receipt.rs index 8683f2ce6..22305c3c1 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -2012,7 +2012,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/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()); } } 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 c0aa5a641..9ebd18178 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] @@ -282,9 +283,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); @@ -324,14 +325,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 } @@ -342,7 +347,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; } /// } @@ -364,22 +369,176 @@ 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. + let retired = std::mem::replace(&mut *guard, Arc::clone(&snapshot)); + self.store_aggregate(&snapshot); + drop(guard); + drop(retired); + 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 { @@ -388,22 +547,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. @@ -416,25 +596,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)); } } } @@ -737,7 +911,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), @@ -1954,8 +2128,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())); @@ -1976,53 +2150,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] @@ -2041,13 +2210,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)); @@ -2069,11 +2240,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. @@ -2107,14 +2279,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), @@ -2130,4 +2303,106 @@ 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 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; + 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); + } }