SWIFT-1329 Use a client per mongos in unified test runner (#667)

kmahar · web-flow · commit 0d5a7fbfda34 · 2021-08-30T14:59:49.000-04:00
diff --git a/Sources/MongoSwift/Operations/RunCommandOperation.swift b/Sources/MongoSwift/Operations/RunCommandOperation.swift
@@ -35,51 +35,19 @@ internal struct RunCommandOperation: Operation {
     private let database: MongoDatabase
     private let command: BSONDocument
     private let options: RunCommandOptions?
-    private let serverAddress: ServerAddress?
 
     internal init(
         database: MongoDatabase,
         command: BSONDocument,
-        options: RunCommandOptions?,
-        serverAddress: ServerAddress? = nil
+        options: RunCommandOptions?
     ) {
         self.database = database
         self.command = command
         self.options = options
-        self.serverAddress = serverAddress
     }
 
     internal func execute(using connection: Connection, session: ClientSession?) throws -> BSONDocument {
-        let serverId: UInt32? = try self.serverAddress.map { address in
-            let id: UInt32? = connection.withMongocConnection { connection in
-                var numServers = 0
-                let sds = mongoc_client_get_server_descriptions(connection, &numServers)
-                defer { mongoc_server_descriptions_destroy_all(sds, numServers) }
-
-                guard numServers > 0 else {
-                    return nil
-                }
-
-                let buffer = UnsafeBufferPointer(start: sds, count: numServers)
-                // the buffer is documented as always containing non-nil pointers (if non-empty).
-                // swiftlint:disable:next force_unwrapping
-                let servers = Array(buffer).map { ServerDescription($0!) }
-                return servers.first { $0.address == address }?.serverId
-            }
-
-            guard let out = id else {
-                throw MongoError.ServerSelectionError(message: "No known host with address \(address)")
-            }
-
-            return out
-        }
-
-        var opts = try encodeOptions(options: self.options, session: session)
-        if let id = serverId {
-            opts = opts ?? [:]
-            opts?["serverId"] = .int64(Int64(id))
-        }
-
+        let opts = try encodeOptions(options: self.options, session: session)
         return try self.database.withMongocDatabase(from: connection) { dbPtr in
             try ReadPreference.withOptionalMongocReadPreference(from: self.options?.readPreference) { rpPtr in
                 try runMongocCommandWithReply(
diff --git a/Sources/TestsCommon/SpecTestUtils.swift b/Sources/TestsCommon/SpecTestUtils.swift
@@ -20,19 +20,6 @@ extension MongoSwiftTestCase {
     }
 }
 
-extension MongoDatabase {
-    @discardableResult
-    public func runCommand(
-        _ command: BSONDocument,
-        on server: ServerAddress,
-        options: RunCommandOptions? = nil,
-        session: ClientSession? = nil
-    ) -> EventLoopFuture<BSONDocument> {
-        let operation = RunCommandOperation(database: self, command: command, options: options, serverAddress: server)
-        return self._client.operationExecutor.execute(operation, client: self._client, on: nil, session: session)
-    }
-}
-
 /// Given a spec folder name (e.g. "crud") and optionally a subdirectory name for a folder (e.g. "read") retrieves an
 /// array of [(filename, file decoded to type T)].
 public func retrieveSpecTestFiles<T: Decodable>(
diff --git a/Tests/MongoSwiftSyncTests/SpecTestRunner/FailPoint.swift b/Tests/MongoSwiftSyncTests/SpecTestRunner/FailPoint.swift
@@ -14,22 +14,19 @@ extension FailPointConfigured {
     /// Sets the active fail point to the provided fail point and enables it.
     internal mutating func activateFailPoint(
         _ failPoint: FailPoint,
-        using client: MongoClient,
-        on serverAddress: ServerAddress? = nil
+        using client: MongoClient
     ) throws {
         self.activeFailPoint = failPoint
-        try self.activeFailPoint?.enable(using: client, on: serverAddress)
-        self.targetedHost = serverAddress
+        try self.activeFailPoint?.enable(using: client)
     }
 
     /// If a fail point is active, it is disabled and cleared.
     internal mutating func disableActiveFailPoint(using client: MongoClient) {
         guard let failPoint = self.activeFailPoint else {
             return
         }
-        failPoint.disable(using: client, on: self.targetedHost)
+        failPoint.disable(using: client)
         self.activeFailPoint = nil
-        self.targetedHost = nil
     }
 }
 
@@ -39,17 +36,14 @@ class FailPointGuard {
     let failPoint: FailPoint
     /// Client to use when disabling the failpoint.
     let client: MongoClient
-    /// Optional server address to disable the failpoint on.
-    let serverAddress: ServerAddress?
 
-    init(failPoint: FailPoint, client: MongoClient, serverAddress: ServerAddress?) {
+    init(failPoint: FailPoint, client: MongoClient) {
         self.failPoint = failPoint
         self.client = client
-        self.serverAddress = serverAddress
     }
 
     deinit {
-        self.failPoint.disable(using: self.client, on: self.serverAddress)
+        self.failPoint.disable(using: self.client)
     }
 }
 
@@ -90,33 +84,27 @@ internal struct FailPoint: Decodable {
 
     internal func enable(
         using client: MongoClient,
-        on serverAddress: ServerAddress? = nil,
         options: RunCommandOptions? = nil
     ) throws {
-        if let address = serverAddress {
-            try client.db("admin").runCommand(self.failPoint, on: address, options: options)
-        } else {
-            try client.db("admin").runCommand(self.failPoint, options: options)
-        }
+        try client.db("admin").runCommand(self.failPoint, options: options)
     }
 
     /// Enables the failpoint, and returns a `FailPointGuard` which will automatically disable the failpoint
     /// upon deinitialization.
     internal func enableWithGuard(
         using client: MongoClient,
-        on serverAddress: ServerAddress? = nil,
         options: RunCommandOptions? = nil
     ) throws -> FailPointGuard {
-        try self.enable(using: client, on: serverAddress, options: options)
-        return FailPointGuard(failPoint: self, client: client, serverAddress: serverAddress)
+        try self.enable(using: client, options: options)
+        return FailPointGuard(failPoint: self, client: client)
     }
 
     internal func enable() throws {
         let client = try MongoClient.makeTestClient()
         try self.enable(using: client)
     }
 
-    internal func disable(using client: MongoClient? = nil, on address: ServerAddress? = nil) {
+    internal func disable(using client: MongoClient? = nil) {
         do {
             let clientToUse: MongoClient
             if let client = client {
@@ -126,15 +114,7 @@ internal struct FailPoint: Decodable {
             }
 
             let command: BSONDocument = ["configureFailPoint": .string(self.name), "mode": "off"]
-
-            if let addr = address {
-                try clientToUse.db("admin").runCommand(command, on: addr)
-            } else {
-                try clientToUse.db("admin").runCommand(command)
-            }
-        } catch _ as MongoError.ServerSelectionError {
-            // this often means the server that the failpoint was set against was marked as unknown
-            // due to the failpoint firing, so we just ignore
+            try clientToUse.db("admin").runCommand(command)
         } catch {
             print("Failed to disable failpoint: \(error)")
         }
diff --git a/Tests/MongoSwiftSyncTests/SyncTestUtils.swift b/Tests/MongoSwiftSyncTests/SyncTestUtils.swift
@@ -154,18 +154,6 @@ internal func captureCommandEvents(
     return monitor.events(withEventTypes: eventTypes, withNames: commandNames)
 }
 
-extension MongoDatabase {
-    @discardableResult
-    public func runCommand(
-        _ command: BSONDocument,
-        on server: ServerAddress,
-        options: RunCommandOptions? = nil,
-        session: MongoSwiftSync.ClientSession? = nil
-    ) throws -> BSONDocument {
-        try self.asyncDB.runCommand(command, on: server, options: options, session: session?.asyncSession).wait()
-    }
-}
-
 extension MongoSwiftSync.MongoCollection {
     public var _client: MongoSwiftSync.MongoClient {
         self.client
diff --git a/Tests/MongoSwiftSyncTests/UnifiedTestRunner/Context.swift b/Tests/MongoSwiftSyncTests/UnifiedTestRunner/Context.swift
@@ -11,9 +11,9 @@ class Context {
     /// Fail points that have been set during test execution and should be disabled on completion.
     var enabledFailPoints: [FailPointGuard] = []
 
-    let internalClient: MongoClient
+    let internalClient: UnifiedTestRunner.InternalClient
 
-    init(path: [String], entities: EntityMap, internalClient: MongoClient) {
+    init(path: [String], entities: EntityMap, internalClient: UnifiedTestRunner.InternalClient) {
         self.path = path
         self.entities = entities
         self.internalClient = internalClient
diff --git a/Tests/MongoSwiftSyncTests/UnifiedTestRunner/SpecialTestOperations.swift b/Tests/MongoSwiftSyncTests/UnifiedTestRunner/SpecialTestOperations.swift
@@ -3,6 +3,7 @@
 import MongoSwiftSync
 @testable import class MongoSwiftSync.ClientSession
 import Nimble
+import TestsCommon
 
 struct UnifiedFailPoint: UnifiedOperationProtocol {
     /// The failpoint to set.
@@ -36,7 +37,7 @@ struct UnifiedAssertCollectionExists: UnifiedOperationProtocol {
     }
 
     func execute(on _: UnifiedOperation.Object, context: Context) throws -> UnifiedOperationResult {
-        let db = context.internalClient.db(self.databaseName)
+        let db = context.internalClient.anyClient.db(self.databaseName)
         expect(try db.listCollectionNames()).to(
             contain(self.collectionName),
             description: "Expected db \(self.databaseName) to contain collection \(self.collectionName)." +
@@ -58,7 +59,7 @@ struct UnifiedAssertCollectionNotExists: UnifiedOperationProtocol {
     }
 
     func execute(on _: UnifiedOperation.Object, context: Context) throws -> UnifiedOperationResult {
-        let db = context.internalClient.db(self.databaseName)
+        let db = context.internalClient.anyClient.db(self.databaseName)
         expect(try db.listCollectionNames()).toNot(
             contain(self.collectionName),
             description: "Expected db \(self.databaseName) to not contain collection \(self.collectionName)." +
@@ -83,7 +84,7 @@ struct UnifiedAssertIndexExists: UnifiedOperationProtocol {
     }
 
     func execute(on _: UnifiedOperation.Object, context: Context) throws -> UnifiedOperationResult {
-        let collection = context.internalClient.db(self.databaseName).collection(self.collectionName)
+        let collection = context.internalClient.anyClient.db(self.databaseName).collection(self.collectionName)
         expect(try collection.listIndexNames()).to(
             contain(self.indexName),
             description: "Expected collection \(collection.namespace) to have index \(self.indexName)."
@@ -108,7 +109,7 @@ struct UnifiedAssertIndexNotExists: UnifiedOperationProtocol {
     }
 
     func execute(on _: UnifiedOperation.Object, context: Context) throws -> UnifiedOperationResult {
-        let collection = context.internalClient.db(self.databaseName).collection(self.collectionName)
+        let collection = context.internalClient.anyClient.db(self.databaseName).collection(self.collectionName)
         expect(try collection.listIndexNames()).toNot(
             contain(self.indexName),
             description: "Expected collection \(collection.namespace) to not have index \(self.indexName)."
@@ -276,11 +277,11 @@ struct UnifiedTargetedFailPoint: UnifiedOperationProtocol {
                 beNil(),
                 description: "Session \(self.session) unexpectedly not pinned to a mongos. Path: \(context.path)"
             )
-        let pinnedMongos = session.pinnedServerAddress!
-        // The test runner SHOULD use the client entity associated with the session to execute the configureFailPoint
-        // command.
-        let client = session.client
-        let fpGuard = try self.failPoint.enableWithGuard(using: client, on: pinnedMongos)
+        let mongosClients = try context.internalClient.asMongosClients()
+        guard let clientForPinnedMongos = mongosClients[session.pinnedServerAddress!] else {
+            throw TestError(message: "Unexpectedly missing client for mongos \(session.pinnedServerAddress!)")
+        }
+        let fpGuard = try self.failPoint.enableWithGuard(using: clientForPinnedMongos)
         // add to context's list of enabled failpoints to ensure we disable it later.
         context.enabledFailPoints.append(fpGuard)
         return .none
diff --git a/Tests/MongoSwiftSyncTests/UnifiedTestRunner/UnifiedTestRunner.swift b/Tests/MongoSwiftSyncTests/UnifiedTestRunner/UnifiedTestRunner.swift
@@ -3,7 +3,36 @@ import Nimble
 import TestsCommon
 
 struct UnifiedTestRunner {
-    let internalClient: MongoClient
+    enum InternalClient {
+        /// For all topologies besides sharded, we use a single client.
+        case single(MongoClient)
+        /// For sharded topologies, we often need to target commands to particular mongoses, so we use a separate
+        /// client for each host we're connecting to.
+        case mongosClients([ServerAddress: MongoClient])
+
+        /// Returns an internal client; for usage in situations where any client connected to the topology will do -
+        /// that is, if the topology is sharded, we do not care about targeting a particular mongos.
+        var anyClient: MongoClient {
+            switch self {
+            case let .single(c):
+                return c
+            case let .mongosClients(clientMap):
+                return clientMap.first!.1
+            }
+        }
+
+        /// If the internal client is a map of per-mongos clients, returns that map; otherwise throws an erorr.
+        func asMongosClients() throws -> [ServerAddress: MongoClient] {
+            guard case let .mongosClients(clientMap) = self else {
+                throw TestError(
+                    message: "Runner unexpectedly did not create per-mongos clients"
+                )
+            }
+            return clientMap
+        }
+    }
+
+    let internalClient: InternalClient
     let serverVersion: ServerVersion
     let topologyType: TestTopologyConfiguration
     let serverParameters: BSONDocument
@@ -12,11 +41,22 @@ struct UnifiedTestRunner {
     static let maxSchemaVersion = SchemaVersion(rawValue: "1.5.0")!
 
     init() throws {
-        let connStr = MongoSwiftTestCase.getConnectionString(singleMongos: false).toString()
-        self.internalClient = try MongoClient.makeTestClient(connStr)
-        self.serverVersion = try self.internalClient.serverVersion()
-        self.topologyType = try self.internalClient.topologyType()
-        self.serverParameters = try self.internalClient.serverParameters()
+        switch MongoSwiftTestCase.topologyType {
+        case .sharded:
+            var mongosClients = [ServerAddress: MongoClient]()
+            for host in MongoSwiftTestCase.getHosts() {
+                let connString = MongoSwiftTestCase.getConnectionString(forHost: host)
+                let client = try MongoClient.makeTestClient(connString.toString())
+                mongosClients[host] = client
+            }
+            self.internalClient = .mongosClients(mongosClients)
+        default:
+            let client = try MongoClient.makeTestClient()
+            self.internalClient = .single(client)
+        }
+        self.serverVersion = try self.internalClient.anyClient.serverVersion()
+        self.topologyType = try self.internalClient.anyClient.topologyType()
+        self.serverParameters = try self.internalClient.anyClient.serverParameters()
     }
 
     func terminateOpenTransactions() throws {
@@ -32,12 +72,12 @@ struct UnifiedTestRunner {
             // SERVER-38335.
             do {
                 let opts = RunCommandOptions(readPreference: .primary)
-                _ = try self.internalClient.db("admin").runCommand(["killAllSessions": []], options: opts)
+                _ = try self.internalClient.anyClient.db("admin").runCommand(["killAllSessions": []], options: opts)
             } catch let commandError as MongoError.CommandError where commandError.code == 11601 {}
         case .sharded, .shardedReplicaSet:
-            for address in MongoSwiftTestCase.getHosts() {
+            for (_, client) in try self.internalClient.asMongosClients() {
                 do {
-                    _ = try self.internalClient.db("admin").runCommand(["killAllSessions": []], on: address)
+                    _ = try client.db("admin").runCommand(["killAllSessions": []])
                 } catch let commandError as MongoError.CommandError where commandError.code == 11601 {
                     continue
                 }
@@ -113,7 +153,7 @@ struct UnifiedTestRunner {
                 // The test runner MUST use the internal MongoClient for these operations.
                 if let initialData = file.initialData {
                     for collData in initialData {
-                        let db = self.internalClient.db(collData.databaseName)
+                        let db = self.internalClient.anyClient.db(collData.databaseName)
                         let collOpts = MongoCollectionOptions(writeConcern: .majority)
                         let coll = db.collection(collData.collectionName, options: collOpts)
                         try coll.drop()
@@ -141,11 +181,10 @@ struct UnifiedTestRunner {
                 // the implementation, test runners MAY execute distinct before every test.
                 if self.topologyType.isSharded && !MongoSwiftTestCase.serverless {
                     let collEntities = context.entities.values.compactMap { try? $0.asCollection() }
-                    for address in MongoSwiftTestCase.getHosts() {
+                    for (_, client) in try self.internalClient.asMongosClients() {
                         for entity in collEntities {
-                            _ = try self.internalClient.db(entity.namespace.db).runCommand(
-                                ["distinct": .string(entity.name), "key": "_id"],
-                                on: address
+                            _ = try client.db(entity.namespace.db).runCommand(
+                                ["distinct": .string(entity.name), "key": "_id"]
                             )
                         }
                     }
@@ -188,6 +227,7 @@ struct UnifiedTestRunner {
                     if let expectedOutcome = test.outcome {
                         for collectionData in expectedOutcome {
                             let collection = self.internalClient
+                                .anyClient
                                 .db(collectionData.databaseName)
                                 .collection(collectionData.collectionName)
                             let opts = FindOptions(
