From f32bb85ddc72266fd40cb15f4366b53424d4c64d Mon Sep 17 00:00:00 2001
From: Niels Andriesse <andriesseniels@gmail.com>
Date: Fri, 14 Jun 2019 10:21:32 +1000
Subject: [PATCH] Remove snodes from a swarm if they failed repeatedly

---
 .../src/Loki/API/LokiAPI+SwarmAPI.swift           | 15 ++++++++++++---
 SignalServiceKit/src/Loki/API/LokiAPI.swift       | 10 ++--------
 .../Loki/Utilities/Dictionary+Description.swift   | 13 +++++++++++++
 SignalServiceKit/src/Messages/OWSMessageSender.m  |  2 +-
 4 files changed, 28 insertions(+), 12 deletions(-)
 create mode 100644 SignalServiceKit/src/Loki/Utilities/Dictionary+Description.swift

diff --git a/SignalServiceKit/src/Loki/API/LokiAPI+SwarmAPI.swift b/SignalServiceKit/src/Loki/API/LokiAPI+SwarmAPI.swift
index fbee22451..9853f0623 100644
--- a/SignalServiceKit/src/Loki/API/LokiAPI+SwarmAPI.swift
+++ b/SignalServiceKit/src/Loki/API/LokiAPI+SwarmAPI.swift
@@ -4,10 +4,12 @@ extension String : Error { }
 
 public extension LokiAPI {
     
+    fileprivate static var failureCount: [LokiAPITarget:UInt] = [:]
+    
     // MARK: Settings
     private static let minimumSnodeCount = 2
     private static let targetSnodeCount = 3
-    private static let defaultSnodePort: UInt16 = 8080
+    fileprivate static let failureThreshold = 3
     
     // MARK: Caching
     private static let swarmCacheKey = "swarmCacheKey"
@@ -100,8 +102,15 @@ internal extension Promise {
             if let error = error as? NetworkManagerError {
                 switch error.statusCode {
                 case 0:
-                    // The snode is unreachable; usually a problem with LokiNet
-                    print("[Loki] Couldn't reach snode at: \(target.address):\(target.port).")
+                    // The snode is unreachable
+                    let oldFailureCount = LokiAPI.failureCount[target] ?? 0
+                    let newFailureCount = oldFailureCount + 1
+                    LokiAPI.failureCount[target] = newFailureCount
+                    print("[Loki] Couldn't reach snode at: \(target.address):\(target.port); setting failure count to: \(newFailureCount).")
+                    if oldFailureCount >= LokiAPI.failureThreshold {
+                        print("[Loki] Failure threshold reached for: \(target); removing it from the swarm cache for: \(hexEncodedPublicKey).")
+                        LokiAPI.dropIfNeeded(target, hexEncodedPublicKey: hexEncodedPublicKey)
+                    }
                 case 421:
                     // The snode isn't associated with the given public key anymore
                     print("[Loki] Invalidating swarm for: \(hexEncodedPublicKey).")
diff --git a/SignalServiceKit/src/Loki/API/LokiAPI.swift b/SignalServiceKit/src/Loki/API/LokiAPI.swift
index 364fe7083..6dd673651 100644
--- a/SignalServiceKit/src/Loki/API/LokiAPI.swift
+++ b/SignalServiceKit/src/Loki/API/LokiAPI.swift
@@ -44,14 +44,8 @@ public final class LokiAPI : NSObject {
         if let headers = headers { request.allHTTPHeaderFields = headers }
         if let timeout = timeout { request.timeoutInterval = timeout ?? defaultTimeout }
         let headers = request.allHTTPHeaderFields ?? [:]
-        let headersDescription = headers.isEmpty ? "no custom headers specified" : headers.description
-        let parametersDescription = "[ " + parameters.map { key, value in
-            let valueDescription = String(describing: value)
-            let maxLength = 20
-            let truncatedValueDescription = valueDescription.count > maxLength ? valueDescription.prefix(maxLength) + "..." : valueDescription
-            return key + " : " + truncatedValueDescription
-        }.joined(separator: ", ") + " ]"
-        print("[Loki] Invoking \(method.rawValue) on \(target) with \(parametersDescription) (\(headersDescription)).")
+        let headersDescription = headers.isEmpty ? "no custom headers specified" : headers.prettifiedDescription
+        print("[Loki] Invoking \(method.rawValue) on \(target) with \(parameters.prettifiedDescription) (\(headersDescription)).")
         return TSNetworkManager.shared().makePromise(request: request).map { $0.responseObject }
             .handlingSwarmSpecificErrorsIfNeeded(for: target, associatedWith: hexEncodedPublicKey).recoveringNetworkErrorsIfNeeded()
     }
diff --git a/SignalServiceKit/src/Loki/Utilities/Dictionary+Description.swift b/SignalServiceKit/src/Loki/Utilities/Dictionary+Description.swift
new file mode 100644
index 000000000..c7aedbfe6
--- /dev/null
+++ b/SignalServiceKit/src/Loki/Utilities/Dictionary+Description.swift
@@ -0,0 +1,13 @@
+
+public extension Dictionary {
+    
+    public var prettifiedDescription: String {
+        return "[ " + map { key, value in
+            let keyDescription = String(describing: key)
+            let valueDescription = String(describing: value)
+            let maxLength = 20
+            let truncatedValueDescription = valueDescription.count > maxLength ? valueDescription.prefix(maxLength) + "..." : valueDescription
+            return keyDescription + " : " + truncatedValueDescription
+        }.joined(separator: ", ") + " ]"
+    }
+}
diff --git a/SignalServiceKit/src/Messages/OWSMessageSender.m b/SignalServiceKit/src/Messages/OWSMessageSender.m
index 62df07abc..27148f8bf 100644
--- a/SignalServiceKit/src/Messages/OWSMessageSender.m
+++ b/SignalServiceKit/src/Messages/OWSMessageSender.m
@@ -1186,7 +1186,7 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
                 }) retainUntilComplete];
             }
         })
-        .catchOn(OWSDispatch.sendingQueue, ^(NSError *error) { // Unreachable snode; usually a problem with LokiNet
+        .catchOn(OWSDispatch.sendingQueue, ^(NSError *error) { // The snode is unreachable
             handleError(error);
         }) retainUntilComplete];