close API updates

Author: kmahar

Sources/MongoSwift/APM.swift

@@ -701,7 +701,8 @@ private func publishEvent<T: MongoSwiftEvent>(type: T.Type, eventPtr: OpaquePoin
 
 /// An extension of `ConnectionPool` to add monitoring capability for commands and server discovery and monitoring.
 extension ConnectionPool {
-    /// Internal function to install monitoring callbacks for this pool.
+    /// Internal function to install monitoring callbacks for this pool. **This method may only be called before any
+    /// connections are checked out from the pool.**
     internal func initializeMonitoring(client: MongoClient) {
         guard let callbacks = mongoc_apm_callbacks_new() else {
             fatalError("failed to initialize new mongoc_apm_callbacks_t")
@@ -722,13 +723,6 @@ extension ConnectionPool {
         mongoc_apm_set_server_heartbeat_succeeded_cb(callbacks, serverHeartbeatSucceeded)
         mongoc_apm_set_server_heartbeat_failed_cb(callbacks, serverHeartbeatFailed)
 
-        // we can pass the MongoClient as unretained because the callbacks are stored on clientHandle, so if the
-        // callback is being executed, this pool and therefore its parent `MongoClient` must still be valid.
-        switch self.state {
-        case let .open(pool):
-            mongoc_client_pool_set_apm_callbacks(pool, callbacks, Unmanaged.passUnretained(client).toOpaque())
-        case .closed:
-            fatalError("ConnectionPool was already closed")
-        }
+        self.setAPMCallbacks(callbacks: callbacks, client: client)
     }
 }

Sources/MongoSwift/ConnectionPool.swift

@@ -1,4 +1,5 @@
 import CLibMongoC
+import NIOConcurrencyHelpers
 
 /// A connection to the database.
 internal class Connection {
@@ -13,11 +14,10 @@ internal class Connection {
     }
 
     deinit {
-        switch self.pool.state {
-        case let .open(pool):
-            mongoc_client_pool_push(pool, self.clientHandle)
-        case .closed:
-            assertionFailure("ConnectionPool was already closed")
+        do {
+            try self.pool.checkIn(self)
+        } catch {
+            assertionFailure("Failed to check connection back in: \(error)")
         }
     }
 }
@@ -28,12 +28,38 @@ internal class ConnectionPool {
     internal enum State {
         /// 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
+        /// no new connections can be checked out.
+        case closing(pool: OpaquePointer)
         /// Indicates that the `ConnectionPool` has been closed and contains no connections.
         case closed
     }
 
     /// The state of this `ConnectionPool`.
-    internal private(set) var state: State
+    private var state: State
+    /// The number of connections currently checked out of the pool.
+    private var connCount = 0
+    /// Lock over `state` and `connCount`.
+    private let stateLock = Lock()
+
+    /// Internal helper for testing purposes that returns whether the pool is in the `closing` state.
+    internal var isClosing: Bool {
+        self.stateLock.withLock {
+            guard case .closing = self.state else {
+                return false
+            }
+            return true
+        }
+    }
+
+    /// Internal helper for testing purposes that returns the number of connections currently checked out from the pool.
+    internal var checkedOutConnections: Int {
+        self.stateLock.withLock {
+            self.connCount
+        }
+    }
+
+    private static let PoolClosedError = LogicError(message: "ConnectionPool was already closed")
 
     /// Initializes the pool using the provided `ConnectionString` and options.
     internal init(from connString: ConnectionString, options: ClientOptions?) throws {
@@ -56,7 +82,7 @@ internal class ConnectionPool {
 
         self.state = .open(pool: pool)
         if let options = options {
-            try self.setTLSOptions(options)
+            self.setTLSOptions(options)
         }
     }
 
@@ -67,38 +93,76 @@ internal class ConnectionPool {
         }
     }
 
-    /// Closes the pool, cleaning up underlying resources. This method blocks as it sends `endSessions` to the server.
-    internal func shutdown() {
-        switch self.state {
-        case let .open(pool):
-            mongoc_client_pool_destroy(pool)
-        case .closed:
-            return
+    /// 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 .closing, .closed:
+                throw Self.PoolClosedError
+            }
+        }
+
+        // continually loop and wait to get all connections back before destroying the pool. release the lock on each
+        // iteration to allow other methods to acquire the lock.
+        var done = false
+        while !done {
+            try self.stateLock.withLock {
+                if self.connCount == 0 {
+                    switch self.state {
+                    case .open, .closed:
+                        throw InternalError(message: "ConnectionPool in unexpected state during close()")
+                    case let .closing(pool):
+                        mongoc_client_pool_destroy(pool)
+                        self.state = .closed
+                    }
+                    done = true
+                }
+            }
         }
-        self.state = .closed
     }
 
     /// Checks out a connection. This connection will return itself to the pool when its reference count drops to 0.
     /// This method will block until a connection is available.
     internal func checkOut() throws -> Connection {
-        switch self.state {
-        case let .open(pool):
-            return Connection(clientHandle: mongoc_client_pool_pop(pool), pool: self)
-        case .closed:
-            throw InternalError(message: "ConnectionPool was already closed")
+        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
+            }
         }
     }
 
     /// Checks out a connection from the pool, or returns `nil` if none are currently available.
     internal func tryCheckOut() throws -> Connection? {
-        switch self.state {
-        case let .open(pool):
-            guard let handle = mongoc_client_pool_try_pop(pool) else {
-                return nil
+        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
+            }
+        }
+    }
+
+    fileprivate func checkIn(_ connection: Connection) throws {
+        try self.stateLock.withLock {
+            switch self.state {
+            case let .open(pool), let .closing(pool):
+                mongoc_client_pool_push(pool, connection.clientHandle)
+                self.connCount -= 1
+            case .closed:
+                throw Self.PoolClosedError
             }
-            return Connection(clientHandle: handle, pool: self)
-        case .closed:
-            throw InternalError(message: "ConnectionPool was already closed")
         }
     }
 
@@ -109,9 +173,9 @@ internal class ConnectionPool {
         return try body(connection)
     }
 
-    // Sets TLS/SSL options that the user passes in through the client level. This must be called from
-    // the ConnectionPool init before the pool is used.
-    private func setTLSOptions(_ options: ClientOptions) throws {
+    // Sets TLS/SSL options that the user passes in through the client level. **This must only be called from
+    // the ConnectionPool initializer**.
+    private func setTLSOptions(_ options: ClientOptions) {
         // return early so we don't set an empty options struct on the libmongoc pool. doing so will make libmongoc
         // attempt to use TLS for connections.
         guard options.tls == true ||
@@ -147,11 +211,15 @@ internal class ConnectionPool {
         if let invalidHosts = options.tlsAllowInvalidHostnames {
             opts.allow_invalid_hostname = invalidHosts
         }
-        switch self.state {
-        case let .open(pool):
-            mongoc_client_pool_set_ssl_opts(pool, &opts)
-        case .closed:
-            throw InternalError(message: "ConnectionPool was already closed")
+
+        self.stateLock.withLock {
+            switch self.state {
+            case let .open(pool):
+                mongoc_client_pool_set_ssl_opts(pool, &opts)
+            case .closing, .closed:
+                // if we get here, we must have called this method outside of `ConnectionPool.init`.
+                fatalError("ConnectionPool unexpectedly in .closing or .closed state")
+            }
         }
     }
 
@@ -187,6 +255,21 @@ internal class ConnectionPool {
             return ConnectionString(copying: uri)
         }
     }
+
+    /// Sets the provided APM callbacks on this pool, using the provided client as the "context" value. **This method
+    /// may only be called before any connections are checked out of the pool.**
+    internal func setAPMCallbacks(callbacks: OpaquePointer, client: MongoClient) {
+        self.stateLock.withLock {
+            switch self.state {
+            case let .open(pool):
+                mongoc_client_pool_set_apm_callbacks(pool, callbacks, Unmanaged.passUnretained(client).toOpaque())
+            case .closing, .closed:
+                // 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 unexpectedly in .closed state")
+            }
+        }
+    }
 }
 
 extension String {

Sources/MongoSwift/MongoClient.swift

@@ -157,8 +157,10 @@ public struct DatabaseOptions: CodingStrategyProvider {
 // sourcery: skipSyncExport
 /// A MongoDB Client providing an asynchronous, SwiftNIO-based API.
 public class MongoClient {
+    /// The pool of connections backing this client.
     internal let connectionPool: ConnectionPool
 
+    /// Executor responsible for executing operations on behalf of this client and its child objects.
     internal let operationExecutor: OperationExecutor
 
     /// Default size for a client's NIOThreadPool.
@@ -167,8 +169,27 @@ public class MongoClient {
     /// Default maximum size for connection pools created by this client.
     internal static let defaultMaxConnectionPoolSize = 100
 
+    /// Indicates the state of this `MongoClient`.
+    private enum State {
+        /// The client is open.
+        case open
+        /// The client is in the process of closing.
+        case closing
+        /// The client has finished closing.
+        case closed
+    }
+
     /// Indicates whether this client has been closed.
-    internal private(set) var isClosed = false
+    private var state = State.open
+
+    /// Lock over `state`.
+    private var stateLock = Lock()
+
+    internal var isOpen: Bool {
+        self.stateLock.withLock {
+            self.state == .open
+        }
+    }
 
     /// Handlers for command monitoring events.
     internal var commandEventHandlers: [CommandEventHandler]
@@ -255,21 +276,46 @@ public class MongoClient {
     }
 
     deinit {
-        assert(self.isClosed, "MongoClient was not closed before deinitialization")
+        assert(
+            self.state == .closed,
+            "MongoClient was not closed before deinitialization. " +
+                "Please call `close()` or `syncClose()` when the client is no longer needed."
+        )
     }
 
-    /// Shuts this `MongoClient` down, closing all connection to the server and cleaning up internal state.
+    /// Closes this `MongoClient`, closing all connections to the server and cleaning up internal state.
     /// Call this method exactly once when you are finished using the client. You must ensure that all operations
     /// using the client have completed before calling this. The returned future must be fulfilled before the
     /// `EventLoopGroup` provided to this client's constructor is shut down.
-    public func shutdown() -> EventLoopFuture<Void> {
-        self.operationExecutor.execute {
-            self.connectionPool.shutdown()
-            self.isClosed = true
+    public func close() -> EventLoopFuture<Void> {
+        let stateError: Error? = self.stateLock.withLock {
+            switch self.state {
+            case .closing, .closed:
+                return Self.ClosedClientError
+            case .open:
+                self.state = .closing
+                return nil
+            }
+        }
+
+        if let stateError = stateError {
+            return self.operationExecutor.makeFailedFuture(stateError)
+        }
+
+        let closeResult = self.operationExecutor.execute {
+            try self.connectionPool.close()
         }
         .flatMap {
-            self.operationExecutor.close()
+            self.operationExecutor.shutdown()
         }
+
+        closeResult.whenComplete { _ in
+            self.stateLock.withLock {
+                self.state = .closed
+            }
+        }
+
+        return closeResult
     }
 
     /**
@@ -281,11 +327,21 @@ public class MongoClient {
      *
      * This method must complete before the `EventLoopGroup` provided to this client's constructor is shut down.
      */
-    public func syncShutdown() {
-        self.connectionPool.shutdown()
-        self.isClosed = true
-        // TODO: SWIFT-349 log any errors encountered here.
-        try? self.operationExecutor.syncClose()
+    public func syncClose() throws {
+        try self.stateLock.withLock {
+            switch self.state {
+            case .closing, .closed:
+                throw Self.ClosedClientError
+            case .open:
+                self.state = .closing
+            }
+        }
+        try self.connectionPool.close()
+        try self.operationExecutor.syncShutdown()
+
+        self.stateLock.withLock {
+            self.state = .closed
+        }
     }
 
     /// Starts a new `ClientSession` with the provided options. When you are done using this session, you must call
@@ -602,15 +658,6 @@ public class MongoClient {
         self.sdamEventHandlers.append(CallbackEventHandler(handlerFunc))
     }
 
-    /// Executes an `Operation` using this `MongoClient` and an optionally provided session.
-    internal func executeOperation<T: Operation>(
-        _ operation: T,
-        using connection: Connection? = nil,
-        session: ClientSession? = nil
-    ) throws -> T.OperationResult {
-        try self.operationExecutor.execute(operation, using: connection, client: self, session: session).wait()
-    }
-
     /// Internal method to check the `ReadConcern` that was ultimately set on this client. **This method may block
     /// and is for testing purposes only**.
     internal func getMongocReadConcern() throws -> ReadConcern? {