SWIFT-1465 Filter servers by max staleness (#735)

Author: isabelatkinson

Sources/MongoSwift/SDAM.swift

@@ -1,6 +1,12 @@
 import CLibMongoC
 import Foundation
 
+internal enum SDAMConstants {
+    internal static let defaultHeartbeatFrequencyMS = 10000
+    internal static let idleWritePeriodMS = 10000
+    internal static let smallestMaxStalenessSeconds = 90
+}
+
 /// A struct representing a network address, consisting of a host and port.
 public struct ServerAddress: Equatable, Hashable {
     /// The hostname or IP address.
@@ -230,20 +236,27 @@ public struct ServerDescription {
         self.tags = hello?.tags ?? [:]
     }
 
-    // For testing purposes
-    internal init(address: ServerAddress, type: ServerType, tags: [String: String]?) {
+    // Used for server selection/max staleness tests.
+    internal init(
+        address: ServerAddress,
+        type: ServerType,
+        tags: [String: String]?,
+        lastWriteDate: Date?,
+        maxWireVersion: Int?,
+        lastUpdateTime: Date?
+    ) {
         self.address = address
         self.type = type
         self.tags = tags ?? [:]
+        self.lastWriteDate = lastWriteDate
+        self.lastUpdateTime = lastUpdateTime ?? Date()
+        self.maxWireVersion = maxWireVersion ?? 0
 
         // 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
-        self.maxWireVersion = 0
         self.me = self.address
         self.setName = nil
         self.setVersion = nil
@@ -418,29 +431,42 @@ public struct TopologyDescription: Equatable {
 }
 
 extension TopologyDescription {
-    internal func findSuitableServers(readPreference: ReadPreference? = nil) -> [ServerDescription] {
+    internal func findSuitableServers(
+        readPreference: ReadPreference? = nil,
+        heartbeatFrequencyMS: Int
+    ) throws -> [ServerDescription] {
+        try readPreference?.validateMaxStalenessSeconds(
+            heartbeatFrequencyMS: heartbeatFrequencyMS,
+            topologyType: self.type
+        )
         switch self.type._topologyType {
         case .unknown:
             return []
-        case .single,
-             .loadBalanced:
+        case .single, .loadBalanced:
             return self.servers
-        case .replicaSetNoPrimary,
-             .replicaSetWithPrimary:
+        case .replicaSetNoPrimary, .replicaSetWithPrimary:
             switch readPreference?.mode {
             case .secondary:
-                let secondaries = self.servers.filter { $0.type == .rsSecondary }
-                return self.filterReplicaSetServers(readPreference: readPreference, servers: secondaries)
+                return self.filterReplicaSetServers(
+                    readPreference: readPreference,
+                    heartbeatFrequencyMS: heartbeatFrequencyMS,
+                    includePrimary: false
+                )
             case .nearest:
-                let secondariesAndPrimary = self.servers.filter { $0.type == .rsSecondary || $0.type == .rsPrimary }
-                return self.filterReplicaSetServers(readPreference: readPreference, servers: secondariesAndPrimary)
+                return self.filterReplicaSetServers(
+                    readPreference: readPreference,
+                    heartbeatFrequencyMS: heartbeatFrequencyMS,
+                    includePrimary: true
+                )
             case .secondaryPreferred:
                 // If mode is 'secondaryPreferred', attempt the selection algorithm with mode 'secondary' and the
                 // user's maxStalenessSeconds and tag_sets. If no server matches, select the primary.
-                let secondaries = self.servers.filter { $0.type == .rsSecondary }
-                let primaries = self.servers.filter { $0.type == .rsPrimary }
-                let matches = self.filterReplicaSetServers(readPreference: readPreference, servers: secondaries)
-                return matches.isEmpty ? primaries : matches
+                let secondaryMatches = self.filterReplicaSetServers(
+                    readPreference: readPreference,
+                    heartbeatFrequencyMS: heartbeatFrequencyMS,
+                    includePrimary: false
+                )
+                return secondaryMatches.isEmpty ? self.servers.filter { $0.type == .rsPrimary } : secondaryMatches
             case .primaryPreferred:
                 // If mode is 'primaryPreferred' or a readPreference is not provided, select the primary if it is known,
                 // otherwise attempt the selection algorithm with mode 'secondary' and the user's
@@ -449,8 +475,11 @@ extension TopologyDescription {
                 if !primaries.isEmpty {
                     return primaries
                 }
-                let secondaries = self.servers.filter { $0.type == .rsSecondary }
-                return self.filterReplicaSetServers(readPreference: readPreference, servers: secondaries)
+                return self.filterReplicaSetServers(
+                    readPreference: readPreference,
+                    heartbeatFrequencyMS: heartbeatFrequencyMS,
+                    includePrimary: false
+                )
             default: // or .primary
                 // the default mode is 'primary'.
                 return self.servers.filter { $0.type == .rsPrimary }
@@ -460,13 +489,33 @@ extension TopologyDescription {
         }
     }
 
-    internal func filterReplicaSetServers(
+    /// Filters the replica set servers in this topology first by max staleness and then by tag sets.
+    private func filterReplicaSetServers(
         readPreference: ReadPreference?,
-        servers: [ServerDescription]
+        heartbeatFrequencyMS: Int,
+        includePrimary: Bool
     ) -> [ServerDescription] {
-        // TODO: Filter out servers staler than maxStalenessSeconds
+        // The initial set of servers from which to filter. Only include the secondaries unless includePrimary is true.
+        var servers = self.servers.filter { ($0.type == .rsPrimary && includePrimary) || $0.type == .rsSecondary }
+
+        // Filter by max staleness. If maxStalenessSeconds is not configured as a positive number, all servers are
+        // eligible.
+        if let maxStalenessSeconds = readPreference?.maxStalenessSeconds, maxStalenessSeconds > 0 {
+            let primary = self.servers.first { $0.type == .rsPrimary }
+            let maxLastWriteDate = self.getMaxLastWriteDate()
+            servers.removeAll {
+                guard let staleness = $0.calculateStalenessSeconds(
+                    primary: primary,
+                    maxLastWriteDate: maxLastWriteDate,
+                    heartbeatFrequencyMS: heartbeatFrequencyMS
+                ) else {
+                    return false
+                }
+                return staleness > maxStalenessSeconds
+            }
+        }
 
-        // Filter by tag_sets
+        // Filter by tag sets.
         guard let tagSets = readPreference?.tagSets else {
             return servers
         }
@@ -476,8 +525,86 @@ extension TopologyDescription {
                 return matches
             }
         }
+
+        // If no matches were found during tag set filtering, return an empty list.
         return []
     }
+
+    /// Returns a `Date` representing the latest `lastWriteDate` configured on a secondary in the topology, or `nil`
+    /// if none is found.
+    private func getMaxLastWriteDate() -> Date? {
+        let secondaryLastWriteDates = self.servers.compactMap {
+            $0.type == .rsSecondary ? $0.lastWriteDate : nil
+        }
+        return secondaryLastWriteDates.max()
+    }
+}
+
+extension ServerDescription {
+    /// Calculates the staleness of this server. If the server is not a secondary, the staleness is 0. Otherwise,
+    /// compare against the primary if one is present, or the maximum last write date seen in the topology if present.
+    /// If staleness cannot be calculated due to an absence of values, `nil` is returned.
+    fileprivate func calculateStalenessSeconds(
+        primary: ServerDescription?,
+        maxLastWriteDate: Date?,
+        heartbeatFrequencyMS: Int
+    ) -> Int? {
+        guard self.type == .rsSecondary else {
+            return 0
+        }
+        guard let lastWriteDate = self.lastWriteDate else {
+            return nil
+        }
+        if let primary = primary {
+            guard let primaryLastWriteDate = primary.lastWriteDate else {
+                return nil
+            }
+            let selfInterval = self.lastUpdateTime.timeIntervalSince(lastWriteDate)
+            let primaryInterval = primary.lastUpdateTime.timeIntervalSince(primaryLastWriteDate)
+            // timeIntervalSince returns a TimeInterval in seconds, so heartbeatFrequencyMS needs to be converted from
+            // milliseconds to seconds.
+            let stalenessSeconds = selfInterval - primaryInterval + Double(heartbeatFrequencyMS) / 1000.0
+            return Int(stalenessSeconds.rounded(.up))
+        } else {
+            guard let maxLastWriteDate = maxLastWriteDate else {
+                return nil
+            }
+            let interval = maxLastWriteDate.timeIntervalSince(lastWriteDate)
+            let stalenessSeconds = interval + Double(heartbeatFrequencyMS) / 1000.0
+            return Int(stalenessSeconds.rounded(.up))
+        }
+    }
+}
+
+extension ReadPreference {
+    fileprivate func validateMaxStalenessSeconds(
+        heartbeatFrequencyMS: Int,
+        topologyType: TopologyDescription.TopologyType
+    ) throws {
+        if let maxStalenessSeconds = self.maxStalenessSeconds {
+            if self.mode == .primary && maxStalenessSeconds > 0 {
+                throw MongoError.InvalidArgumentError(
+                    message: "A positive maxStalenessSeconds cannot be specified when the read preference mode is"
+                        + " primary"
+                )
+            }
+            if topologyType == .replicaSetWithPrimary || topologyType == .replicaSetNoPrimary {
+                if maxStalenessSeconds * 1000 < heartbeatFrequencyMS + SDAMConstants.idleWritePeriodMS {
+                    throw MongoError.InvalidArgumentError(
+                        message: "maxStalenessSeconds must be at least the sum of the heartbeatFrequencyMS configured"
+                            + " on the client (\(heartbeatFrequencyMS)) and the idleWritePeriodMS"
+                            + " (\(SDAMConstants.idleWritePeriodMS))"
+                    )
+                }
+                if maxStalenessSeconds < SDAMConstants.smallestMaxStalenessSeconds {
+                    throw MongoError.InvalidArgumentError(
+                        message: "The maxStalenessSeconds configured for a replica set must be at least"
+                            + " \(SDAMConstants.smallestMaxStalenessSeconds)"
+                    )
+                }
+            }
+        }
+    }
 }
 
 extension ServerDescription {

Sources/TestsCommon/CodableExtensions.swift

@@ -1,3 +1,4 @@
+import Foundation
 @testable import MongoSwift
 
 /// Allows a type to specify a set of known keys and check whether any unknown top-level keys are found in a decoder.
@@ -146,9 +147,15 @@ extension ReadPreference: StrictDecodable {
     public init(from decoder: Decoder) throws {
         if let container = try? decoder.container(keyedBy: CodingKeys.self) {
             try Self.checkKeys(using: decoder)
-            let mode = try container.decode(Mode.self, forKey: .mode)
-            let tagSets = try container.decodeIfPresent([BSONDocument].self, forKey: .tagSets)
-            try self.init(mode, tagSets: tagSets)
+            // Some tests specify a read preference with no fields to indicate the default read preference (i.e.
+            // primary). Because this is not representable in the Swift driver due to the mode field not being
+            // optional, this sets the mode to be primary explicitly if one is not present.
+            let mode = try container.decodeIfPresent(Mode.self, forKey: .mode) ?? Mode.primary
+            self.init(mode)
+            // The init method that takes in these fields also performs validation, so these fields are set manually to
+            // allow decoding to succeed and ensure that validation occurs during server selection.
+            self.tagSets = try container.decodeIfPresent([BSONDocument].self, forKey: .tagSets)
+            self.maxStalenessSeconds = try container.decodeIfPresent(Int.self, forKey: .maxStalenessSeconds)
         } else { // sometimes the spec tests only specify the mode as a string
             let container = try decoder.singleValueContainer()
             let mode = try container.decode(ReadPreference.Mode.self)
@@ -157,7 +164,7 @@ extension ReadPreference: StrictDecodable {
     }
 
     internal enum CodingKeys: String, CodingKey, CaseIterable {
-        case mode, tagSets = "tag_sets"
+        case mode, tagSets = "tag_sets", maxStalenessSeconds
     }
 }
 
@@ -205,12 +212,40 @@ 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-1456: decode and set averageRoundTripTimeMS
+        let maxWireVersion = try values.decodeIfPresent(Int.self, forKey: .maxWireVersion)
+
+        var lastUpdateTime: Date?
+        if let lastUpdateTimeMS = try values.decodeIfPresent(Int64.self, forKey: .lastUpdateTime) {
+            lastUpdateTime = Date(msSinceEpoch: lastUpdateTimeMS)
+        }
 
-        self.init(address: address, type: type, tags: tags)
+        // lastWriteDate is specified in a document in the form described in the error message below
+        var lastWriteDate: Date?
+        if let lastWrite = try values.decodeIfPresent(BSONDocument.self, forKey: .lastWrite) {
+            guard let lastWriteDateMS = lastWrite["lastWriteDate"]?.int64Value else {
+                throw DecodingError.dataCorruptedError(
+                    forKey: .lastWrite,
+                    in: values,
+                    debugDescription: "lastWrite should be specified in the form"
+                        + " lastWrite: { lastWriteDate: { \"$numberLong\": value } }"
+                )
+            }
+            lastWriteDate = Date(msSinceEpoch: lastWriteDateMS)
+        }
+
+        // TODO: SWIFT-1461: decode and set averageRoundTripTimeMS
+
+        self.init(
+            address: address,
+            type: type,
+            tags: tags,
+            lastWriteDate: lastWriteDate,
+            maxWireVersion: maxWireVersion,
+            lastUpdateTime: lastUpdateTime
+        )
     }
 
     internal enum CodingKeys: String, CodingKey, CaseIterable {
-        case address, type, tags, averageRoundTripTimeMS = "avg_rtt_ms"
+        case address, type, tags, averageRoundTripTimeMS = "avg_rtt_ms", lastWrite, maxWireVersion, lastUpdateTime
     }
 }