OXYGENHUB

Coverage / crates/jet/src/swarm.rs 91.4%

308 of 337 lines run, 82 of 94 functions · run #210 on main · run missed

1
//! The swarm (oxygenflow#47, docs/jetstream.md §8.1): runners fetch image
2
//! chunks from each other as well as from the hub.
3
//!
4
//! This module holds what both ends agree on and what can be tested
5
//! without a network: the have-map a peer offers ([`Bits`]), the ticket the
6
//! hub signs to admit a runner to its peers ([`Ticket`]), and the order in
7
//! which a fetcher asks its sources for chunks ([`Picker`]): rarest first,
8
//! the hub only for what no peer has, and an endgame at the close.
9
10
use oxy_core::hash::KeyId;
11
use oxy_core::identity::{verify_challenge, Identity, SecretKey, Signature};
12
use serde::{Deserialize, Serialize};
13
use std::collections::HashMap;
14
use std::time::{Duration, Instant};
15
16
/// How long a ticket admits a runner to its peers.
17
pub const TICKET_LIFE_MS: u64 = 15 * 60 * 1000;
18
/// Peers make no progress for this long and the hub is asked for anything.
19
pub const STALL: Duration = Duration::from_secs(2);
20
/// The most neighbouring chunks a fetcher takes in order where rarity does
21
/// not decide (about 5.6 MB of an image's 88 KB chunks); a small image's
22
/// runs are shorter, so it still has a few hundred to spread runners over.
23
const RUN: usize = 64;
24
/// Most sources asked for the same chunk in the endgame.
25
const ENDGAME_COPIES: u8 = 2;
26
27
/// One bit per chunk of a manifest's needed list, in its order.
28
#[derive(Clone, Debug, Default, PartialEq, Eq)]
29
pub struct Bits {
30
    words: Vec<u64>,
31
    len: usize,
32
}
33
34
impl Bits {
3576×
    pub fn new(len: usize) -> Bits {
3676×
        Bits { words: vec![0; len.div_ceil(64)], len }
3776×
    }
38
39
    /// Every bit set: a peer that holds the whole image.
407×
    pub fn full(len: usize) -> Bits {
417×
        let mut b = Bits::new(len);
4216593×
        (0..len).for_each(|i| b.set(i));
437×
        b
447×
    }
45
467×
    pub fn len(&self) -> usize {
477×
        self.len
487×
    }
49
500×
    pub fn is_empty(&self) -> bool {
510×
        self.len == 0
520×
    }
53
54101892×
    pub fn get(&self, i: usize) -> bool {
55101892×
        i < self.len && self.words[i / 64] & (1 << (i % 64)) != 0
56101892×
    }
57
5817296×
    pub fn set(&mut self, i: usize) {
5917296×
        if i < self.len {
6017296×
            self.words[i / 64] |= 1 << (i % 64);
6117296×
        }
6217296×
    }
63
64228×
    pub fn clear(&mut self, i: usize) {
65228×
        if i < self.len {
66228×
            self.words[i / 64] &= !(1 << (i % 64));
67228×
        }
68228×
    }
69
701441×
    pub fn count(&self) -> usize {
7118018×
        self.words.iter().map(|w| w.count_ones() as usize).sum()
721441×
    }
73
74
    /// Little-endian bytes, as `GET …/have` sends them.
7520×
    pub fn to_bytes(&self) -> Vec<u8> {
7622×
        let mut out: Vec<u8> = self.words.iter().flat_map(|w| w.to_le_bytes()).collect();
7720×
        out.truncate(self.len.div_ceil(8));
7820×
        out
7920×
    }
80
81
    /// The bits of a `have` answer for a list of `len` chunks; bits past
82
    /// `len` are ignored, and a short answer means the rest are unset.
8321×
    pub fn from_bytes(len: usize, bytes: &[u8]) -> Bits {
8421×
        let mut b = Bits::new(len);
85101×
        for (i, byte) in bytes.iter().enumerate().take(len.div_ceil(8)) {
86101×
            b.words[i / 8] |= (*byte as u64) << (8 * (i % 8));
87101×
        }
8821×
        if let Some(last) = b.words.last_mut() {
8921×
            if !len.is_multiple_of(64) {
9021×
                *last &= (1u64 << (len % 64)) - 1;
9121×
            }
920×
        }
9321×
        b
9421×
    }
95
}
96
97
/// What the hub signs to admit a runner to its peers.
98
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
99
pub struct Ticket {
100
    /// The hub's key, hex: a peer accepts only its own hub's tickets.
101
    pub hub: String,
102
    pub runner: i64,
103
    pub name: String,
104
    pub site: String,
105
    pub expires_ms: u64,
106
}
107
108
impl Ticket {
109
    /// The ticket as text: its JSON and the hub's signature over it, hex,
110
    /// joined by a dot.
11118×
    pub fn seal(&self, key: &KeyId, secret: &SecretKey) -> String {
11218×
        let body = serde_json::to_vec(self).expect("a ticket serialises");
11318×
        let sig = secret.sign_challenge(key, &body).expect("a ticket is longer than 16 bytes");
11418×
        format!("{}.{}", hex::encode(&body), hex::encode(Signature::encode_list(&[sig])))
11518×
    }
116
117
    /// Read and check a ticket: signed by `hub`, not expired at `now_ms`.
11811×
    pub fn open(text: &str, hub: &Identity, now_ms: u64) -> Result<Ticket, String> {
11911×
        let (body, sig) = text.trim().split_once('.').ok_or("not a ticket")?;
12010×
        let body = hex::decode(body).map_err(|_| "not a ticket")?;
12110×
        let sigs = hex::decode(sig).ok().and_then(|b| Signature::decode_list(&b).ok()).ok_or("not a ticket's signature")?;
12210×
        let sig = sigs.first().ok_or("an unsigned ticket")?;
12310×
        verify_challenge(hub, sig, &body).map_err(|_| "the ticket is not signed by this runner's hub")?;
1248×
        let t: Ticket = serde_json::from_slice(&body).map_err(|_| "not a ticket")?;
1258×
        if t.hub != hub.hash().to_hex() {
1260×
            return Err("the ticket names another hub".into());
1278×
        }
1288×
        if t.expires_ms <= now_ms {
1291×
            return Err("the ticket has expired".into());
1307×
        }
1317×
        Ok(t)
13211×
    }
133
}
134
135
/// Where a chunk comes from: the hub, or a peer by its index.
136
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
137
pub enum Source {
138
    Hub,
139
    Peer(usize),
140
}
141
142
/// What a source should do next.
143
#[derive(Debug, PartialEq, Eq)]
144
pub enum Pick {
145
    /// Ask for this chunk (an index into the needed list).
146
    Chunk(usize),
147
    /// Nothing for this source now; ask again shortly.
148
    Wait,
149
    /// Every chunk is here.
150
    Done,
151
}
152
153
/// The order a fetcher asks its sources for chunks (docs/jetstream.md
154
/// §8.1): the rarest chunk a peer holds first; the hub for chunks no peer
155
/// holds, or for anything once the peers stall; and in the endgame, a
156
/// second source for a chunk already asked for.
157
pub struct Picker {
158
    have: Bits,
159
    asked: Vec<u8>,
160
    /// The source first asked for each chunk in flight: the endgame asks
161
    /// a different one.
162
    first: Vec<Option<Source>>,
163
    /// Each peer's map: what it has, and what it is fetching now.
164
    peers: HashMap<usize, (Bits, Bits)>,
165
    /// Missing chunks, rarest first; rebuilt when the peers' maps change.
166
    order: Vec<usize>,
167
    /// Where each chunk sits in `order`.
168
    pos: Vec<usize>,
169
    /// No chunk before this place in `order` can be picked fresh: each is
170
    /// here or asked for. Moved back when one fails.
171
    skip: usize,
172
    /// How many peers hold each chunk, as of the last rebuild.
173
    holders: Vec<u16>,
174
    /// Whether some peer is fetching each chunk now: the hub is not asked
175
    /// for it, the peer will have it shortly.
176
    coming: Vec<bool>,
177
    progress: Instant,
178
    seed: u64,
179
}
180
181
impl Picker {
182
    /// A picker for `len` chunks, of which `have` are already here.
18316×
    pub fn new(len: usize, have: Bits, seed: u64) -> Picker {
18416×
        let mut p = Picker {
18516×
            have,
18616×
            asked: vec![0; len],
18716×
            first: vec![None; len],
18816×
            peers: HashMap::new(),
18916×
            order: Vec::new(),
19016×
            pos: vec![0; len],
19116×
            skip: 0,
19216×
            holders: vec![0; len],
19316×
            coming: vec![false; len],
19416×
            progress: Instant::now(),
19516×
            seed: seed | 1,
19616×
        };
19716×
        p.reorder();
19816×
        p
19916×
    }
200
2011546×
    pub fn len(&self) -> usize {
2021546×
        self.asked.len()
2031546×
    }
204
2050×
    pub fn is_empty(&self) -> bool {
2060×
        self.asked.is_empty()
2070×
    }
208
2091438×
    pub fn missing(&self) -> usize {
2101438×
        self.len() - self.have.count()
2111438×
    }
212
2130×
    pub fn have(&self) -> &Bits {
2140×
        &self.have
2150×
    }
216
217
    /// A peer's have-map, new or refreshed.
2187×
    pub fn peer(&mut self, peer: usize, bits: Bits) {
2197×
        let none = Bits::new(bits.len());
2207×
        self.peer_fetching(peer, bits, none);
2217×
    }
222
223
    /// A peer's have-map and the chunks it is fetching now.
22419×
    pub fn peer_fetching(&mut self, peer: usize, have: Bits, fetching: Bits) {
22519×
        self.peers.insert(peer, (have, fetching));
22619×
        self.reorder();
22719×
    }
228
229
    /// A peer that failed: it is asked for nothing more.
2301×
    pub fn drop_peer(&mut self, peer: usize) {
2311×
        self.peers.remove(&peer);
2321×
        self.reorder();
2331×
    }
234
2351354×
    fn next_rand(&mut self) -> u64 {
236
        // xorshift64: enough to spread runners that start together.
2371354×
        self.seed ^= self.seed << 13;
2381354×
        self.seed ^= self.seed >> 7;
2391354×
        self.seed ^= self.seed << 17;
2401354×
        self.seed
2411354×
    }
242
243
    /// Rarest first; among chunks as rare as each other, runs of up to
244
    /// [`RUN`] neighbours in order, the runs shuffled. A peer reads them out of a
245
    /// whole image, and on a spinning disk scattered reads cost a seek each
246
    /// (Earth sent 2.3 GiB in 20 minutes, picked one chunk at a time at
247
    /// random); runners that start together still spread over the runs.
24836×
    fn reorder(&mut self) {
24936×
        let run = (self.len() / 256).clamp(1, RUN);
2501354×
        let shuffle: Vec<u64> = (0..self.len().div_ceil(run)).map(|_| self.next_rand()).collect();
25136×
        let mut keyed: Vec<(u16, u64, usize)> = Vec::new();
25233610×
        for i in 0..self.len() {
25333610×
            self.holders[i] = self.peers.values().filter(|(b, _)| b.get(i)).count().min(u16::MAX as usize) as u16;
25433610×
            self.coming[i] = self.peers.values().any(|(_, f)| f.get(i));
25533610×
            if !self.have.get(i) {
25633472×
                keyed.push((self.holders[i], shuffle[i / run], i));
25733472×
            }
258
        }
25936×
        keyed.sort_unstable();
26036×
        self.order = keyed.into_iter().map(|(_, _, i)| i).collect();
26133472×
        for (at, &i) in self.order.iter().enumerate() {
26233472×
            self.pos[i] = at;
26333472×
        }
26436×
        self.skip = 0;
26536×
    }
266
267
    /// The next chunk for `src` to fetch.
2681294×
    pub fn pick(&mut self, src: Source) -> Pick {
2691294×
        if self.missing() == 0 {
27064×
            return Pick::Done;
2711230×
        }
2729479×
        let holds = |i: usize| match src {
273759×
            Source::Hub => true,
2748720×
            Source::Peer(p) => self.peers.get(&p).is_some_and(|(b, _)| b.get(i)),
2759479×
        };
2761230×
        let stalled = self.progress.elapsed() >= STALL;
277
        // Not yet asked for: a peer takes the rarest it holds; the hub
278
        // takes what no peer holds, or anything once the peers stall.
2791501×
        while self.skip < self.order.len() && (self.have.get(self.order[self.skip]) || self.asked[self.order[self.skip]] > 0) {
280271×
            self.skip += 1;
281271×
        }
28219213×
        let fresh = self.order[self.skip..].iter().copied().find(|&i| {
28319213×
            !self.have.get(i)
28412422×
                && self.asked[i] == 0
2859477×
                && holds(i)
286933×
                && (src != Source::Hub || stalled || self.peers.is_empty() || (self.holders[i] == 0 && !self.coming[i]))
28719213×
        });
2881230×
        if let Some(i) = fresh {
289322×
            self.asked[i] += 1;
290322×
            self.first[i] = Some(src);
291322×
            return Pick::Chunk(i);
292908×
        }
293
        // The endgame: everything missing is asked for already. The hub
294
        // joins it only once the peers stall.
295908×
        if src == Source::Hub && !stalled && !self.peers.is_empty() {
296265×
            return Pick::Wait;
297643×
        }
2985168×
        let mut left = self.order.iter().copied().filter(|&i| !self.have.get(i));
2991959×
        if left.clone().all(|i| self.asked[i] > 0) {
300110×
            if let Some(i) = left.find(|&i| self.asked[i] < ENDGAME_COPIES && self.first[i] != Some(src) && holds(i)) {
3012×
                self.asked[i] += 1;
3022×
                return Pick::Chunk(i);
30341×
            }
304600×
        }
305641×
        Pick::Wait
3061294×
    }
307
308
    /// Chunk `i` arrived and was checked.
309234×
    pub fn got(&mut self, i: usize) {
310234×
        self.have.set(i);
311234×
        self.asked[i] = self.asked[i].saturating_sub(1);
312234×
        self.first[i] = None;
313234×
        self.progress = Instant::now();
314234×
    }
315
316
    /// Chunk `i` did not come: it can be asked for again.
3173×
    pub fn failed(&mut self, i: usize) {
3183×
        self.asked[i] = self.asked[i].saturating_sub(1);
3193×
        if self.asked[i] == 0 {
3202×
            self.first[i] = None;
3212×
            if !self.have.get(i) {
3222×
                self.skip = self.skip.min(self.pos[i]);
3232×
            }
3241×
        }
3253×
    }
326
}
327
328
#[cfg(test)]
329
mod tests {
330
    #[test]
3311×
    fn from_one_peer_neighbouring_chunks_come_in_order() {
3321×
        let n = RUN * 256;
3331×
        let mut p = Picker::new(n, Bits::new(n), 3);
3341×
        p.peer(0, Bits::full(n));
3351×
        let first: Vec<usize> = (0..RUN)
33664×
            .map(|_| match p.pick(Source::Peer(0)) {
33764×
                Pick::Chunk(i) => i,
3380×
                other => panic!("{other:?}"),
33964×
            })
3401×
            .collect();
3411×
        assert_eq!(first[0] % RUN, 0, "a run starts at its first chunk: {first:?}");
34263×
        assert!(first.windows(2).all(|w| w[1] == w[0] + 1), "one run, in order: {first:?}");
343
        // A failed chunk is asked for again before going further.
3441×
        p.failed(first[3]);
3451×
        assert_eq!(p.pick(Source::Peer(0)), Pick::Chunk(first[3]));
3461×
    }
347
348
    /// A fetch the size of the 14 GB Rust image (142,244 chunks) from one
349
    /// peer that holds them all: how long picking takes.
350
    #[test]
351
    #[ignore]
3520×
    fn picking_a_big_image_from_one_peer_is_quick() {
3530×
        let n = 142_244;
3540×
        let mut p = Picker::new(n, Bits::new(n), 7);
3550×
        p.peer(0, Bits::full(n));
3560×
        let started = std::time::Instant::now();
3570×
        let mut in_flight = std::collections::VecDeque::new();
358
        loop {
3590×
            match p.pick(Source::Peer(0)) {
3600×
                Pick::Chunk(i) => {
3610×
                    in_flight.push_back(i);
3620×
                    if in_flight.len() >= 4 {
3630×
                        p.got(in_flight.pop_front().unwrap());
3640×
                    }
365
                }
3660×
                Pick::Wait => p.got(in_flight.pop_front().unwrap()),
3670×
                Pick::Done => break,
368
            }
369
        }
3700×
        eprintln!("picked {n} chunks in {:?}", started.elapsed());
3710×
    }
372
373
    use super::*;
374
375
    #[test]
3761×
    fn bits_travel_as_bytes() {
3771×
        let mut b = Bits::new(70);
3785×
        [0, 9, 63, 64, 69].iter().for_each(|&i| b.set(i));
3791×
        let back = Bits::from_bytes(70, &b.to_bytes());
3801×
        assert_eq!(back, b);
3811×
        assert_eq!(b.to_bytes().len(), 9);
3821×
        assert_eq!(Bits::from_bytes(10, &[0xff, 0xff, 0xff]).count(), 10, "bits past the list are ignored");
3831×
        assert_eq!(Bits::from_bytes(10, &[]).count(), 0, "a short answer leaves the rest unset");
3841×
        assert_eq!(Bits::full(130).count(), 130);
3851×
    }
386
387
    #[test]
3881×
    fn a_ticket_opens_only_with_its_hub_and_in_time() {
3891×
        let sk = SecretKey::generate();
3901×
        let hub = sk.identity("hub", "hub@example", 1);
3911×
        let t = Ticket { hub: hub.hash().to_hex(), runner: 7, name: "earth-1".into(), site: "home".into(), expires_ms: 10_000 };
3921×
        let text = t.seal(&hub.hash(), &sk);
3931×
        assert_eq!(Ticket::open(&text, &hub, 9_999).unwrap(), t);
3941×
        assert!(Ticket::open(&text, &hub, 10_000).unwrap_err().contains("expired"));
3951×
        let other_sk = SecretKey::generate();
3961×
        let other = other_sk.identity("other", "o@example", 1);
3971×
        assert!(Ticket::open(&text, &other, 1).is_err(), "another hub's key");
3981×
        let mut forged = text.clone();
3991×
        forged.replace_range(2..4, "ff");
4001×
        assert!(Ticket::open(&forged, &hub, 1).is_err(), "a changed body");
4011×
        assert!(Ticket::open("nonsense", &hub, 1).is_err());
4021×
    }
403
404
    #[test]
4051×
    fn peers_are_asked_for_the_rarest_chunk_they_hold() {
4061×
        let mut p = Picker::new(4, Bits::new(4), 42);
4071×
        let mut a = Bits::new(4);
4083×
        [0, 1, 2].iter().for_each(|&i| a.set(i));
4091×
        let mut b = Bits::new(4);
4102×
        [0, 1].iter().for_each(|&i| b.set(i));
4111×
        p.peer(0, a);
4121×
        p.peer(1, b);
413
        // Chunk 2 only peer 0 holds; chunk 3 nobody does.
4141×
        assert_eq!(p.pick(Source::Peer(0)), Pick::Chunk(2));
4151×
        assert_eq!(p.pick(Source::Hub), Pick::Chunk(3), "the hub seeds what no peer holds");
4161×
        assert_eq!(p.pick(Source::Hub), Pick::Wait, "and leaves the rest to the peers");
4171×
        let first = p.pick(Source::Peer(1));
4181×
        assert!(matches!(first, Pick::Chunk(0) | Pick::Chunk(1)));
4191×
    }
420
421
    #[test]
4221×
    fn the_hub_is_not_asked_for_what_a_peer_is_fetching() {
4231×
        let mut p = Picker::new(2, Bits::new(2), 4);
4241×
        let mut fetching = Bits::new(2);
4251×
        fetching.set(0);
4261×
        p.peer_fetching(0, Bits::new(2), fetching);
4271×
        assert_eq!(p.pick(Source::Hub), Pick::Chunk(1));
4281×
        assert_eq!(p.pick(Source::Hub), Pick::Wait, "chunk 0 is coming from the peer");
4291×
        let mut have = Bits::new(2);
4301×
        have.set(0);
4311×
        p.peer(0, have);
4321×
        assert_eq!(p.pick(Source::Peer(0)), Pick::Chunk(0));
4331×
    }
434
435
    #[test]
4361×
    fn without_peers_the_hub_serves_everything() {
4371×
        let mut p = Picker::new(3, Bits::new(3), 1);
4381×
        let mut seen = Vec::new();
4394×
        while let Pick::Chunk(i) = p.pick(Source::Hub) {
4403×
            seen.push(i);
4413×
            p.got(i);
4423×
        }
4431×
        seen.sort();
4441×
        assert_eq!(seen, vec![0, 1, 2]);
4451×
        assert_eq!(p.pick(Source::Hub), Pick::Done);
4461×
    }
447
448
    #[test]
4491×
    fn the_endgame_never_asks_the_same_source_twice() {
4501×
        let mut p = Picker::new(1, Bits::new(1), 9);
4511×
        assert_eq!(p.pick(Source::Hub), Pick::Chunk(0));
4521×
        assert_eq!(p.pick(Source::Hub), Pick::Wait);
4531×
        p.peer(0, Bits::full(1));
4541×
        assert_eq!(p.pick(Source::Peer(0)), Pick::Chunk(0), "a peer may race the hub for it");
4551×
        assert_eq!(p.pick(Source::Hub), Pick::Wait, "but not the hub a peer, while the peers make progress");
4561×
    }
457
458
    #[test]
4591×
    fn the_endgame_asks_a_second_source_and_a_failure_frees_the_chunk() {
4601×
        let mut p = Picker::new(1, Bits::new(1), 3);
4611×
        p.peer(0, Bits::full(1));
4621×
        p.peer(1, Bits::full(1));
4631×
        assert_eq!(p.pick(Source::Peer(0)), Pick::Chunk(0));
4641×
        assert_eq!(p.pick(Source::Peer(1)), Pick::Chunk(0), "the endgame");
4651×
        assert_eq!(p.pick(Source::Peer(1)), Pick::Wait, "at most two copies");
4661×
        p.failed(0);
4671×
        p.failed(0);
4681×
        p.drop_peer(0);
4691×
        assert_eq!(p.pick(Source::Peer(1)), Pick::Chunk(0));
4701×
        p.got(0);
4711×
        assert_eq!(p.pick(Source::Peer(1)), Pick::Done);
4721×
    }
473
474
    #[test]
4751×
    fn chunks_already_here_are_never_asked_for() {
4761×
        let mut have = Bits::new(3);
4771×
        have.set(1);
4781×
        let mut p = Picker::new(3, have, 5);
4791×
        assert_eq!(p.missing(), 2);
4801×
        let mut seen = Vec::new();
4813×
        while let Pick::Chunk(i) = p.pick(Source::Hub) {
4822×
            seen.push(i);
4832×
            p.got(i);
4842×
        }
4851×
        assert!(!seen.contains(&1));
4861×
    }
487
488
    #[test]
4891×
    fn runners_starting_together_ask_the_hub_for_different_chunks() {
4901×
        let a: Vec<_> = {
4911×
            let mut p = Picker::new(64, Bits::new(64), 11);
4928×
            (0..8).filter_map(|_| if let Pick::Chunk(i) = p.pick(Source::Hub) { Some(i) } else { None }).collect()
493
        };
4941×
        let b: Vec<_> = {
4951×
            let mut p = Picker::new(64, Bits::new(64), 12);
4968×
            (0..8).filter_map(|_| if let Pick::Chunk(i) = p.pick(Source::Hub) { Some(i) } else { None }).collect()
497
        };
4981×
        assert_ne!(a, b);
4991×
    }
500
}