SWIFT-1045 Ensure MongoClient initializer does not perform blocking DNS request (#596)

Author: patrickfreed

Sources/MongoSwift/APM.swift

@@ -742,7 +742,6 @@ extension ConnectionPool {
         guard let callbacks = mongoc_apm_callbacks_new() else {
             fatalError("failed to initialize new mongoc_apm_callbacks_t")
         }
-        defer { mongoc_apm_callbacks_destroy(callbacks) }
 
         mongoc_apm_set_command_started_cb(callbacks, commandStarted)
         mongoc_apm_set_command_succeeded_cb(callbacks, commandSucceeded)

Sources/MongoSwift/ConnectionPool.swift

@@ -1,5 +1,6 @@
 import CLibMongoC
 import Foundation
+import NIO
 import NIOConcurrencyHelpers
 
 /// A connection to the database.
@@ -39,6 +40,8 @@ extension NSCondition {
 internal class ConnectionPool {
     /// Represents the state of a `ConnectionPool`.
     private enum State {
+        /// Indicates the `ConnectionPool` is still starting up in the background.
+        case opening(future: EventLoopFuture<OpaquePointer>)
         /// Indicates that the `ConnectionPool` is open and using the associated pointer to a `mongoc_client_pool_t`.
         case open(pool: OpaquePointer)
         /// Indicates that the `ConnectionPool` is in the process of closing. Connections can be checked back in, but
@@ -75,23 +78,25 @@ internal class ConnectionPool {
     internal static let PoolClosedError = MongoError.LogicError(message: "ConnectionPool was already closed")
 
     /// Initializes the pool using the provided `ConnectionString`.
-    internal init(from connString: ConnectionString) throws {
-        let pool: OpaquePointer = try connString.withMongocURI { uriPtr in
-            guard let pool = mongoc_client_pool_new(uriPtr) else {
-                throw MongoError.InternalError(message: "Failed to initialize libmongoc client pool")
-            }
-            return pool
-        }
+    internal init(from connString: ConnectionString, executor: OperationExecutor) throws {
+        let poolFut = executor.execute(on: nil) { () -> OpaquePointer in
+            try connString.withMongocURI { uriPtr in
+                guard let pool = mongoc_client_pool_new(uriPtr) else {
+                    throw MongoError.InternalError(message: "Failed to initialize libmongoc client pool")
+                }
 
-        self.state = .open(pool: pool)
+                guard mongoc_client_pool_set_error_api(pool, MONGOC_ERROR_API_VERSION_2) else {
+                    fatalError("Could not configure error handling on client pool")
+                }
 
-        guard mongoc_client_pool_set_error_api(pool, MONGOC_ERROR_API_VERSION_2) else {
-            fatalError("Could not configure error handling on client pool")
+                // We always set min_heartbeat_frequency because the hard-coded default in the vendored mongoc
+                // was lowered to 50. Setting it here brings it to whatever was specified, or 500 if it wasn't.
+                mongoc_client_pool_set_min_heartbeat_frequency_msec(pool, UInt64(connString.minHeartbeatFrequencyMS))
+                return pool
+            }
         }
 
-        // We always set min_heartbeat_frequency because the hard-coded default in the vendored mongoc
-        // was lowered to 50. Setting it here brings it to whatever was specified, or 500 if it wasn't.
-        mongoc_client_pool_set_min_heartbeat_frequency_msec(pool, UInt64(connString.minHeartbeatFrequencyMS))
+        self.state = .opening(future: poolFut)
     }
 
     deinit {
@@ -101,13 +106,36 @@ internal class ConnectionPool {
         }
     }
 
+    /// Execute the provided closure with the pointer to the created `mongoc_client_pool_t`, potentially waiting
+    /// for the pool to finish starting up first. The closure will be executed while the stateLock is held, so
+    /// closure MUST NOT attempt to acquire the lock.
+    ///
+    /// Throws an error if the pool is closed or closing.
+    private func withResolvedPool<T>(body: (OpaquePointer) throws -> T) throws -> T {
+        try self.stateLock.withLock {
+            switch self.state {
+            case let .open(pool):
+                return try body(pool)
+            case let .opening(future):
+                let pool = try future.wait()
+                self.state = .open(pool: pool)
+                return try body(pool)
+            case .closed, .closing:
+                throw Self.PoolClosedError
+            }
+        }
+    }
+
     /// Closes the pool, cleaning up underlying resources. **This method blocks until all connections are returned to
     /// the pool.**
     internal func close() throws {
         try self.stateLock.withLock {
             switch self.state {
             case let .open(pool):
                 self.state = .closing(pool: pool)
+            case let .opening(future):
+                let pool = try future.wait()
+                self.state = .closing(pool: pool)
             case .closing, .closed:
                 throw Self.PoolClosedError
             }
@@ -118,7 +146,7 @@ internal class ConnectionPool {
             }
 
             switch self.state {
-            case .open, .closed:
+            case .open, .closed, .opening:
                 throw MongoError.InternalError(
                     message: "ConnectionPool in unexpected state \(self.state) during close()"
                 )
@@ -133,32 +161,22 @@ internal class ConnectionPool {
     /// This method will block until a connection is available. Throws an error if the pool is in the process of
     /// closing or has finished closing.
     internal func checkOut() throws -> Connection {
-        try self.stateLock.withLock {
-            switch self.state {
-            case let .open(pool):
-                self.connCount += 1
-                return Connection(clientHandle: mongoc_client_pool_pop(pool), pool: self)
-            case .closing, .closed:
-                throw Self.PoolClosedError
-            }
+        try self.withResolvedPool { pool in
+            self.connCount += 1
+            return Connection(clientHandle: mongoc_client_pool_pop(pool), pool: self)
         }
     }
 
     /// Checks out a connection from the pool, or returns `nil` if none are currently available. Throws an error if the
     /// pool is not open. This method may block waiting on the state lock as well as libmongoc's locks and thus must be
     // run within the thread pool.
     internal func tryCheckOut() throws -> Connection? {
-        try self.stateLock.withLock {
-            switch self.state {
-            case let .open(pool):
-                guard let handle = mongoc_client_pool_try_pop(pool) else {
-                    return nil
-                }
-                self.connCount += 1
-                return Connection(clientHandle: handle, pool: self)
-            case .closing, .closed:
-                throw Self.PoolClosedError
+        try self.withResolvedPool { pool in
+            guard let handle = mongoc_client_pool_try_pop(pool) else {
+                return nil
             }
+            self.connCount += 1
+            return Connection(clientHandle: handle, pool: self)
         }
     }
 
@@ -178,6 +196,8 @@ internal class ConnectionPool {
                 self.stateLock.signal()
             case .closed:
                 throw Self.PoolClosedError
+            case .opening:
+                fatalError("ConnectionPool in unexpected state \(self.state) while checking in a connection")
             }
         }
     }
@@ -230,15 +250,31 @@ internal class ConnectionPool {
     /// may only be called before any connections are checked out of the pool.** Ideally this code would just live in
     /// `ConnectionPool.init`. However, the client we accept here has to be fully initialized before we can pass it
     /// as the context. In order for it to be fully initialized its pool must exist already.
+    ///
+    /// This method takes ownership of the provided `mongoc_apm_callbacks_t`, so the pointer MUST NOT be used or freed
+    /// after it is passed to this method.
     internal func setAPMCallbacks(callbacks: OpaquePointer, client: MongoClient) {
         // lock isn't needed as this is called before pool is in use.
-        switch self.state {
-        case let .open(pool):
-            mongoc_client_pool_set_apm_callbacks(pool, callbacks, Unmanaged.passUnretained(client).toOpaque())
-        case .closing, .closed:
+        guard case let .opening(future) = self.state else {
             // this method is called via `initializeMonitoring()`, which is called from `MongoClient.init`.
             // unless we have a bug it's impossible that the pool is already closed.
             fatalError("ConnectionPool in unexpected state \(self.state) while setting APM callbacks")
         }
+
+        // to ensure the callbacks get set before any waiting on this future becomes unblocked, use
+        // `map` to create a new future instead of .whenSuccess
+        let newFut = future.map { pool -> OpaquePointer in
+            mongoc_client_pool_set_apm_callbacks(pool, callbacks, Unmanaged.passUnretained(client).toOpaque())
+            mongoc_apm_callbacks_destroy(callbacks)
+            return pool
+        }
+
+        // this shouldn't be reached but just in case we still free the memory
+        newFut.whenFailure { _ in
+            mongoc_apm_callbacks_destroy(callbacks)
+        }
+
+        // update the stored future to refer to the one that sets the callbacks
+        self.state = .opening(future: newFut)
     }
 }

Sources/MongoSwift/MongoClient.swift

@@ -285,11 +285,11 @@ public class MongoClient {
         initializeMongoc()
 
         let connString = try ConnectionString(connectionString, options: options)
-        self.connectionPool = try ConnectionPool(from: connString)
         self.operationExecutor = OperationExecutor(
             eventLoopGroup: eventLoopGroup,
             threadPoolSize: options?.threadPoolSize ?? MongoClient.defaultThreadPoolSize
         )
+        self.connectionPool = try ConnectionPool(from: connString, executor: self.operationExecutor)
 
         let rc = connString.readConcern
         if !rc.isDefault {