diff --git a/wacore/benches/send_receive_benchmark.rs b/wacore/benches/send_receive_benchmark.rs index ffd97155c..4035b65cd 100644 --- a/wacore/benches/send_receive_benchmark.rs +++ b/wacore/benches/send_receive_benchmark.rs @@ -655,7 +655,16 @@ fn setup_dm_recv() -> DmRecvData { struct GrpSendData { alice: User, group_jid: Jid, - participants: Vec, + /// Built in setup, not per iteration: production resolves the group once + /// and holds the result behind an `Arc` across sends (`ensure_self_in_group` + /// hands the same `Arc` straight back whenever we are already a member, the + /// steady state), so a send never constructs or drops a participant list. + /// Building it in the measured body charged every group send an + /// N-participant construct + teardown that no send performs — 26.8K + /// instructions at 512 members, and the entire reason this benchmark + /// appeared to scale with group size while `prepare_group_stanza` itself is + /// flat (334.0K at 8 members, 334.1K at 512). + group_info: GroupInfo, /// Warm-send fixture: the resolved set with its phash memo pre-warmed in /// setup, like the per-group device memo serves production repeat sends. resolved_for_phash: Option>, @@ -705,10 +714,18 @@ fn setup_group_send(n: usize) -> GrpSendData { .phash(&alice.jid) .expect("phash must warm in setup"); + // Self-append happens once here for the same reason production does it once + // per resolution: `prepare_group_stanza` expects the sender in the list. + let own_base = alice.jid.to_non_ad(); + if !participants.iter().any(|p| p.is_same_user_as(&own_base)) { + participants.push(own_base); + } + let group_info = GroupInfo::new(participants, AddressingMode::Pn); + GrpSendData { alice, group_jid, - participants, + group_info, resolved_for_phash: Some(resolved), force_skdm: false, resolver: MockResolver(devices), @@ -1045,15 +1062,7 @@ fn run_group_send(d: &mut GrpSendData) { // only emits a phash if it gets the full device set. Mirror the real // warm-send caller by passing it; the cold/force_skdm path resolves the set // itself and keeps None. - let mut group_info = GroupInfo::new(std::mem::take(&mut d.participants), AddressingMode::Pn); - let own_base = own_jid.to_non_ad(); - if !group_info - .participants - .iter() - .any(|p| p.is_same_user_as(&own_base)) - { - group_info.participants.push(own_base); - } + let group_info = &d.group_info; let mut stores = SignalStores { sender_key_store: &mut d.alice.sender_keys, session_store: &mut d.alice.sessions, @@ -1067,7 +1076,7 @@ fn run_group_send(d: &mut GrpSendData) { &mut stores, &d.resolver, GroupStanzaRequest { - group: &group_info, + group: group_info, own_jid: &own_jid, own_lid: &own_jid, account: Some(&d.account), diff --git a/wacore/binary/Cargo.toml b/wacore/binary/Cargo.toml index 716daa462..a182d042d 100644 --- a/wacore/binary/Cargo.toml +++ b/wacore/binary/Cargo.toml @@ -48,3 +48,7 @@ harness = false [lints] workspace = true + +[[bench]] +name = "group_fanout_benchmark" +harness = false diff --git a/wacore/binary/benches/group_fanout_benchmark.rs b/wacore/binary/benches/group_fanout_benchmark.rs new file mode 100644 index 000000000..bc17ac006 --- /dev/null +++ b/wacore/binary/benches/group_fanout_benchmark.rs @@ -0,0 +1,96 @@ +//! Encoder cost of the group sender-key distribution fan-out, swept across +//! recipient counts. +//! +//! Its own target rather than a section of `binary_benchmark`: the fixture +//! below instantiates the typed-JID attribute path, and adding it to that +//! crate root changed inlining enough to cost `create_attr_node`'s builder an +//! extra `SmallVec` reallocation — ~300 instructions on `bench_attr_parser`, +//! a benchmark with no connection to this one. A separate crate root keeps a +//! new fixture from moving an unrelated baseline. + +use divan::black_box; +use wacore_binary::builder::NodeBuilder; +use wacore_binary::jid::Jid; +use wacore_binary::marshal::marshal_exact; +use wacore_binary::node::Node; + +fn main() { + divan::main(); +} + +/// The `` shape a sender-key distribution puts on the wire: one +/// `` per recipient device, each wrapping its own ``. This is +/// the only group stanza whose encode cost is proportional to the recipient +/// count — a steady-state group send distributes to our own companions only +/// and carries nothing per member (pinned by +/// `warm_group_stanza_size_tracks_own_devices_not_group_size` in `wacore`), +/// which is why the width is swept here rather than assumed. +/// +/// Recipients are a typed [`Jid`] per device, as `build_participant_node` +/// passes them, and are spread over distinct users with a handful of devices +/// each, the way a real fanout resolves. Both details decide the encoding: +/// a typed JID skips the string classifier production never runs here, and a +/// device id must stay within `u8` to take the `AD_JID` path — a single user +/// numbered up to 511 would silently encode half the sweep as `JID_PAIR` and +/// measure two wire shapes at once. +/// +/// The ciphertexts are `type="msg"`, the shape a redistribution to devices that +/// already hold a pairwise session emits — a membership change or a rotation. +/// **This sweep does not characterize a first-contact fan-out.** Those get +/// `type="pkmsg"`, whose `PreKeySignalMessage` carries an identity key, a base +/// key and the registration id on top of the same inner message, and that +/// larger payload is paid once *per recipient* — `marshal_exact` copies every +/// payload through the writer — so it raises the slope, not the intercept. Read +/// this sweep as a lower bound there, or measure a `pkmsg` payload separately; +/// do not extrapolate the cold cost from these numbers. +fn create_skdm_fanout_node(width: usize) -> Node { + const DEVICES_PER_USER: usize = 4; + let recipients: Vec = (0..width) + .map(|i| { + let user = 5511999990000u64 + (i / DEVICES_PER_USER) as u64; + let device = (i % DEVICES_PER_USER) as u16; + NodeBuilder::new("to") + .attr("jid", Jid::pn_device(user.to_string(), device)) + .children(vec![ + NodeBuilder::new("enc") + .attr("v", "2") + .attr("type", "msg") + .bytes(vec![0xAB; 128]) + .build(), + ]) + .build() + }) + .collect(); + NodeBuilder::new("message") + .attr("to", "120363000000000001@g.us") + .attr("id", "3EB0A1B2C3D4E5F60718") + .attr("type", "text") + .children(vec![ + NodeBuilder::new("participants") + .children(recipients) + .build(), + ]) + .build() +} + +// Group sender-key distribution, swept across the recipient count reported for +// real groups. Marshalling is linear in the fan-out width, so this is what a +// redistribution — a membership change, or a rotation — pays in the encoder. +// A first-contact fan-out pays a steeper per-recipient term over the larger +// `pkmsg` payload (see the fixture), so it is not what these numbers measure. +// The steady-state send that follows carries no `` at all. +// Keeping both facts measurable is what tells a group-size regression ("the +// warm stanza grew a per-participant node") apart from a group that is merely +// redistributing. +// +// `marshal_exact`, not `marshal_auto`: every outbound stanza goes through +// `Client::marshal_node_for_send`, which picks the two-pass exact strategy. +// The two differ in exactly what this sweep is measuring — one-pass reserves +// and grows, two-pass plans the size first and replays a hint tape — so the +// wrong one would track a path no group send takes. +#[divan::bench(args = [8, 32, 128, 512])] +fn bench_marshal_exact_group_fanout(bencher: divan::Bencher, width: usize) { + bencher + .with_inputs(|| create_skdm_fanout_node(width)) + .bench_refs(|node| black_box(marshal_exact(black_box(node)).unwrap())); +} diff --git a/wacore/src/send/tests.rs b/wacore/src/send/tests.rs index 56b0c893a..dd595da01 100644 --- a/wacore/src/send/tests.rs +++ b/wacore/src/send/tests.rs @@ -3921,6 +3921,253 @@ mod mark_full_distribution_list { not aborting the whole cohort" ); } + + /// A steady-state group stanza's size tracks our OWN device count, never + /// the group's. + /// + /// `` carries sender-key distributions only. A warm send + /// distributes none to members — but it does re-distribute to our own + /// companions on every send, because own devices are never memoized warm + /// (WA Web `!isMeDevice`, see `update_sender_key_devices`). So the steady + /// state is one `` per own companion and nothing per member, and the + /// `phash` covering the whole device set is a fixed-width digest ("2:" plus + /// 8 base64 chars) memoized on the resolved set. Both the single-device and + /// the multi-device steady state are pinned below at 8 and at 512 members: + /// the encoded stanza is the same size either way, so a repeat group send + /// has no per-member encoding to cache between sends. + /// + /// Pinned as a test rather than left to the group benchmarks because the + /// claim is about the *shape* of the stanza: a future change that folded + /// member state into it would still benchmark fine on a small group. + #[tokio::test] + async fn warm_group_stanza_size_tracks_own_devices_not_group_size() { + // Our own other devices, which receive a fresh SKDM on every send. + // Shared with the assertions so they can name the exact JIDs the stanza + // must address, not merely how many. + fn companion_jids(companions: usize) -> Vec { + (1..=companions) + .map(|d| format!("12025550111:{d}@s.whatsapp.net").parse().unwrap()) + .collect() + } + + // `members` is the group; `companions` are our own other devices. + async fn warm_stanza(members: usize, companions: usize) -> Node { + let own_jid: Jid = "12025550111:0@s.whatsapp.net".parse().unwrap(); + let own_lid: Jid = "100000000000001:0@lid".parse().unwrap(); + let group: Jid = "120363000000000001@g.us".parse().unwrap(); + + let participants: Vec = (0..members) + .map(|i| { + format!("{}@s.whatsapp.net", 12025550200u64 + i as u64) + .parse() + .unwrap() + }) + .collect(); + let own_companions: Vec = companion_jids(companions); + + let mut rng = rand::make_rng::(); + let mut sks = MemSenderKeyStore::default(); + // A warm send never creates the chain, so seed it exactly as the + // first (cold) send to this group would have. + let sk_name = make_sender_key_name(&group, &own_jid.to_protocol_address()); + crate::libsignal::protocol::create_sender_key_distribution_message( + &sk_name, &mut sks, &mut rng, + ) + .await + .expect("seed the sender key chain"); + + // Sessions already exist for the companions, as they do in the + // steady state, so the SKDM encrypts to `msg` (not `pkmsg`) and no + // prekey fetch or device-identity node enters the stanza. + let mut ss = MemSessionStore::default(); + let mut is = MemIdentityStore { + pair: IdentityKeyPair::generate(&mut rng), + reg_id: 7, + known: Default::default(), + }; + for companion in &own_companions { + let addr = companion.to_protocol_address(); + process_prekey_bundle( + &addr, + &mut ss, + &mut is, + &signed_prekey_bundle(), + &mut rng, + UsePQRatchet::No, + ) + .await + .expect("establish the companion session"); + // `process_prekey_bundle` alone leaves the session holding a + // pending pre-key, so its next encryption is still a `pkmsg` + // first contact. The steady state this fixture models is the + // one after the companion has answered, which is what clears + // the pending key — so clear it, and let the `enc type` + // assertion below hold the fixture to it. + let mut record = ss + .load_session(&addr) + .await + .expect("load") + .expect("session present"); + record + .session_state_mut() + .expect("session state") + .clear_unacknowledged_pre_key_message(); + ss.store_session(&addr, record).await.expect("store"); + } + let mut pks = UnusedPreKeyStore; + let spks = UnusedSignedPreKeyStore; + let mut stores = SignalStores { + sender_key_store: &mut sks, + session_store: &mut ss, + identity_store: &mut is, + prekey_store: &mut pks, + signed_prekey_store: &spks, + }; + + let mut group_participants = participants.clone(); + group_participants.push(own_jid.to_non_ad()); + let group_info = GroupInfo::new(group_participants, AddressingMode::Pn); + // The full resolved device set the warm send hashes into `phash`. + // The companions belong inside it, not beside it: production filters + // the SKDM targets out of this very set (`filter_skdm_targets` over + // `all_devices_for_phash`), and the server validates the phash against + // every recipient device — so a stanza whose `` named a + // device the phash did not cover is a shape no send produces. + let mut resolved_devices = participants; + resolved_devices.extend(own_companions.iter().cloned()); + let resolved = ResolvedGroupDevices::new(resolved_devices); + let msg = wa::Message { + conversation: Some("steady state".into()), + ..Default::default() + }; + + prepare_group_stanza( + &TokioTestRuntime, + &mut stores, + &MockSendContextResolver::new(), + GroupStanzaRequest { + group: &group_info, + own_jid: &own_jid, + own_lid: &own_lid, + account: None, + to: &group, + message: &msg, + message_id: "WARMGROUPSCALE1", + force_distribution: false, + distribution_targets: (!own_companions.is_empty()) + .then(|| own_companions.clone()), + distribution_policy: SenderKeyDistributionPolicy::BestEffort, + phash_devices: Some(&resolved), + edit: None, + extra_nodes: &[], + pre_encoded: None, + }, + ) + .await + .expect("warm group send") + .node + } + + // Every ciphertext in the stanza varies in length run to run (WA pads + // each plaintext by a random 1..=16 bytes), so sizes are only + // comparable with the payloads normalised. What is under test is the + // stanza's structure and attributes, not the ciphertext. + fn with_fixed_payloads(node: &Node) -> Node { + use wacore_binary::node::NodeContent; + let mut out = node.clone(); + out.content = match out.content { + Some(NodeContent::Bytes(_)) => Some(NodeContent::Bytes(vec![0u8; 96])), + Some(NodeContent::Nodes(children)) => Some(NodeContent::Nodes( + children.iter().map(with_fixed_payloads).collect(), + )), + other => other, + }; + out + } + + // The whole hierarchy, not just the root's children: a `` or `` + // subtree that grew with the group would otherwise slip past, and a + // rename that happens to preserve the encoded length would slip past the + // size comparison too. Attribute *keys* only — the values legitimately + // differ (the phash digests two different device sets), and the phash is + // asserted on its own below. + fn shape(node: &Node) -> String { + let mut attrs: Vec<&str> = node.attrs.0.iter().map(|(k, _)| k.as_ref()).collect(); + attrs.sort_unstable(); + let children: Vec = node.children().unwrap_or(&[]).iter().map(shape).collect(); + format!("{}[{}]({})", node.tag, attrs.join(","), children.join(" ")) + } + + // Single-device account (no companions) and a two-companion one: the + // two steady states this client actually produces. + for companions in [0usize, 2] { + let small = warm_stanza(8, companions).await; + let large = warm_stanza(512, companions).await; + + for (label, node) in [("8-member", &small), ("512-member", &large)] { + // The JIDs, not just how many: a list of the right length that + // addressed group members instead of our companions would be + // exactly the regression this test exists to catch. + let distributed: Vec = node + .get_optional_child("participants") + .and_then(Node::children) + .unwrap_or(&[]) + .iter() + .map(|to| to.attrs().jid("jid")) + .collect(); + assert_eq!( + distributed, + companion_jids(companions), + "{label} warm send distributes to our own companions only, \ + never to the group's members" + ); + // The enc type is the whole premise of the fixture, so it is + // checked rather than asserted in a comment. + for to in node + .get_optional_child("participants") + .and_then(Node::children) + .unwrap_or(&[]) + { + let enc = to + .get_optional_child("enc") + .unwrap_or_else(|| panic!("{label} participant carries an enc")); + assert_eq!( + enc.attrs().optional_string("type").as_deref(), + Some("msg"), + "{label} companion SKDM ciphertext type" + ); + } + // Version tag plus 8 base64 chars — the width is what makes the + // stanza size independent of the set hashed, and the `2:` is + // what makes it the phash the server expects rather than some + // other ten-character attribute. + let phash = node + .attrs() + .optional_string("phash") + .unwrap_or_else(|| panic!("{label} warm send must carry a phash")); + assert!( + phash.starts_with("2:") && phash.len() == 10, + "{label} phash is a fixed-width v2 digest, got {phash:?}" + ); + } + + assert_eq!( + shape(&small), + shape(&large), + "same stanza shape with {companions} companions" + ); + assert_eq!( + wacore_binary::marshal::marshal(&with_fixed_payloads(&small)) + .unwrap() + .len(), + wacore_binary::marshal::marshal(&with_fixed_payloads(&large)) + .unwrap() + .len(), + "the encoded warm group stanza is the same size at 8 and 512 members \ + with {companions} companions" + ); + } + } } /// Item 3 — phash device-set construction. The set hashed is the full