summaryrefslogtreecommitdiff
path: root/client/src/download.rs
blob: 8e1686a6b1b347c1a59700436a38b9ed75324e42 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
use crate::network::Network;
use anyhow::Result;
use log::debug;
use std::collections::HashSet;
use weareshared::{
    packets::{Packet, Resource},
    store::{ResourceStore, sha256},
};

pub struct Downloader {
    have: HashSet<Resource>,
    need: HashSet<Resource>,
    pending: HashSet<Resource>,
    store: ResourceStore,
}

impl Downloader {
    pub fn new(store: ResourceStore) -> Self {
        Self {
            have: HashSet::new(),
            need: HashSet::new(),
            pending: HashSet::new(),
            store,
        }
    }
    pub fn try_get(&mut self, hash: Resource) -> Result<Option<Vec<u8>>> {
        if self.have.contains(&hash) {
            self.store.get(hash)
        } else {
            self.need.insert(hash);
            Ok(None)
        }
    }
    pub fn packet(&mut self, p: &Packet) -> Result<()> {
        match p {
            Packet::RespondResource(d) => {
                let key = Resource(sha256(&d.0));
                self.store.set(&d.0)?;
                self.need.remove(&key);
                self.pending.remove(&key);
                if self.have.insert(key) {
                    debug!("have {key}");
                }
            }
            _ => (),
        }
        Ok(())
    }
    pub fn update(&mut self, network: &mut Network) {
        let mut new_pending = Vec::new();
        for n in self.need.difference(&self.pending) {
            network
                .packet_send
                .send(Packet::RequestResource(*n))
                .unwrap();
            debug!("need {n}");
            new_pending.push(*n);
        }
        self.pending.extend(new_pending);
    }
}