From 8187f1e184446af2382bf76366c32f3d15f65aac Mon Sep 17 00:00:00 2001
From: Teascade <teascade@gmail.com>
Date: Mon, 24 Aug 2020 20:33:28 +0300
Subject: [PATCH] Make erronous locations requeue

---
 src/api.rs  | 37 +++++++++++++++++++++++++++++++------
 src/main.rs |  4 +++-
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/src/api.rs b/src/api.rs
index 7afcf54..758509d 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -160,7 +160,7 @@ impl API {
         &mut self,
     ) -> Receiver<Result<Result<Vec<LocationModel>, ErrorModel>, GenericError>> {
         let (sender, receiver) = mpsc::channel();
-        let locations = self.location_req_que.clone();
+        let mut locations = self.location_req_que.clone();
         let config = self.config.clone();
 
         // sleep for a second to make sure no overlaps with previous requests.
@@ -193,7 +193,7 @@ impl API {
                         let location = locations[idx].clone();
                         awaiting += 1;
                         thread::spawn(move || {
-                            let response = minreq::get(location)
+                            let response = minreq::get(location.clone())
                                 .with_header("content-type", "application/json")
                                 .with_header("x-api-key", &config.api_key)
                                 .send()
@@ -202,21 +202,29 @@ impl API {
                                     Err(_) => Ok(Err(r.json()?)),
                                 })
                                 .map_err(GenericError::from);
-                            i_sender.send(response).ok();
+                            i_sender.send((location, response)).ok();
                         });
                         timer += interval;
                         idx += 1;
                     } else {
                         while awaiting > 0 {
-                            if let Ok(rec) = i_receiver.try_recv() {
+                            if let Ok((loc, rec)) = i_receiver.try_recv() {
+                                if API::add_error_back(loc, &rec, &mut locations) {
+                                    println!("Sent back to queue");
+                                }
                                 awaiting -= 1;
                                 sender.send(rec).ok();
                             }
                         }
-                        break;
+                        if locations.len() <= idx {
+                            break;
+                        }
                     }
                 } else {
-                    if let Ok(rec) = i_receiver.try_recv() {
+                    if let Ok((loc, rec)) = i_receiver.try_recv() {
+                        if API::add_error_back(loc, &rec, &mut locations) {
+                            println!("Sent back to queue");
+                        }
                         awaiting -= 1;
                         sender.send(rec).ok();
                     }
@@ -227,6 +235,23 @@ impl API {
         receiver
     }
 
+    fn add_error_back(
+        location: String,
+        rec: &Result<Result<Vec<LocationModel>, ErrorModel>, GenericError>,
+        locations: &mut Vec<String>,
+    ) -> bool {
+        if let Ok(inner) = rec {
+            if let Err(_) = inner {
+                locations.push(location);
+                true
+            } else {
+                false
+            }
+        } else {
+            false
+        }
+    }
+
     pub fn get_between<T: Timestamped + Clone>(
         list: &Vec<T>,
         from: Option<NaiveDateTime>,
diff --git a/src/main.rs b/src/main.rs
index 145ee6c..810089a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -119,7 +119,9 @@ fn run(
         match receiver.try_recv() {
             Ok(res) => match res {
                 Ok(loc) => {
-                    counter += 1;
+                    if loc.is_ok() {
+                        counter += 1;
+                    }
                     let remaining = exp_time(&api, states.len() as u64 - counter as u64);
                     print!(
                         "Done: {:<5.2}%    Remaining: {:<5.2} seconds\r",