SWIFT-1466 Implement updateAverageRoundTripTime for ServerDescription (#736)

Author: isabelatkinson

Sources/MongoSwift/SDAM.swift

@@ -142,6 +142,9 @@ public struct ServerDescription {
     /// The duration in milliseconds of the server's last "hello" call.
     public let roundTripTime: Int?
 
+    /// The average duration in milliseconds of the server's "hello" calls.
+    internal var averageRoundTripTimeMS: Double?
+
     /// The date of the most recent write operation seen by this server.
     public var lastWriteDate: Date?
 
@@ -194,6 +197,9 @@ public struct ServerDescription {
         self.address = ServerAddress(mongoc_server_description_host(description))
         self.serverId = mongoc_server_description_id(description)
         self.roundTripTime = Int(mongoc_server_description_round_trip_time(description))
+        // averageRoundTripTimeMS is an internally-calculated value used for server selection. Because we still rely
+        // upon libmongoc for server selection, this value is not relevant.
+        self.averageRoundTripTimeMS = nil
         self.lastUpdateTime = Date(msSinceEpoch: mongoc_server_description_last_update_time(description))
         self.type = ServerType(rawValue: String(cString: mongoc_server_description_type(description))) ?? .unknown
 
@@ -233,6 +239,7 @@ public struct ServerDescription {
         // these fields are not used by the server selection tests
         self.serverId = 0
         self.roundTripTime = 0
+        self.averageRoundTripTimeMS = nil
         self.lastUpdateTime = Date()
         self.lastWriteDate = nil
         self.minWireVersion = 0
@@ -247,6 +254,31 @@ public struct ServerDescription {
         self.passives = []
         self.arbiters = []
     }
+
+    // Used for RTT calculation tests.
+    internal init(averageRoundTripTimeMS: Double?) {
+        self.averageRoundTripTimeMS = averageRoundTripTimeMS
+
+        // these fields are not used by the RTT calculation tests
+        self.address = ServerAddress(host: "test", port: 0)
+        self.type = .unknown
+        self.tags = [:]
+        self.serverId = 0
+        self.roundTripTime = nil
+        self.lastWriteDate = nil
+        self.minWireVersion = 0
+        self.maxWireVersion = 0
+        self.me = nil
+        self.hosts = []
+        self.arbiters = []
+        self.passives = []
+        self.setName = nil
+        self.setVersion = nil
+        self.electionID = nil
+        self.primary = nil
+        self.lastUpdateTime = Date()
+        self.logicalSessionTimeoutMinutes = nil
+    }
 }
 
 extension ServerDescription: Equatable {
@@ -447,3 +479,14 @@ extension TopologyDescription {
         return []
     }
 }
+
+extension ServerDescription {
+    internal mutating func updateAverageRoundTripTime(roundTripTime: Double) {
+        if let oldAverageRTT = self.averageRoundTripTimeMS {
+            let alpha = 0.2
+            self.averageRoundTripTimeMS = alpha * roundTripTime + (1 - alpha) * oldAverageRTT
+        } else {
+            self.averageRoundTripTimeMS = roundTripTime
+        }
+    }
+}

Sources/TestsCommon/CodableExtensions.swift

@@ -205,7 +205,7 @@ extension ServerDescription: StrictDecodable {
         let address = try ServerAddress(try values.decode(String.self, forKey: .address))
         let type = try values.decode(ServerType.self, forKey: .type)
         let tags = try values.decodeIfPresent([String: String].self, forKey: .tags) ?? [:]
-        // TODO: SWIFT-1461: decode and set averageRoundTripTimeMS
+        // TODO: SWIFT-1456: decode and set averageRoundTripTimeMS
 
         self.init(address: address, type: type, tags: tags)
     }

Tests/MongoSwiftTests/ServerSelectionTests.swift

@@ -22,6 +22,26 @@ private enum OperationType: String, Decodable {
     case read, write
 }
 
+private struct RTTCalculationTestFile: Decodable {
+    let averageRoundTripTimeMS: Double?
+    let newRoundTripTimeMS: Double
+    let newAverageRoundTripTimeMS: Double
+
+    enum CodingKeys: String, CodingKey {
+        case averageRoundTripTimeMS = "avg_rtt_ms", newRoundTripTimeMS = "new_rtt_ms",
+             newAverageRoundTripTimeMS = "new_avg_rtt"
+    }
+
+    internal init(from decoder: Decoder) throws {
+        let values = try decoder.container(keyedBy: CodingKeys.self)
+        // The tests specify a non-present initial average RTT as "NULL", so if decoding to a Double fails, ignore and
+        // set to nil.
+        self.averageRoundTripTimeMS = try? values.decode(Double.self, forKey: .averageRoundTripTimeMS)
+        self.newRoundTripTimeMS = try values.decode(Double.self, forKey: .newRoundTripTimeMS)
+        self.newAverageRoundTripTimeMS = try values.decode(Double.self, forKey: .newAverageRoundTripTimeMS)
+    }
+}
+
 final class ServerSelectionTests: MongoSwiftTestCase {
     func testServerSelectionLogic() throws {
         let tests = try retrieveSpecTestFiles(
@@ -38,4 +58,18 @@ final class ServerSelectionTests: MongoSwiftTestCase {
             expect(selectedServers).to(contain(test.suitableServers))
         }
     }
+
+    func testRoundTripTimeCalculation() throws {
+        let tests = try retrieveSpecTestFiles(
+            specName: "server-selection",
+            subdirectory: "rtt",
+            asType: RTTCalculationTestFile.self
+        )
+        for (filename, test) in tests {
+            print("Running test from \(filename)...")
+            var serverDescription = ServerDescription(averageRoundTripTimeMS: test.averageRoundTripTimeMS)
+            serverDescription.updateAverageRoundTripTime(roundTripTime: test.newRoundTripTimeMS)
+            expect(serverDescription.averageRoundTripTimeMS).to(equal(test.newAverageRoundTripTimeMS))
+        }
+    }
 }