diff --git a/import_osm/Cargo.lock b/import_osm/Cargo.lock index d8f913f..3367a43 100644 --- a/import_osm/Cargo.lock +++ b/import_osm/Cargo.lock @@ -83,6 +83,19 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "console" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3993e6445baa160675931ec041a5e03ca84b9c6e32a056150d3aa2bdda0a1f45" +dependencies = [ + "encode_unicode", + "lazy_static", + "libc", + "terminal_size", + "winapi", +] + [[package]] name = "cpufeatures" version = "0.1.5" @@ -120,6 +133,12 @@ dependencies = [ "generic-array 0.14.4", ] +[[package]] +name = "encode_unicode" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f" + [[package]] name = "fallible-iterator" version = "0.2.0" @@ -350,6 +369,7 @@ name = "import_osm" version = "0.0.0" dependencies = [ "geo", + "indicatif", "num_cpus", "osm_pbf_iter", "postgres", @@ -357,6 +377,18 @@ dependencies = [ "serde_json", ] +[[package]] +name = "indicatif" +version = "0.16.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d207dc617c7a380ab07ff572a6e52fa202a2a8f355860ac9c38e23f8196be1b" +dependencies = [ + "console", + "lazy_static", + "number_prefix", + "regex", +] + [[package]] name = "instant" version = "0.1.10" @@ -479,6 +511,12 @@ dependencies = [ "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "opaque-debug" version = "0.3.0" @@ -702,6 +740,21 @@ dependencies = [ "bitflags", ] +[[package]] +name = "regex" +version = "1.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d07a8629359eb56f1e2fb1652bb04212c072a87ba68546a04065d525673ac461" +dependencies = [ + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.6.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f497285884f3fcff424ffc933e56d7cbca511def0c9831a7f9b5f6153e3cc89b" + [[package]] name = "robust" version = "0.2.3" @@ -837,6 +890,16 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "terminal_size" +version = "0.1.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "633c1a546cee861a1a6d0dc69ebeca693bf4296661ba7852b9d21d159e0506df" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "tinyvec" version = "1.3.1" diff --git a/import_osm/Cargo.toml b/import_osm/Cargo.toml index eec3e2e..de9443e 100644 --- a/import_osm/Cargo.toml +++ b/import_osm/Cargo.toml @@ -18,3 +18,4 @@ postgres = { version = "0.19", features = ["with-geo-types-0_7", "with-serde_jso serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" geo = "0.18" +indicatif = "0.16" diff --git a/import_osm/src/main.rs b/import_osm/src/main.rs index 45b90fb..2e32215 100644 --- a/import_osm/src/main.rs +++ b/import_osm/src/main.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::thread; use osm_pbf_iter::*; +use indicatif::{ProgressBar, ProgressStyle}; pub struct PrimSource { req_rx: Receiver @@ -26,6 +27,20 @@ impl PrimSource { } fn process_osm R + 'static + Send + Clone, R: Send + 'static>(filename: &str, f: F) -> Vec { + let progress = ProgressBar::new(0) + .with_style( + ProgressStyle::default_bar() + .template("{prefix:yellow} [{wide_bar:green/red}] {msg:lightblue}") + ).with_prefix(filename.to_string()); + + // open file + progress.set_message("Opening"); + let mut file = File::open(filename).unwrap(); + let filesize = file.seek(SeekFrom::End(0)).unwrap(); + progress.set_length(filesize); + file.seek(SeekFrom::Start(0)).unwrap(); + let mut reader = BlobReader::new(BufReader::new(file)); + let cpus = num_cpus::get(); let mut worker_results = Vec::with_capacity(cpus); @@ -44,27 +59,33 @@ fn process_osm R + 'static + Send + Clone, R: Send + 's }); } - // open file - println!("Open {}", filename); - let f = File::open(filename).unwrap(); - let mut reader = BlobReader::new(BufReader::new(f)); - let start = Instant::now(); - // feed + progress.set_message("Parsing"); + let start = Instant::now(); let mut w = 0; + let mut pos = 0; for blob in &mut reader { + pos += match &blob { + Blob::Raw(bytes) => bytes.len(), + Blob::Zlib(bytes) => bytes.len(), + } as u64; + let req_tx = &workers[w].0; w = (w + 1) % cpus; req_tx.send(blob).unwrap(); + + progress.set_position(pos); } // receive results + progress.set_message("Finishing"); for (req_tx, res_rx) in workers.into_iter() { drop(req_tx); let worker_res = res_rx.recv().unwrap(); worker_results.push(worker_res); } + progress.set_position(filesize); // stats let stop = Instant::now(); @@ -74,11 +95,14 @@ fn process_osm R + 'static + Send + Clone, R: Send + 's match f.seek(SeekFrom::Current(0)) { Ok(pos) => { let rate = pos as f64 / 1024f64 / 1024f64 / duration; - println!("Processed {} MB in {:.2} seconds ({:.2} MB/s)", - pos / 1024 / 1024, duration, rate); + progress.set_message(format!( + "Processed {} MB in {:.2} seconds ({:.2} MB/s)", + pos / 1024 / 1024, duration, rate + )); }, Err(_) => (), } + progress.finish(); worker_results }