Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ event-listener = { version = "5", default-features = false }
flate2 = { version = "1.1.9", default-features = false, features = ["zlib-rs"] }
futures = { version = "0.3", default-features = false, features = ["alloc", "async-await"] }
getrandom = { version = "0.4", default-features = false }
hashbrown = { version = "0.15.5", default-features = false, features = ["equivalent"] }
heck = "0.5"
hex = { version = "0.4", default-features = false, features = ["alloc"] }
hkdf = { version = "0.13.0", default-features = false }
Expand Down Expand Up @@ -172,6 +173,7 @@ bytes = { workspace = true }
chrono = { workspace = true, features = ["clock"] }
event-listener = { workspace = true }
futures = { workspace = true, features = ["std"] }
hashbrown = { workspace = true }
hex = { workspace = true }
itoa = { workspace = true }
log = { workspace = true }
Expand Down
3 changes: 3 additions & 0 deletions agent_docs/observability.md
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,9 @@ figures come from the `wacore::stats::HeapSize` trait:
`LidPnEntry`, `ResolvedGroupDevices`, ...).
- Store-backed caches (Redis etc.) report `bytes: 0` — their entries are not
process memory.
- In-flight history sync reports queued/running task count, retained compressed
payload storage, and lifetime peaks. Inline payloads count while queued;
external payloads contribute their `Vec` capacity once materialized.

Semantics: honest estimates for attribution and leak detection, not
byte-exact accounting. The e2e `memory_soak.rs` logs the byte totals next to
Expand Down
39 changes: 32 additions & 7 deletions src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -239,6 +239,13 @@ pub struct MemoryReport {
pub undecryptable_dispatched: u64,
pub pdo_pending_requests: u64,
pub pdo_requested: u64,
/// Queued/running history-sync tasks and their retained compressed-payload
/// allocation estimate.
pub history_sync_tasks: CollectionStats,
/// Lifetime high-water mark of queued/running history-sync tasks.
pub history_sync_tasks_peak: u64,
/// Lifetime high-water mark of retained compressed-payload storage.
pub history_sync_payload_bytes_peak: u64,
// -- Capacity-only caches (coordination, counts only) --
pub session_locks: u64,
pub chat_lanes: u64,
Expand Down Expand Up @@ -267,7 +274,7 @@ 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).
fn collections(&self) -> [(&'static str, &CollectionStats); 10] {
fn collections(&self) -> [(&'static str, &CollectionStats); 11] {
[
("group_cache:", &self.group_cache),
("device_registry_cache:", &self.device_registry_cache),
Expand All @@ -279,6 +286,7 @@ impl MemoryReport {
("signal_sessions:", &self.signal_sessions),
("signal_identities:", &self.signal_identities),
("signal_sender_keys:", &self.signal_sender_keys),
("history_sync_tasks:", &self.history_sync_tasks),
]
}

Expand All @@ -298,8 +306,10 @@ impl std::fmt::Display for MemoryReport {
writeln!(f, " {name:<22} {:>7} entries {:>10} B", c.entries, c.bytes)
}
// First TTL_BOUNDED entries of collections() are the TTL-bounded
// caches; the rest are the Signal store caches.
// caches; the next SIGNAL_CACHES are Signal store caches. The final
// entry is transient history-sync retention.
const TTL_BOUNDED: usize = 7;
const SIGNAL_CACHES: usize = 3;
let collections = self.collections();
writeln!(f, "=== Memory Report ===")?;
writeln!(f, "--- TTL-bounded caches ---")?;
Expand Down Expand Up @@ -345,9 +355,25 @@ impl std::fmt::Display for MemoryReport {
)?;
writeln!(f, " app_state_syncing: {}", self.app_state_syncing)?;
writeln!(f, "--- Signal store caches ---")?;
for (name, c) in &collections[TTL_BOUNDED..] {
for (name, c) in &collections[TTL_BOUNDED..TTL_BOUNDED + SIGNAL_CACHES] {
line(f, name, c)?;
}
writeln!(f, "--- In-flight history sync ---")?;
line(
f,
collections[TTL_BOUNDED + SIGNAL_CACHES].0,
&self.history_sync_tasks,
)?;
writeln!(
f,
" peak tasks: {}",
self.history_sync_tasks_peak
)?;
writeln!(
f,
" peak payload storage: {} B",
self.history_sync_payload_bytes_peak
)?;
writeln!(f, "--- Misc ---")?;
writeln!(f, " chatstate_handlers: {}", self.chatstate_handlers)?;
writeln!(f, " custom_enc_handlers: {}", self.custom_enc_handlers)?;
Expand Down Expand Up @@ -706,10 +732,9 @@ pub struct Client {
/// Empty (zero capacity) outside the offline window.
pub(crate) offline_receipt_buffer:
std::sync::Mutex<Vec<Arc<crate::types::message::MessageInfo>>>,
/// Number of history sync tasks currently queued or running.
pub(crate) history_sync_tasks_in_flight: Arc<AtomicUsize>,
/// Notifier triggered when history sync work becomes idle.
pub(crate) history_sync_idle_notifier: Arc<event_listener::Event>,
/// Task count, retained payload storage, peaks, and idle notification for
/// history sync work.
pub(crate) history_sync_activity: Arc<crate::sync_task::HistorySyncActivity>,
/// Flushed by `disconnect()`/`reconnect()` before tearing down the transport
/// so in-flight delivery receipts aren't dropped with `NotConnected`
/// (issue #571).
Expand Down
7 changes: 7 additions & 0 deletions src/client/accessors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,10 @@ impl Client {
let app_state_key_requests = self.app_state_key_requests.lock().await.len();
let app_state_syncing = self.app_state_syncing.lock().await.len();
let chatstate_handlers = self.chatstate_handlers.read().await.len();
let history_sync_tasks = CollectionStats::new(
self.history_sync_activity.tasks() as u64,
self.history_sync_activity.payload_bytes() as u64,
);

MemoryReport {
group_cache,
Expand All @@ -177,6 +181,9 @@ impl Client {
undecryptable_dispatched: self.undecryptable_dispatched.entry_count(),
pdo_pending_requests: self.pdo_pending_requests.entry_count(),
pdo_requested: self.pdo_requested.entry_count(),
history_sync_tasks,
history_sync_tasks_peak: self.history_sync_activity.tasks_peak() as u64,
history_sync_payload_bytes_peak: self.history_sync_activity.payload_bytes_peak() as u64,
session_locks: self.session_locks.entry_count(),
chat_lanes: self.chat_lanes.entry_count(),
group_distribution_locks: group_distribution_locks.entries,
Expand Down
8 changes: 6 additions & 2 deletions src/client/app_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -287,9 +287,13 @@ impl Client {
message_id,
notification,
} => {
self.process_history_sync_task(message_id, *notification)
let retained_payload_bytes = notification
.inline_payload
.as_ref()
.map_or(0, bytes::Bytes::len);
let mut tracker = self.track_history_sync_task(retained_payload_bytes);
self.process_history_sync_task_tracked(message_id, *notification, &mut tracker)
.await;
self.finish_history_sync_task();
}
crate::sync_task::MajorSyncTask::AppStateSync { name, full_sync } => {
if let Err(e) = self.process_app_state_sync_task(name, full_sync).await {
Expand Down
7 changes: 2 additions & 5 deletions src/client/lifecycle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -228,8 +228,7 @@ impl Client {
offline_sync_finish_started: Arc::new(AtomicBool::new(false)),
offline_receipt_buffer: std::sync::Mutex::new(Vec::new()),
inbound_commit_batch: Default::default(),
history_sync_tasks_in_flight: Arc::new(AtomicUsize::new(0)),
history_sync_idle_notifier: Arc::new(event_listener::Event::new()),
history_sync_activity: Arc::new(crate::sync_task::HistorySyncActivity::new()),
outbound_flush: Arc::new(crate::flush_scope::FlushScope::new()),
delivery_receipt_queue: std::sync::OnceLock::new(),
presence_subscriptions: Arc::new(async_lock::Mutex::new(HashSet::new())),
Expand Down Expand Up @@ -899,9 +898,7 @@ impl Client {
Ok(mut guard) => *guard = None,
Err(poison) => *poison.into_inner() = None,
}
self.history_sync_tasks_in_flight
.store(0, Ordering::Relaxed);
self.history_sync_idle_notifier.notify(usize::MAX);
self.history_sync_activity.reset();
// Drain all pending IQ waiters so they fail fast with InternalChannelClosed
// instead of hanging until the 75s timeout.
// Scoped so the sync guard is dropped before the awaits below (a
Expand Down
40 changes: 24 additions & 16 deletions src/client/node_io.rs
Original file line number Diff line number Diff line change
Expand Up @@ -490,40 +490,48 @@ impl Client {
return;
}

// Most messages do not need a transport <ack> from this generic gate.
// Move those nodes into their chat lane instead of retaining a second
// Arc in this dispatcher while decryption starts. Besides removing an
// atomic refcount pair, this lets a large uniquely-owned pkmsg donate
// its receive buffer to authenticated in-place decryption. Newsletter
// and status messages keep the extra owner until their deferred ack is
// encoded, preserving the existing acknowledgement semantics.
let should_ack = self.should_ack(nr);
let deferred_ack_node = should_ack.then(|| Arc::clone(&node));

// Bypass async_trait's boxed future for the hot built-in handlers while
// retaining router registration for direct router callers.
let handled = match nr.tag.as_ref() {
match nr.tag.as_ref() {
"ack" => {
self.handle_ack_response_arc(&node);
true
}
"receipt" => {
self.handle_receipt_inline(Arc::clone(&node));
true
self.handle_receipt_inline(node);
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
"message" => {
crate::handlers::message::MessageHandler::handle_inline(
self.clone(),
Arc::clone(&node),
node,
&mut cancelled,
)
.await
.await;
}
_ => {
self.stanza_router
let handled = self
.stanza_router
.dispatch(self.clone(), Arc::clone(&node), &mut cancelled)
.await
.await;
if !handled {
warn!(
"Received unknown top-level node: {}",
DisplayableNodeRef(node.get())
);
}
}
};
if !handled {
warn!(
"Received unknown top-level node: {}",
DisplayableNodeRef(nr)
);
}

// Send the deferred ACK if applicable and not cancelled by handler
if self.should_ack(nr) && !cancelled {
if !cancelled && let Some(node) = deferred_ack_node {
self.maybe_deferred_ack(node).await;
}
}
Expand Down
36 changes: 15 additions & 21 deletions src/client/sessions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -213,25 +213,19 @@ impl Client {
}
}

pub(crate) fn begin_history_sync_task(&self) {
self.history_sync_tasks_in_flight
.fetch_add(1, Ordering::Relaxed);
}

pub(crate) fn finish_history_sync_task(&self) {
// The `previous <= 1` clamp is also the underflow guard: a detached task
// that finishes after cleanup_connection_state() reset the counter to 0
// hits fetch_sub-from-0, which momentarily wraps the stored value to
// usize::MAX — but previous == 0 takes this branch and stores 0, so the
// wrap never sticks (and it never leaves the idle waiter blocked).
let previous = self
.history_sync_tasks_in_flight
.fetch_sub(1, Ordering::Relaxed);
if previous <= 1 {
self.history_sync_tasks_in_flight
.store(0, Ordering::Relaxed);
self.history_sync_idle_notifier.notify(usize::MAX);
}
pub(crate) fn begin_history_sync_task(&self, retained_payload_bytes: usize) {
self.history_sync_activity.begin(retained_payload_bytes);
}

pub(crate) fn track_history_sync_task(
&self,
retained_payload_bytes: usize,
) -> crate::sync_task::HistorySyncTaskTracker {
self.history_sync_activity.tracker(retained_payload_bytes)
}

pub(crate) fn finish_history_sync_task(&self, retained_payload_bytes: usize) {
drop(self.track_history_sync_task(retained_payload_bytes));
}

pub async fn wait_for_startup_sync(&self, timeout: std::time::Duration) -> Result<()> {
Expand All @@ -251,8 +245,8 @@ impl Client {
}

loop {
let history_fut = self.history_sync_idle_notifier.listen();
if self.history_sync_tasks_in_flight.load(Ordering::Relaxed) == 0 {
let history_fut = self.history_sync_activity.listen();
if self.history_sync_activity.tasks() == 0 {
return Ok(());
}

Expand Down
31 changes: 2 additions & 29 deletions src/client/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3409,33 +3409,6 @@ async fn terminal_disconnect_propagates_to_per_connection_signal() {
);
}

// Counting allocator for the empirical allocation guard below. Process-wide
// for the unit-test binary, but the only cost is one relaxed atomic add per
// alloc; tests that never read the counter are unaffected. Host-only harness:
// std's 64-bit atomic is fine here since this never compiles for embedded.
#[allow(clippy::disallowed_types)]
mod counting_alloc {
use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicU64, Ordering};

pub(super) static ALLOCS: AtomicU64 = AtomicU64::new(0);

struct CountingAlloc;

unsafe impl GlobalAlloc for CountingAlloc {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
ALLOCS.fetch_add(1, Ordering::Relaxed);
unsafe { System.alloc(layout) }
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
unsafe { System.dealloc(ptr, layout) }
}
}

#[global_allocator]
static GLOBAL: CountingAlloc = CountingAlloc;
}

/// Locks the zero-allocation property of the ack miss path: id resolution and
/// the waiter probe must borrow from the node buffer. An `into_owned()` here
/// costs one String per received ack, which the e2e dhat profile caught live.
Expand All @@ -3450,9 +3423,9 @@ async fn ack_miss_path_does_not_heap_allocate() {
// in every window, so the minimum only reaches 0 when the path is clean.
let mut min_delta = u64::MAX;
for _ in 0..100 {
let before = counting_alloc::ALLOCS.load(std::sync::atomic::Ordering::Relaxed);
let before = crate::test_alloc::ALLOCS.load(std::sync::atomic::Ordering::Relaxed);
let handled = client.handle_ack_response_arc(&node);
let after = counting_alloc::ALLOCS.load(std::sync::atomic::Ordering::Relaxed);
let after = crate::test_alloc::ALLOCS.load(std::sync::atomic::Ordering::Relaxed);
assert!(!handled, "no waiter is registered for this id");
min_delta = min_delta.min(after - before);
}
Expand Down
1 change: 1 addition & 0 deletions src/features/chat_actions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -826,6 +826,7 @@ impl Client {
/// like [`ChatActions`] and [`Labels`](crate::Labels) wrap it.
///
/// ```no_run
/// # #![recursion_limit = "512"]
/// # async fn ex(client: &whatsapp_rust::Client) -> anyhow::Result<()> {
/// use whatsapp_rust::schemas;
/// use whatsapp_rust::waproto::whatsapp as wa;
Expand Down
Loading
Loading