diff --git a/src/bot.rs b/src/bot.rs index 3928960d0..49edccb8b 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -15,6 +15,7 @@ use std::future::Future; use std::marker::PhantomData; use std::pin::Pin; use std::sync::Arc; +use std::time::Duration; use thiserror::Error; use wacore::proto_helpers::MessageBuilderExt; use wacore::runtime::Runtime; @@ -565,6 +566,7 @@ pub struct BotBuilder< cache_config: CacheConfig, wanted_pre_key_count: Option, resend_rate_limit: Option<(u32, u32)>, + app_state_key_wait: Option, task_instrument: Option>, alloc_meter: Option>, _marker: PhantomData<(B, T, H, R)>, @@ -589,6 +591,7 @@ impl BotBuilder BotBuilder { cache_config: self.cache_config, wanted_pre_key_count: self.wanted_pre_key_count, resend_rate_limit: self.resend_rate_limit, + app_state_key_wait: self.app_state_key_wait, task_instrument: self.task_instrument, alloc_meter: self.alloc_meter, _marker: PhantomData, @@ -1031,6 +1035,16 @@ impl BotBuilder { self } + /// Override how long the initial critical app-state sync waits for the + /// encrypted app-state key-share before attempting the snapshot (default 5s). + /// Raise it when a large concurrent history sync can delay the key-share past + /// the default, which otherwise fails critical sync with "didn't find app + /// state key" and drops the account's saved contact names. + pub fn with_app_state_key_wait(mut self, wait: Duration) -> Self { + self.app_state_key_wait = Some(wait); + self + } + /// Set an initial push name on the device before connecting. /// /// This is included in the `ClientPayload` during registration, allowing the @@ -1185,6 +1199,10 @@ impl BotBuilder { client.set_resend_rate_limit(burst, refill_per_min); } + if let Some(wait) = self.app_state_key_wait { + client.set_app_state_key_wait(wait); + } + Ok(Bot { client, sync_task_receiver: Some(sync_task_receiver), @@ -1636,6 +1654,37 @@ mod tests { ); } + #[tokio::test] + async fn test_bot_builder_app_state_key_wait() { + let backend = create_test_sqlite_backend().await; + let bot = Bot::builder() + .with_backend_arc(backend) + .with_transport_factory(TokioWebSocketTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_app_state_key_wait(Duration::from_secs(30)) + .with_runtime(TokioRuntime) + .build() + .await + .expect("Failed to build bot with custom app-state key wait"); + + assert_eq!(bot.client().app_state_key_wait(), Duration::from_secs(30)); + } + + #[tokio::test] + async fn test_bot_builder_default_app_state_key_wait() { + let backend = create_test_sqlite_backend().await; + let bot = Bot::builder() + .with_backend_arc(backend) + .with_transport_factory(TokioWebSocketTransportFactory::new()) + .with_http_client(MockHttpClient) + .with_runtime(TokioRuntime) + .build() + .await + .expect("Failed to build bot"); + + assert_eq!(bot.client().app_state_key_wait(), Duration::from_secs(5)); + } + #[tokio::test] async fn registered_handlers_accumulate_instead_of_replacing() { let backend = create_test_sqlite_backend().await; diff --git a/src/client.rs b/src/client.rs index bbd3461be..e5fbe93d4 100644 --- a/src/client.rs +++ b/src/client.rs @@ -792,6 +792,14 @@ pub struct Client { /// Clamped to the protocol-safe range at upload time. pub(crate) wanted_pre_key_count: AtomicUsize, + /// How long the initial critical app-state sync waits for the encrypted + /// app-state key-share before attempting the snapshot fetch. Default 5s; set + /// via [`BotBuilder::with_app_state_key_wait`] or [`Client::set_app_state_key_wait`]. + /// Raise it when a large concurrent history sync can delay the key-share past + /// the default, which otherwise fails critical sync with "didn't find app + /// state key" and drops the account's saved contact names. + pub(crate) app_state_key_wait_ms: AtomicU64, + /// Cache configuration for TTL and capacity of all caches. /// Stored for use by lazily-initialized caches (group_cache). pub(crate) cache_config: CacheConfig, diff --git a/src/client/accessors.rs b/src/client/accessors.rs index 6a200f01b..25fbc5f3f 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -69,6 +69,19 @@ impl Client { self.wanted_pre_key_count.load(Ordering::Relaxed) } + /// Set how long the initial critical app-state sync waits for the encrypted + /// app-state key-share before attempting the snapshot. Read once when the + /// initial sync starts, so set it before connecting. + pub fn set_app_state_key_wait(&self, wait: Duration) { + self.app_state_key_wait_ms + .store(wait.as_millis() as u64, Ordering::Relaxed); + } + + /// The configured initial app-state key-share wait. + pub fn app_state_key_wait(&self) -> Duration { + Duration::from_millis(self.app_state_key_wait_ms.load(Ordering::Relaxed)) + } + /// Retune the per-chat outbound resend rate limiter live (no reconnect). /// /// Outbound resends to a chat are bounded by a token bucket: `burst` is the diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index d2c34c4e7..992b4b4b4 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -269,6 +269,7 @@ impl Client { override_version, skip_history_sync: AtomicBool::new(false), wanted_pre_key_count: AtomicUsize::new(crate::prekeys::DEFAULT_WANTED_PRE_KEY_COUNT), + app_state_key_wait_ms: AtomicU64::new(5_000), cache_config, self_weak: std::sync::OnceLock::new(), saver_handle: std::sync::OnceLock::new(), diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 412879046..02fa2bf68 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -918,13 +918,14 @@ impl Client { .initial_app_state_keys_received .load(Ordering::Relaxed) { + let key_wait = client_clone.app_state_key_wait(); debug!( target: "Client/AppState", - "Waiting up to 5s for app state keys..." + "Waiting up to {key_wait:?} for app state keys..." ); let _ = rt_timeout( &*client_clone.runtime, - Duration::from_secs(5), + key_wait, client_clone.initial_keys_synced_notifier.listen(), ) .await;