map_gui/tools/
updater.rs

1use std::collections::BTreeSet;
2
3use anyhow::Result;
4use fs_err::File;
5use futures_channel::mpsc;
6
7use abstio::{DataPacks, Manifest, MapName};
8use abstutil::prettyprint_bytes;
9use widgetry::tools::{ChooseSomething, FutureLoader, PopupMsg};
10use widgetry::{EventCtx, Key, Transition};
11
12use crate::AppLike;
13
14pub fn prompt_to_download_missing_data<A: AppLike + 'static>(
15    ctx: &mut EventCtx,
16    map_name: MapName,
17    on_load: Box<dyn FnOnce(&mut EventCtx, &mut A) -> Transition<A>>,
18) -> Transition<A> {
19    let manifest = files_to_download(&map_name);
20    let bytes = manifest
21        .entries
22        .iter()
23        .map(|(_, e)| e.compressed_size_bytes)
24        .sum();
25
26    Transition::Push(ChooseSomething::new_state(
27        ctx,
28        format!(
29            "Missing data. Download {} for {}?",
30            prettyprint_bytes(bytes),
31            map_name.describe()
32        ),
33        vec![
34            widgetry::Choice::string("Yes, download"),
35            widgetry::Choice::string("Never mind").key(Key::Escape),
36        ],
37        Box::new(move |resp, ctx, _| {
38            if resp == "Never mind" {
39                return Transition::Pop;
40            }
41
42            let (outer_progress_tx, outer_progress_rx) = futures_channel::mpsc::channel(1000);
43            let (inner_progress_tx, inner_progress_rx) = futures_channel::mpsc::channel(1000);
44            Transition::Replace(FutureLoader::<A, Result<()>>::new_state(
45                ctx,
46                Box::pin(async {
47                    let result =
48                        download_files(manifest, outer_progress_tx, inner_progress_tx).await;
49                    let wrap: Box<dyn Send + FnOnce(&A) -> Result<()>> =
50                        Box::new(move |_: &A| result);
51                    Ok(wrap)
52                }),
53                outer_progress_rx,
54                inner_progress_rx,
55                "Downloading missing files",
56                Box::new(|ctx, app, maybe_result| {
57                    let error_msg = match maybe_result {
58                        Ok(Ok(())) => None,
59                        Ok(Err(err)) => Some(err.to_string()),
60                        Err(err) => Some(format!("Something went very wrong: {}", err)),
61                    };
62                    if let Some(err) = error_msg {
63                        Transition::Replace(PopupMsg::new_state(ctx, "Download failed", vec![err]))
64                    } else {
65                        on_load(ctx, app)
66                    }
67                }),
68            ))
69        }),
70    ))
71}
72
73fn files_to_download(map: &MapName) -> Manifest {
74    let mut data_packs = DataPacks {
75        runtime: BTreeSet::new(),
76        input: BTreeSet::new(),
77    };
78    data_packs.runtime.insert(map.to_data_pack_name());
79    let mut manifest = Manifest::load().filter(data_packs);
80    // Don't download files that already exist
81    manifest
82        .entries
83        .retain(|path, _| !abstio::file_exists(&abstio::path(path.strip_prefix("data/").unwrap())));
84
85    // DataPacks are an updater tool concept, but we don't want everything from that city, just the
86    // one map (and it's scenarios, prebaked_results, and maybe the city.bin overview)
87    manifest.entries.retain(|path, _| {
88        // TODO This reinvents a bit of abst_data.rs
89        let parts = path.split('/').collect::<Vec<_>>();
90        parts[4] == "city.bin"
91            || (parts[4] == "maps" && parts[5] == format!("{}.bin", map.map))
92            || (parts.len() >= 6 && parts[5] == map.map)
93    });
94
95    manifest
96}
97
98async fn download_files(
99    manifest: Manifest,
100    mut outer_progress: mpsc::Sender<String>,
101    mut inner_progress: mpsc::Sender<String>,
102) -> Result<()> {
103    let num_files = manifest.entries.len();
104    let mut messages = Vec::new();
105    let mut files_so_far = 0;
106
107    for (path, entry) in manifest.entries {
108        files_so_far += 1;
109        let local_path = abstio::path(path.strip_prefix("data/").unwrap());
110        let url = format!(
111            "https://play.abstreet.org/{}/{}.gz",
112            crate::tools::version(),
113            path
114        );
115        if let Err(err) = outer_progress.try_send(format!(
116            "Downloading file {}/{}: {} ({})",
117            files_so_far,
118            num_files,
119            url,
120            prettyprint_bytes(entry.compressed_size_bytes)
121        )) {
122            warn!("Couldn't send progress: {}", err);
123        }
124
125        match abstio::download_bytes(&url, None, &mut inner_progress)
126            .await
127            .and_then(|bytes| {
128                // TODO Instead of holding everything in memory like this, we could also try to
129                // stream the gunzipping and output writing
130                info!("Decompressing {}", path);
131                fs_err::create_dir_all(std::path::Path::new(&local_path).parent().unwrap())
132                    .unwrap();
133                let mut out = File::create(&local_path).unwrap();
134                let mut decoder = flate2::read::GzDecoder::new(&bytes[..]);
135                std::io::copy(&mut decoder, &mut out).map_err(|err| err.into())
136            }) {
137            Ok(_) => {}
138            Err(err) => {
139                let msg = format!("Problem with {}: {}", url, err);
140                error!("{}", msg);
141                messages.push(msg);
142            }
143        }
144    }
145    if !messages.is_empty() {
146        bail!("{} errors: {}", messages.len(), messages.join(", "));
147    }
148    Ok(())
149}