From a440a08a040728ed1597ee672ac4298854421708 Mon Sep 17 00:00:00 2001
From: Niels Andriesse <andriesseniels@gmail.com>
Date: Mon, 27 May 2019 09:58:34 +1000
Subject: [PATCH] Implement retrying

---
 .../src/Loki/API/LokiAPI+Convenience.swift       |  9 ---------
 SignalServiceKit/src/Loki/API/LokiAPI.swift      |  9 ++++++---
 .../Loki/Utilities/AnyPromise+Conversion.swift   | 10 ++++++++++
 .../src/Loki/Utilities/Promise+Retrying.swift    | 16 ++++++++++++++++
 4 files changed, 32 insertions(+), 12 deletions(-)
 create mode 100644 SignalServiceKit/src/Loki/Utilities/AnyPromise+Conversion.swift
 create mode 100644 SignalServiceKit/src/Loki/Utilities/Promise+Retrying.swift

diff --git a/SignalServiceKit/src/Loki/API/LokiAPI+Convenience.swift b/SignalServiceKit/src/Loki/API/LokiAPI+Convenience.swift
index 94b68a0be..c829f5a0f 100644
--- a/SignalServiceKit/src/Loki/API/LokiAPI+Convenience.swift
+++ b/SignalServiceKit/src/Loki/API/LokiAPI+Convenience.swift
@@ -46,12 +46,3 @@ internal extension Promise {
         }
     }
 }
-
-internal extension AnyPromise {
-
-    internal static func from<T : Any>(_ promise: Promise<T>) -> AnyPromise {
-        let result = AnyPromise(promise)
-        result.retainUntilComplete()
-        return result
-    }
-}
diff --git a/SignalServiceKit/src/Loki/API/LokiAPI.swift b/SignalServiceKit/src/Loki/API/LokiAPI.swift
index 760f67127..182409db0 100644
--- a/SignalServiceKit/src/Loki/API/LokiAPI.swift
+++ b/SignalServiceKit/src/Loki/API/LokiAPI.swift
@@ -6,6 +6,7 @@ import PromiseKit
     // MARK: Settings
     private static let version = "v1"
     public static let defaultMessageTTL: UInt64 = 1 * 24 * 60 * 60 * 1000
+    private static let maxRetryCount: UInt = 3
     
     // MARK: Types
     public typealias RawResponse = Any
@@ -44,7 +45,7 @@ import PromiseKit
                 let newRawMessages = removeDuplicates(from: rawMessages)
                 return parseProtoEnvelopes(from: newRawMessages)
             }
-        }.map { Set($0) }
+        }.retryingIfNeeded(maxRetryCount: maxRetryCount).map { Set($0) }
     }
     
     public static func sendSignalMessage(_ signalMessage: SignalMessage, to destination: String, timestamp: UInt64) -> Promise<Set<Promise<RawResponse>>> {
@@ -58,7 +59,8 @@ import PromiseKit
             // TODO: Send using P2P protocol
         } else {
             let parameters = lokiMessage.toJSON()
-            return getTargetSnodes(for: lokiMessage.destination).mapValues { invoke(.sendMessage, on: $0, associatedWith: lokiMessage.destination, parameters: parameters) }.map { Set($0) }
+            return getTargetSnodes(for: lokiMessage.destination).mapValues { invoke(.sendMessage, on: $0, associatedWith: lokiMessage.destination, parameters: parameters) }
+                .retryingIfNeeded(maxRetryCount: maxRetryCount).map { Set($0) }
         }
     }
     
@@ -68,7 +70,8 @@ import PromiseKit
             // TODO: Send using P2P protocol
         } else {
             let parameters: [String:Any] = [ "pubKey" : hexEncodedPublicKey ] // TODO: Figure out correct parameters
-            return getTargetSnodes(for: hexEncodedPublicKey).mapValues { invoke(.sendMessage, on: $0, associatedWith: hexEncodedPublicKey, parameters: parameters) }.map { Set($0) }
+            return getTargetSnodes(for: hexEncodedPublicKey).mapValues { invoke(.sendMessage, on: $0, associatedWith: hexEncodedPublicKey, parameters: parameters) }
+                .retryingIfNeeded(maxRetryCount: maxRetryCount).map { Set($0) }
         }
     }
     
diff --git a/SignalServiceKit/src/Loki/Utilities/AnyPromise+Conversion.swift b/SignalServiceKit/src/Loki/Utilities/AnyPromise+Conversion.swift
new file mode 100644
index 000000000..25480fd6b
--- /dev/null
+++ b/SignalServiceKit/src/Loki/Utilities/AnyPromise+Conversion.swift
@@ -0,0 +1,10 @@
+import PromiseKit
+
+internal extension AnyPromise {
+    
+    internal static func from<T : Any>(_ promise: Promise<T>) -> AnyPromise {
+        let result = AnyPromise(promise)
+        result.retainUntilComplete()
+        return result
+    }
+}
diff --git a/SignalServiceKit/src/Loki/Utilities/Promise+Retrying.swift b/SignalServiceKit/src/Loki/Utilities/Promise+Retrying.swift
new file mode 100644
index 000000000..eb340081e
--- /dev/null
+++ b/SignalServiceKit/src/Loki/Utilities/Promise+Retrying.swift
@@ -0,0 +1,16 @@
+import PromiseKit
+
+internal extension Promise {
+    
+    internal func retryingIfNeeded(maxRetryCount: UInt) -> Promise<T> {
+        var retryCount = 0
+        func retryIfNeeded() -> Promise<T> {
+            return recover { error -> Promise<T> in
+                guard retryCount != maxRetryCount else { throw error }
+                retryCount += 1
+                return retryIfNeeded()
+            }
+        }
+        return retryIfNeeded()
+    }
+}