Add a Transformer iterator to apply a limiter to jittering

6 files changed, 93 insertions, 11 deletions

diff --git a/ios/MullvadREST/RetryStrategy/ExponentialBackoff.swift b/ios/MullvadREST/RetryStrategy/ExponentialBackoff.swift
index 57a54e07a9..fc0e3da2ee 100644
--- a/ios/MullvadREST/RetryStrategy/ExponentialBackoff.swift
+++ b/ios/MullvadREST/RetryStrategy/ExponentialBackoff.swift
@@ -12,9 +12,9 @@ import MullvadTypes
 struct ExponentialBackoff: IteratorProtocol {
     private var _next: Duration
     private let multiplier: UInt64
-    private let maxDelay: Duration?
+    private let maxDelay: Duration
 
-    init(initial: Duration, multiplier: UInt64, maxDelay: Duration? = nil) {
+    init(initial: Duration, multiplier: UInt64, maxDelay: Duration) {
         _next = initial
         self.multiplier = multiplier
         self.maxDelay = maxDelay
@@ -23,7 +23,7 @@ struct ExponentialBackoff: IteratorProtocol {
     mutating func next() -> Duration? {
         let next = _next
 
-        if let maxDelay, next > maxDelay {
+        if next > maxDelay {
             return maxDelay
         }
 
diff --git a/ios/MullvadREST/RetryStrategy/Jittered.swift b/ios/MullvadREST/RetryStrategy/Jittered.swift
index 271e088bbe..0f930b3aab 100644
--- a/ios/MullvadREST/RetryStrategy/Jittered.swift
+++ b/ios/MullvadREST/RetryStrategy/Jittered.swift
@@ -27,3 +27,19 @@ struct Jittered<InnerIterator: IteratorProtocol>: IteratorProtocol
         return .milliseconds(millisWithJitter)
     }
 }
+
+/// Iterator that applies a transform function to the result of another iterator.
+struct Transformer<Inner: IteratorProtocol>: IteratorProtocol {
+    typealias Element = Inner.Element
+    private var inner: Inner
+    private let transformer: (Inner.Element?) -> Inner.Element?
+
+    init(inner: Inner, transform: @escaping (Inner.Element?) -> Inner.Element?) {
+        self.inner = inner
+        self.transformer = transform
+    }
+
+    mutating func next() -> Inner.Element? {
+        transformer(inner.next())
+    }
+}
diff --git a/ios/MullvadREST/RetryStrategy/RetryStrategy.swift b/ios/MullvadREST/RetryStrategy/RetryStrategy.swift
index d7cd047b26..18e3cd69f3 100644
--- a/ios/MullvadREST/RetryStrategy/RetryStrategy.swift
+++ b/ios/MullvadREST/RetryStrategy/RetryStrategy.swift
@@ -25,7 +25,17 @@ extension REST {
             let inner = delay.makeIterator()
 
             if applyJitter {
-                return AnyIterator(Jittered(inner))
+                return switch delay {
+                case .never:
+                    AnyIterator(inner)
+                case .constant:
+                    AnyIterator(Jittered(inner))
+                case let .exponentialBackoff(_, _, maxDelay):
+                    AnyIterator(Transformer(inner: Jittered(inner)) { nextValue in
+                        guard let nextValue else { return maxDelay }
+                        return nextValue >= maxDelay ? maxDelay : nextValue
+                    })
+                }
             } else {
                 return AnyIterator(inner)
             }
@@ -68,7 +78,7 @@ extension REST {
         case constant(Duration)
 
         /// Exponential backoff.
-        case exponentialBackoff(initial: Duration, multiplier: UInt64, maxDelay: Duration?)
+        case exponentialBackoff(initial: Duration, multiplier: UInt64, maxDelay: Duration)
 
         func makeIterator() -> AnyIterator<Duration> {
             switch self {
diff --git a/ios/MullvadRESTTests/ExponentialBackoffTests.swift b/ios/MullvadRESTTests/ExponentialBackoffTests.swift
index 4fc5a99bb6..24dd00b54b 100644
--- a/ios/MullvadRESTTests/ExponentialBackoffTests.swift
+++ b/ios/MullvadRESTTests/ExponentialBackoffTests.swift
@@ -12,7 +12,7 @@ import XCTest
 
 final class ExponentialBackoffTests: XCTestCase {
     func testExponentialBackoff() {
-        var backoff = ExponentialBackoff(initial: .seconds(2), multiplier: 3)
+        var backoff = ExponentialBackoff(initial: .seconds(2), multiplier: 3, maxDelay: .seconds(18))
 
         XCTAssertEqual(backoff.next(), .seconds(2))
         XCTAssertEqual(backoff.next(), .seconds(6))
@@ -20,7 +20,7 @@ final class ExponentialBackoffTests: XCTestCase {
     }
 
     func testAtMaximumValue() {
-        var backoff = ExponentialBackoff(initial: .milliseconds(.max - 1), multiplier: 2)
+        var backoff = ExponentialBackoff(initial: .milliseconds(.max - 1), multiplier: 2, maxDelay: .seconds(.max - 1))
 
         XCTAssertEqual(backoff.next(), .milliseconds(.max - 1))
         XCTAssertEqual(backoff.next(), .milliseconds(.max))
@@ -40,20 +40,20 @@ final class ExponentialBackoffTests: XCTestCase {
     }
 
     func testMinimumValue() {
-        var backoff = ExponentialBackoff(initial: .milliseconds(0), multiplier: 10)
+        var backoff = ExponentialBackoff(initial: .milliseconds(0), multiplier: 10, maxDelay: .milliseconds(0))
 
         XCTAssertEqual(backoff.next(), .milliseconds(0))
         XCTAssertEqual(backoff.next(), .milliseconds(0))
 
-        backoff = ExponentialBackoff(initial: .milliseconds(1), multiplier: 0)
+        backoff = ExponentialBackoff(initial: .milliseconds(1), multiplier: 0, maxDelay: .zero)
 
-        XCTAssertEqual(backoff.next(), .milliseconds(1))
+        XCTAssertEqual(backoff.next(), .milliseconds(0))
         XCTAssertEqual(backoff.next(), .milliseconds(0))
     }
 
     func testJitter() {
         let initial: Duration = .milliseconds(500)
-        var iterator = Jittered(ExponentialBackoff(initial: initial, multiplier: 3))
+        var iterator = Jittered(ExponentialBackoff(initial: initial, multiplier: 3, maxDelay: .milliseconds(1500)))
 
         XCTAssertGreaterThanOrEqual(iterator.next()!, initial)
     }
diff --git a/ios/MullvadRESTTests/RetryStrategyTests.swift b/ios/MullvadRESTTests/RetryStrategyTests.swift
new file mode 100644
index 0000000000..41c7576229
--- /dev/null
+++ b/ios/MullvadRESTTests/RetryStrategyTests.swift
@@ -0,0 +1,52 @@
+//
+//  RetryStrategyTests.swift
+//  MullvadRESTTests
+//
+//  Created by Marco Nikic on 2024-06-07.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+@testable import MullvadREST
+@testable import MullvadTypes
+import XCTest
+
+class RetryStrategyTests: XCTestCase {
+    func testJitteredBackoffDoesNotGoBeyondMaxDelay() throws {
+        let maxDelay = Duration(secondsComponent: 10, attosecondsComponent: 0)
+        let retryDelay = REST.RetryDelay.exponentialBackoff(initial: .seconds(1), multiplier: 2, maxDelay: maxDelay)
+        let retry = REST.RetryStrategy(maxRetryCount: 0, delay: retryDelay, applyJitter: true)
+        let iterator = retry.makeDelayIterator()
+        var previousDelay = Duration(secondsComponent: 0, attosecondsComponent: 0)
+
+        for _ in 0 ... 10 {
+            let currentDelay = try XCTUnwrap(iterator.next())
+            XCTAssertLessThanOrEqual(previousDelay, currentDelay)
+            XCTAssertLessThanOrEqual(currentDelay, maxDelay)
+            previousDelay = currentDelay
+        }
+    }
+
+    func testJitteredConstantCannotBeMoreThanDouble() throws {
+        let retryDelay = REST.RetryDelay.constant(.seconds(10))
+        let retry = REST.RetryStrategy(maxRetryCount: 0, delay: retryDelay, applyJitter: true)
+        let iterator = retry.makeDelayIterator()
+        let minimumDelay = Duration(secondsComponent: 10, attosecondsComponent: 0)
+        let maximumDelay = Duration(secondsComponent: 20, attosecondsComponent: 0)
+
+        for _ in 0 ... 10 {
+            let currentDelay = try XCTUnwrap(iterator.next())
+            let maximumJitterRange = minimumDelay ... maximumDelay
+            print(currentDelay)
+            XCTAssertLessThanOrEqual(maximumJitterRange.lowerBound, currentDelay)
+            XCTAssertGreaterThanOrEqual(maximumJitterRange.upperBound, currentDelay)
+        }
+    }
+
+    func testCannotApplyJitterToNeverRetry() throws {
+        let retryDelay = REST.RetryDelay.never
+        let retry = REST.RetryStrategy(maxRetryCount: 0, delay: retryDelay, applyJitter: true)
+        let iterator = retry.makeDelayIterator()
+        XCTAssertNil(iterator.next())
+    }
+}
diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 1f3846a78b..bb540b34be 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -688,6 +688,7 @@
 		A917352129FAAA5200D5DCFD /* TransportStrategyTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = A917352029FAAA5200D5DCFD /* TransportStrategyTests.swift */; };
 		A91D78E32B03BDF200FCD5D3 /* TunnelObfuscation.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 5840231F2A406BF5007B27AC /* TunnelObfuscation.framework */; };
 		A91D78E42B03C01600FCD5D3 /* MullvadSettings.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 58B2FDD32AA71D2A003EB5C6 /* MullvadSettings.framework */; };
+		A91EBEDA2C1337040004A84D /* RetryStrategyTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = A91EBED92C1337040004A84D /* RetryStrategyTests.swift */; };
 		A93181A12B727ED700E341D2 /* TunnelSettingsV4.swift in Sources */ = {isa = PBXBuildFile; fileRef = A93181A02B727ED700E341D2 /* TunnelSettingsV4.swift */; };
 		A932D9EF2B5ADD0700999395 /* ProxyConfigurationTransportProvider.swift in Sources */ = {isa = PBXBuildFile; fileRef = A932D9EE2B5ADD0700999395 /* ProxyConfigurationTransportProvider.swift */; };
 		A932D9F32B5EB61100999395 /* HeadRequestTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = A932D9F22B5EB61100999395 /* HeadRequestTests.swift */; };
@@ -2029,6 +2030,7 @@
 		A91614D02B108D1B00F416EB /* TransportLayer.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TransportLayer.swift; sourceTree = "<group>"; };
 		A91614D52B10B26B00F416EB /* TunnelControlViewModel.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TunnelControlViewModel.swift; sourceTree = "<group>"; };
 		A917352029FAAA5200D5DCFD /* TransportStrategyTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TransportStrategyTests.swift; sourceTree = "<group>"; };
+		A91EBED92C1337040004A84D /* RetryStrategyTests.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = RetryStrategyTests.swift; sourceTree = "<group>"; };
 		A92962582B1F4FDB00DFB93B /* PrivacyInfo.xcprivacy */ = {isa = PBXFileReference; lastKnownFileType = text.xml; path = PrivacyInfo.xcprivacy; sourceTree = "<group>"; };
 		A92ECC202A77FFAF0052F1B1 /* TunnelSettings.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TunnelSettings.swift; sourceTree = "<group>"; };
 		A92ECC232A7802520052F1B1 /* StoredAccountData.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = StoredAccountData.swift; sourceTree = "<group>"; };
@@ -3749,6 +3751,7 @@
 				A932D9F22B5EB61100999395 /* HeadRequestTests.swift */,
 				58BDEB9E2A98F6B400F578F2 /* Mocks */,
 				58B4656F2A98C53300467203 /* RequestExecutorTests.swift */,
+				A91EBED92C1337040004A84D /* RetryStrategyTests.swift */,
 				F0164EC22B4C49D30020268D /* ShadowsocksLoaderStub.swift */,
 				A917352029FAAA5200D5DCFD /* TransportStrategyTests.swift */,
 			);
@@ -6039,6 +6042,7 @@
 				58BDEB9D2A98F69E00F578F2 /* MemoryCache.swift in Sources */,
 				58BDEB9B2A98F58600F578F2 /* TimeServerProxy.swift in Sources */,
 				A932D9F52B5EBB9D00999395 /* RESTTransportStub.swift in Sources */,
+				A91EBEDA2C1337040004A84D /* RetryStrategyTests.swift in Sources */,
 				58BDEB992A98F4ED00F578F2 /* AnyTransport.swift in Sources */,
 				A932D9F32B5EB61100999395 /* HeadRequestTests.swift in Sources */,
 			);