From 4e827e4f8ab4c5d3daf22de8753d9d585772dcdb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Garc=C3=ADa?=
 <dani-garcia@users.noreply.github.com>
Date: Wed, 12 Dec 2018 22:15:54 +0100
Subject: [PATCH] Implement better retry and use it while saving device

---
 src/db/models/attachment.rs | 34 ++++++++--------------------------
 src/db/models/device.rs     | 12 +++++++++---
 src/util.rs                 | 29 ++++++++++++++++++++++++++++-
 3 files changed, 45 insertions(+), 30 deletions(-)

diff --git a/src/db/models/attachment.rs b/src/db/models/attachment.rs
index 5cec96c7..6530fafc 100644
--- a/src/db/models/attachment.rs
+++ b/src/db/models/attachment.rs
@@ -64,33 +64,15 @@ impl Attachment {
     }
 
     pub fn delete(self, conn: &DbConn) -> QueryResult<()> {
-        use crate::util;
-        use std::{thread, time};
+        crate::util::retry(
+            || {
+                diesel::delete(attachments::table.filter(attachments::id.eq(&self.id)))
+                    .execute(&**conn)
+            },
+            10,
+        )?;
 
-        let mut retries = 10;
-
-        loop {
-            match diesel::delete(
-                attachments::table.filter(
-                    attachments::id.eq(&self.id)
-                )
-            ).execute(&**conn) {
-                Ok(_) => break,
-                Err(err) => {
-                    if retries < 1 {
-                        error!("Failed with 10 retries");
-                        return Err(err)
-                    } else {
-                        retries -= 1;
-                        info!("Had to retry! Retries left: {}", retries);
-                        thread::sleep(time::Duration::from_millis(500));
-                        continue
-                    }
-                }
-            }
-        }
-
-        util::delete_file(&self.get_file_path());
+        crate::util::delete_file(&self.get_file_path());
         Ok(())
     }
 
diff --git a/src/db/models/device.rs b/src/db/models/device.rs
index b415d5e7..7f2552e4 100644
--- a/src/db/models/device.rs
+++ b/src/db/models/device.rs
@@ -101,7 +101,6 @@ impl Device {
             amr: vec!["Application".into()],
         };
 
-
         (encode_jwt(&claims), DEFAULT_VALIDITY.num_seconds())
     }
 }
@@ -116,8 +115,15 @@ impl Device {
     pub fn save(&mut self, conn: &DbConn) -> QueryResult<()> {
         self.updated_at = Utc::now().naive_utc();
 
-        diesel::replace_into(devices::table)
-            .values(&*self).execute(&**conn).and(Ok(()))
+        crate::util::retry(
+            || {
+                diesel::replace_into(devices::table)
+                    .values(&*self)
+                    .execute(&**conn)
+            },
+            10,
+        )
+        .and(Ok(()))
     }
 
     pub fn delete(self, conn: &DbConn) -> QueryResult<()> {
diff --git a/src/util.rs b/src/util.rs
index e43dbcac..a4d96a3c 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -252,6 +252,33 @@ fn upcase_value(value: &Value) -> Value {
 fn _process_key(key: &str) -> String {
     match key.to_lowercase().as_ref() {
         "ssn" => "SSN".into(),
-        _ => self::upcase_first(key)
+        _ => self::upcase_first(key),
+    }
+}
+
+//
+// Retry methods
+//
+
+pub fn retry<F, T, E>(func: F, max_tries: i32) -> Result<T, E>
+where
+    F: Fn() -> Result<T, E>,
+{
+    use std::{thread::sleep, time::Duration};
+    let mut tries = 0;
+
+    loop {
+        match func() {
+            ok @ Ok(_) => return ok,
+            err @ Err(_) => {
+                tries += 1;
+
+                if tries >= max_tries {
+                    return err;
+                }
+
+                sleep(Duration::from_millis(500));
+            }
+        }
     }
 }