diff --git a/wacore/benches/send_receive_benchmark.rs b/wacore/benches/send_receive_benchmark.rs index 7b61503d0..ffd97155c 100644 --- a/wacore/benches/send_receive_benchmark.rs +++ b/wacore/benches/send_receive_benchmark.rs @@ -16,8 +16,8 @@ use wacore::client::context::{GroupInfo, SendContextResolver}; use wacore::messages::MessageUtils; use wacore::runtime::{AbortHandle, Runtime}; use wacore::send::{ - GroupStanzaRequest, SenderKeyDistributionPolicy, SignalStores, prepare_group_stanza, - prepare_peer_stanza, + DmStanzaRequest, GroupStanzaRequest, ResolvedDmDevices, SenderKeyDistributionPolicy, + SignalStores, prepare_dm_stanza, prepare_group_stanza, prepare_peer_stanza, }; use wacore::types::jid::{JidExt, make_sender_key_name}; use wacore::types::message::AddressingMode; @@ -344,6 +344,15 @@ impl User { } } + /// The same account on another device: same keys, distinct address. Real + /// DM fan-out targets several devices of one user, which `User::new` alone + /// cannot express since it derives the address from user+server. + fn with_device(mut self, device: u16) -> Self { + self.jid.device = device; + self.address = self.jid.to_protocol_address(); + self + } + fn prekey_bundle(&self) -> PreKeyBundle { PreKeyBundle::new( self.identity.reg_id, @@ -488,6 +497,30 @@ fn extract_skmsg_bytes(stanza: &Node) -> Vec { } } +/// Encrypt one padded text message from `alice` to `to`, returning the wire +/// bytes and the enc type the stanza would carry. Shared by both receive +/// fixtures so the ratchet and steady-state benches cannot drift apart on +/// padding or enc-type mapping. +fn encrypt_one(alice: &mut User, to: &ProtocolAddress) -> (Vec, String) { + futures::executor::block_on(async { + let ct = message_encrypt( + &MessageUtils::pad_message_v2(text_msg().encode_to_vec()), + to, + &mut alice.sessions, + &mut alice.identity, + ) + .await + .unwrap(); + match ct { + CiphertextMessage::SignalMessage(m) => (m.serialized().to_vec(), "msg".to_string()), + CiphertextMessage::PreKeySignalMessage(m) => { + (m.serialized().to_vec(), "pkmsg".to_string()) + } + _ => panic!("unexpected ciphertext type in benchmark fixture"), + } + }) +} + fn decrypt_dm( ciphertext: &[u8], enc_type: &str, @@ -542,10 +575,37 @@ fn decrypt_group( // DM setups // --------------------------------------------------------------------------- +/// The `` blob a paired device carries. +/// +/// Every send path takes `Option<&ADVSignedDeviceIdentity>`, and `None` changes +/// what they do — differently per path, so the distinction matters here: +/// +/// - `prepare_peer_stanza` refuses, and pays a whole extra session checkout to +/// find out (`account.is_none() && pkmsg_would_be_emitted(..)`); +/// - `prepare_dm_stanza` and the `BestEffort` group policy are lenient: they +/// drop the `` child and send anyway; +/// - only `SenderKeyDistributionPolicy::Required` propagates the error. +/// +/// A paired device always has one (`device_snapshot.account.as_deref()`), so a +/// bench passing `None` measures a shape production never sends — and on the +/// lenient paths it silently omits a stanza child rather than failing loudly. +/// +/// Only the bytes matter here — it is serialised into the stanza, never +/// validated by the sender. +fn bench_account() -> wa::ADVSignedDeviceIdentity { + wa::ADVSignedDeviceIdentity { + details: Some(vec![0xAD; 32]), + account_signature_key: Some(vec![0xAC; 32]), + account_signature: Some(vec![0x51; 64]), + device_signature: Some(vec![0xD5; 64]), + } +} + struct DmSendData { alice: User, bob_jid: Jid, msg: wa::Message, + account: wa::ADVSignedDeviceIdentity, } fn setup_dm_send() -> DmSendData { @@ -556,6 +616,7 @@ fn setup_dm_send() -> DmSendData { alice, bob_jid: bob.jid, msg: text_msg(), + account: bench_account(), } } @@ -571,25 +632,13 @@ fn setup_dm_recv() -> DmRecvData { let mut bob = User::new("5511999000002", "s.whatsapp.net"); establish_bidirectional(&mut alice, &mut bob); - // Encrypt a message (subsequent, not PreKey — realistic steady-state) - let (ciphertext, enc_type) = futures::executor::block_on(async { - let ct = message_encrypt( - &MessageUtils::pad_message_v2(text_msg().encode_to_vec()), - &bob.address, - &mut alice.sessions, - &mut alice.identity, - ) - .await - .unwrap(); - - match ct { - CiphertextMessage::SignalMessage(msg) => (msg.serialized().to_vec(), "msg".to_string()), - CiphertextMessage::PreKeySignalMessage(msg) => { - (msg.serialized().to_vec(), "pkmsg".to_string()) - } - _ => panic!("unexpected type"), - } - }); + // A `msg` rather than a `pkmsg`, but NOT steady state: Bob has not decrypted + // anything on this sending chain yet, so this first one costs him a DH ratchet + // step (two X25519 agreements plus a key generation). That is a real shape — + // the first message after the peer replies — but it is ~40x a steady-state + // decrypt, so the two need separate benches or the common case is invisible. + let bob_addr = bob.address.clone(); + let (ciphertext, enc_type) = encrypt_one(&mut alice, &bob_addr); DmRecvData { bob, @@ -613,6 +662,11 @@ struct GrpSendData { force_skdm: bool, resolver: MockResolver, msg: wa::Message, + /// Built in setup, not in the measured body: production borrows the cached + /// account rather than allocating one per send, so building it inside the + /// closure would charge every group baseline four `Vec` allocations that no + /// send performs. + account: wa::ADVSignedDeviceIdentity, // Built once in setup so the measured body excludes thread-pool startup // (building the pool inside the bench body would charge its syscalls to // the encrypt path). @@ -659,6 +713,7 @@ fn setup_group_send(n: usize) -> GrpSendData { force_skdm: false, resolver: MockResolver(devices), msg: text_msg(), + account: bench_account(), runtime: BenchRuntime, } } @@ -737,6 +792,7 @@ fn setup_group_recv() -> GrpRecvData { let runtime = BenchRuntime; let message = text_msg(); + let account = bench_account(); let result = futures::executor::block_on(prepare_group_stanza( &runtime, &mut stores, @@ -745,7 +801,7 @@ fn setup_group_recv() -> GrpRecvData { group: &group_info, own_jid: &own_jid, own_lid: &own_jid, - account: None, + account: Some(&account), to: &group_jid, message: &message, message_id: "bench-grp-recv", @@ -774,8 +830,149 @@ fn setup_group_recv() -> GrpRecvData { // Benchmarks // =========================================================================== +struct DmFanoutData { + alice: User, + to_jid: Jid, + msg: wa::Message, + account: wa::ADVSignedDeviceIdentity, + devices: ResolvedDmDevices, + resolver: MockResolver, + runtime: BenchRuntime, +} + +/// The real 1:1 send: `prepare_dm_stanza`, which fans out to the recipient's +/// devices and our own companions and builds the `` tree. +/// +/// `n_recipient_devices` counts the contact's devices; one own companion is +/// always added on top, because a paired account has at least the phone and +/// this bench should not measure the degenerate zero-companion shape. The +/// total fan-out is therefore `n_recipient_devices + 1` pairwise encryptions. +/// +/// Sessions are established **bidirectionally**. `establish_session` alone +/// leaves `pending_pre_key` set, so every outbound stays a `pkmsg` and the +/// bench would measure first-message prekey wrapping plus `` +/// serialisation forever — the opposite of the repeat-send shape intended here. +fn setup_dm_fanout(n_recipient_devices: usize) -> DmFanoutData { + let mut alice = User::new("5511999000001", "s.whatsapp.net"); + let own_jid = alice.jid.clone(); + let to_base: Jid = "5511999000002@s.whatsapp.net".parse().unwrap(); + + let mut all_devices = Vec::with_capacity(n_recipient_devices + 1); + for i in 0..n_recipient_devices { + let mut device = to_base.clone(); + device.device = i as u16; + let mut peer = User::new("5511999000002", "s.whatsapp.net").with_device(i as u16); + establish_bidirectional(&mut alice, &mut peer); + all_devices.push(device); + } + // One own companion, so the own-device half of the partition is exercised. + let mut companion = own_jid.clone(); + companion.device = 42; + let mut own_peer = User::new("5511999000001", "s.whatsapp.net").with_device(42); + establish_bidirectional(&mut alice, &mut own_peer); + all_devices.push(companion); + + // Assert the sessions really are acknowledged, so a fixture regression cannot + // quietly turn this back into a first-message benchmark: with `pending_pre_key` + // still set every outbound is a pkmsg, which measures prekey wrapping and + // `` serialisation instead of the repeat send. + for device in &all_devices { + let addr = device.to_protocol_address(); + let emits_pkmsg = futures::executor::block_on(wacore::send::pkmsg_would_be_emitted( + &mut alice.sessions, + &addr, + )) + .expect("session must be loadable in setup"); + assert!( + !emits_pkmsg, + "fixture must use acknowledged sessions; {addr} still emits pkmsg" + ); + } + + let resolver = MockResolver(all_devices.clone()); + let devices = ResolvedDmDevices::new(all_devices, &own_jid, None); + // Warm the phash memo in setup, as the per-recipient memo does for repeat + // sends in production; otherwise every iteration measures the cold path. + devices.phash(); + + DmFanoutData { + alice, + to_jid: to_base, + msg: text_msg(), + account: bench_account(), + devices, + resolver, + runtime: BenchRuntime, + } +} + +/// One recipient device plus our companion: two pairwise encryptions. +fn setup_dm_fanout_1() -> DmFanoutData { + setup_dm_fanout(1) +} +/// Four recipient devices plus our companion: five pairwise encryptions. +fn setup_dm_fanout_4() -> DmFanoutData { + setup_dm_fanout(4) +} + +fn run_dm_fanout(d: &mut DmFanoutData) { + let own_jid = d.alice.jid.clone(); + let mut stores = SignalStores { + sender_key_store: &mut d.alice.sender_keys, + session_store: &mut d.alice.sessions, + identity_store: &mut d.alice.identity, + prekey_store: &mut d.alice.prekeys, + signed_prekey_store: &d.alice.signed_prekeys, + }; + let prepared = futures::executor::block_on(prepare_dm_stanza( + &d.runtime, + &mut stores, + &d.resolver, + DmStanzaRequest { + own_jid: &own_jid, + account: Some(&d.account), + to: &d.to_jid, + message: &d.msg, + message_id: "b-dm", + edit: None, + extra_nodes: &[], + devices: &d.devices, + pre_encoded: None, + }, + )) + .unwrap(); + black_box(marshal(&prepared.node).unwrap()); +} + +/// The 1:1 path an actual message takes. Nothing covered this before: the only +/// DM-shaped bench measured `prepare_peer_stanza`, which is the own-device peer +/// path and skips the fan-out, the participants tree and the reporting token. #[divan::bench] fn bench_dm_send(bencher: divan::Bencher) { + bencher + .with_inputs(setup_dm_fanout_1) + .bench_refs(run_dm_fanout); +} + +/// Multi-device recipient: the shape that actually spawns encrypt tasks. +/// Named for the total fan-out — four recipient devices plus one own companion, +/// so five pairwise encryptions. +#[divan::bench] +fn bench_dm_send_5_way_fanout(bencher: divan::Bencher) { + bencher + .with_inputs(setup_dm_fanout_4) + .bench_refs(run_dm_fanout); +} + +/// A peer stanza: what we send to our OWN other devices, not a DM to a contact. +/// See `bench_dm_send` for the contact path. +/// +/// `account` is `Some`, matching a paired device. With `None` the `&&` in +/// `prepare_peer_stanza` does not short-circuit and every send pays an extra +/// `SessionCheckout::load` + `commit()` just to pre-flight the pkmsg case — +/// work production never does, and enough of it to move the baseline. +#[divan::bench] +fn bench_peer_send(bencher: divan::Bencher) { bencher.with_inputs(setup_dm_send).bench_refs(|d| { let signal_addr = d.bob_jid.to_protocol_address(); let node = futures::executor::block_on(prepare_peer_stanza( @@ -785,7 +982,7 @@ fn bench_dm_send(bencher: divan::Bencher) { &signal_addr, &d.msg, "b-001", - None, + Some(&d.account), )) .unwrap(); black_box(marshal(&node).unwrap()); @@ -804,6 +1001,44 @@ fn bench_dm_recv(bencher: divan::Bencher) { }); } +/// Steady state: Bob has already ratcheted into Alice's current sending chain, +/// so this decrypt is symmetric-only — no curve operation at all. This is what +/// an active conversation actually pays per message, and it is roughly an order +/// of magnitude cheaper than the ratchet step `bench_dm_recv` measures. +fn setup_dm_recv_steady() -> DmRecvData { + let mut alice = User::new("5511999000001", "s.whatsapp.net"); + let mut bob = User::new("5511999000002", "s.whatsapp.net"); + establish_bidirectional(&mut alice, &mut bob); + + let bob_addr = bob.address.clone(); + + // Burn the ratchet step on a throwaway message, so the one we hand back is + // decrypted with the chain already advanced. + let (first, first_type) = encrypt_one(&mut alice, &bob_addr); + decrypt_dm(&first, &first_type, &alice.address, &mut bob); + + let (ciphertext, enc_type) = encrypt_one(&mut alice, &bob_addr); + + DmRecvData { + bob, + alice_addr: alice.address, + ciphertext, + enc_type, + } +} + +#[divan::bench] +fn bench_dm_recv_steady(bencher: divan::Bencher) { + bencher.with_inputs(setup_dm_recv_steady).bench_refs(|d| { + black_box(decrypt_dm( + &d.ciphertext, + &d.enc_type, + &d.alice_addr, + &mut d.bob, + )); + }); +} + fn run_group_send(d: &mut GrpSendData) { let own_jid = d.alice.jid.clone(); // Warm sends (force_skdm=false) distribute no SKDM, so prepare_group_stanza @@ -835,7 +1070,7 @@ fn run_group_send(d: &mut GrpSendData) { group: &group_info, own_jid: &own_jid, own_lid: &own_jid, - account: None, + account: Some(&d.account), to: &d.group_jid, message: &d.msg, message_id: "b-grp",