SWIFT-674 Make MongoSwift.ClientSession async, and implement MongoSwiftSync.ClientSession (#391)

Author: kmahar

.swiftlint.yml

@@ -25,7 +25,6 @@ opt_in_rules:
   - missing_docs
   - modifier_order
   - multiline_arguments
-  - multiline_function_chains
   - multiline_literal_brackets
   - multiline_parameters
   - operator_usage_whitespace

Package.resolved

@@ -1,5 +1,6 @@
 import CLibMongoC
 import Foundation
+import NIO
 
 /**
  * A MongoDB client session.
@@ -111,23 +112,23 @@ public final class ClientSession {
     public let options: ClientSessionOptions?
 
     /// Initializes a new client session.
-    internal init(client: MongoClient, options: ClientSessionOptions? = nil) throws {
+    internal init(client: MongoClient, options: ClientSessionOptions? = nil) {
         self.options = options
         self.client = client
         self.state = .notStarted(opTime: nil, clusterTime: nil)
     }
 
     /// Starts this session's corresponding libmongoc session, if it has not been started already. Throws an error if
     /// this session has already been ended.
-    internal func startIfNeeded() throws {
+    internal func startIfNeeded() -> EventLoopFuture<Void> {
         switch self.state {
         case .notStarted:
             let operation = StartSessionOperation(session: self)
-            try self.client.executeOperation(operation)
+            return self.client.operationExecutor.execute(operation, client: self.client, session: nil)
         case .started:
-            return
+            return self.client.operationExecutor.makeSucceededFuture(Void())
         case .ended:
-            throw ClientSession.SessionInactiveError
+            return self.client.operationExecutor.makeFailedFuture(ClientSession.SessionInactiveError)
         }
     }
 
@@ -143,18 +144,28 @@ public final class ClientSession {
         return connection
     }
 
-    /// Destroy the underlying `mongoc_client_session_t` and ends this session. Has no effect if this session is
-    /// already ended.
-    internal func end() {
-        if case let .started(session, _) = self.state {
-            mongoc_client_session_destroy(session)
+    /// Ends this `ClientSession`. Call this method when you are finished using the session. You must ensure that all
+    /// operations using this session have completed before calling this. The returned future must be fulfilled before
+    /// this session's parent `MongoClient` is closed.
+    public func end() -> EventLoopFuture<Void> {
+        switch self.state {
+        case .notStarted, .ended:
+            self.state = .ended
+            return self.client.operationExecutor.makeSucceededFuture(Void())
+        case let .started(session, _):
+            return self.client.operationExecutor.execute {
+                mongoc_client_session_destroy(session)
+                self.state = .ended
+            }
         }
-        self.state = .ended
     }
 
     /// Cleans up internal state.
     deinit {
-        self.end()
+        guard case .ended = self.state else {
+            assertionFailure("ClientSession was not ended before going out of scope; please call ClientSession.end()")
+            return
+        }
     }
 
     /**

Sources/MongoSwift/ClientSession.swift

@@ -292,7 +292,7 @@ public class MongoClient {
 
     /// Closes this `MongoClient`. 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 client goes out of scope.
+    /// fulfilled before the `EventLoopGroup` provided to this client's constructor is shut down.
     public func close() -> EventLoopFuture<Void> {
         return self.operationExecutor.execute {
             self.connectionPool.close()
@@ -303,14 +303,9 @@ public class MongoClient {
         }
     }
 
-    /**
-     * Starts a new `ClientSession` with the provided options.
-     *
-     * - Throws:
-     *   - `RuntimeError.compatibilityError` if the deployment does not support sessions.
-     */
-    public func startSession(options: ClientSessionOptions? = nil) throws -> ClientSession {
-        return try ClientSession(client: self, options: options)
+    /// Starts a new `ClientSession` with the provided options.
+    public func startSession(options: ClientSessionOptions? = nil) -> ClientSession {
+        return ClientSession(client: self, options: options)
     }
 
     /**
@@ -322,11 +317,30 @@ public class MongoClient {
      */
     public func withSession<T>(
         options: ClientSessionOptions? = nil,
-        _ sessionBody: (ClientSession) throws -> T
-    ) throws -> T {
-        let session = try ClientSession(client: self, options: options)
-        defer { session.end() }
-        return try sessionBody(session)
+        _ sessionBody: (ClientSession) throws -> EventLoopFuture<T>
+    ) -> EventLoopFuture<T> {
+        let promise = self.operationExecutor.makePromise(of: T.self)
+        let session = self.startSession(options: options)
+        do {
+            let bodyFuture = try sessionBody(session)
+            // regardless of whether body's returned future succeeds we want to call session.end() once its complete.
+            // only once session.end() finishes can we fulfill the returned promise. otherwise the user can't tell if
+            // it is safe to close the parent client of this session, and they could inadvertently close it before the
+            // session is actually ended and its parent `mongoc_client_t` is returned to the pool.
+            bodyFuture.flatMap { _ in
+                session.end()
+            }.flatMapError { _ in
+                session.end()
+            }.whenComplete { _ in
+                promise.completeWith(bodyFuture)
+            }
+        } catch {
+            session.end().whenComplete { _ in
+                promise.fail(error)
+            }
+        }
+
+        return promise.futureResult
     }
 
     /**

Sources/MongoSwift/MongoClient.swift

@@ -47,34 +47,46 @@ internal class OperationExecutor {
             return self.makeFailedFuture(MongoClient.ClosedClientError)
         }
 
+        let doOperation = { () -> T.OperationResult in
+            // select a connection in following order of priority:
+            // 1. connection specifically provided for use with this operation
+            // 2. if a session was provided, use its underlying connection
+            // 3. a new connection from the pool
+            let connection = try connection ?? resolveConnection(client: client, session: session)
+            return try operation.execute(using: connection, session: session)
+        }
+
         if let session = session {
             if case .ended = session.state {
                 return self.makeFailedFuture(ClientSession.SessionInactiveError)
             }
             guard session.client == client else {
                 return self.makeFailedFuture(ClientSession.ClientMismatchError)
             }
-        }
 
-        return self.execute {
-            // if a session was provided, start it if it hasn't been started already.
-            try session?.startIfNeeded()
-            // select a connection in following order of priority:
-            // 1. connection specifically provided for use with this operation
-            // 2. if a session was provided, use its underlying connection
-            // 3. a new connection from the pool
-            let connection = try connection ?? resolveConnection(client: client, session: session)
-            return try operation.execute(using: connection, session: session)
+            // start the session if needed (which generates a new operation itself), and then execute the operation.
+            return session.startIfNeeded().flatMap { self.execute(doOperation) }
         }
+
+        // no session was provided, so we can just jump to executing the operation.
+        return self.execute(doOperation)
     }
 
-    internal func execute<T>(body: @escaping () throws -> T) -> EventLoopFuture<T> {
+    internal func execute<T>(_ body: @escaping () throws -> T) -> EventLoopFuture<T> {
         return self.threadPool.runIfActive(eventLoop: self.eventLoopGroup.next(), body)
     }
 
     internal func makeFailedFuture<T>(_ error: Error) -> EventLoopFuture<T> {
         return self.eventLoopGroup.next().makeFailedFuture(error)
     }
+
+    internal func makeSucceededFuture<T>(_ value: T) -> EventLoopFuture<T> {
+        return self.eventLoopGroup.next().makeSucceededFuture(value)
+    }
+
+    internal func makePromise<T>(of type: T.Type) -> EventLoopPromise<T> {
+        return self.eventLoopGroup.next().makePromise(of: type)
+    }
 }
 
 /// Given a client and optionally a session associated which are to be associated with an operation, returns a

Sources/MongoSwift/Operations/Operation.swift

@@ -37,29 +37,40 @@ public final class ClientSession {
     /// The most recent cluster time seen by this session. This value will be nil if either of the following are true:
     /// - No operations have been executed using this session and `advanceClusterTime` has not been called.
     /// - This session has been ended.
-    public var clusterTime: Document? {
-        fatalError("unimplemented")
-    }
+    public var clusterTime: Document? { return self.asyncSession.clusterTime }
 
     /// The operation time of the most recent operation performed using this session. This value will be nil if either
     /// of the following are true:
     /// - No operations have been performed using this session and `advanceOperationTime` has not been called.
     /// - This session has been ended.
-    public var operationTime: Timestamp? {
-        fatalError("unimplemented")
-    }
+    public var operationTime: Timestamp? { return self.asyncSession.operationTime }
 
     /// The options used to start this session.
-    public let options: ClientSessionOptions?
+    public var options: ClientSessionOptions? { return self.asyncSession.options }
 
     /// Initializes a new client session.
-    internal init(wrapping session: MongoSwift.ClientSession) throws {
-        fatalError("unimplemented")
+    internal init(client: MongoClient, options: ClientSessionOptions?) {
+        self.client = client
+        self.asyncSession = client.asyncClient.startSession(options: options)
+    }
+
+    /// Ends the underlying async session.
+    internal func end() {
+        // we only call this method from places that we can't throw (deinit, defers) so we handle the error here
+        // instead. the async method will only fail if the async client, thread pool, or event loop group have been
+        // closed/ended. we manage the lifetimes of all of those ourselves, so if we hit the assertionFailure it's due
+        // to a bug in our own code.
+        do {
+            try self.asyncSession.end().wait()
+        } catch {
+            assertionFailure("Error ending async session: \(error)")
+        }
     }
 
     /// Cleans up internal state.
     deinit {
-        fatalError("unimplemented")
+        // a repeated call to `end` is a no-op so it's ok to call this even if `end()` was already called explicitly.
+        self.end()
     }
 
     /**
@@ -71,7 +82,7 @@ public final class ClientSession {
      *   - clusterTime: The session's new cluster time, as a `Document` like `["cluster time": Timestamp(...)]`
      */
     public func advanceClusterTime(to clusterTime: Document) {
-        fatalError("unimplemented")
+        self.asyncSession.advanceClusterTime(to: clusterTime)
     }
 
     /**
@@ -83,6 +94,6 @@ public final class ClientSession {
      *   - operationTime: The session's new operationTime
      */
     public func advanceOperationTime(to operationTime: Timestamp) {
-        fatalError("unimplemented")
+        self.asyncSession.advanceOperationTime(to: operationTime)
     }
 }

Sources/MongoSwiftSync/ClientSession.swift

@@ -22,7 +22,7 @@ public class MongoClient {
     private let eventLoopGroup: MultiThreadedEventLoopGroup
 
     /// The underlying async client.
-    private let asyncClient: MongoSwift.MongoClient
+    internal let asyncClient: MongoSwift.MongoClient
 
     /**
      * Create a new client connection to a MongoDB server. For options that included in both the connection string URI
@@ -65,14 +65,9 @@ public class MongoClient {
         }
     }
 
-    /**
-     * Starts a new `ClientSession` with the provided options.
-     *
-     * - Throws:
-     *   - `RuntimeError.compatibilityError` if the deployment does not support sessions.
-     */
-    public func startSession(options: ClientSessionOptions? = nil) throws -> ClientSession {
-        fatalError("unimplemented")
+    /// Starts a new `ClientSession` with the provided options.
+    public func startSession(options: ClientSessionOptions? = nil) -> ClientSession {
+        return ClientSession(client: self, options: options)
     }
 
     /**
@@ -85,8 +80,10 @@ public class MongoClient {
     public func withSession<T>(
         options: ClientSessionOptions? = nil,
         _ sessionBody: (ClientSession) throws -> T
-    ) throws -> T {
-        fatalError("unimplemented")
+    ) rethrows -> T {
+        let session = self.startSession(options: options)
+        defer { session.end() }
+        return try sessionBody(session)
     }
 
     /**
@@ -160,7 +157,7 @@ public class MongoClient {
      * - Returns: a `MongoDatabase` corresponding to the provided database name
      */
     public func db(_ name: String, options: DatabaseOptions? = nil) -> MongoDatabase {
-        return MongoDatabase(asyncDB: self.asyncClient.db(name, options: options))
+        return MongoDatabase(client: self, asyncDB: self.asyncClient.db(name, options: options))
     }
 
     /**

Sources/MongoSwiftSync/MongoClient.swift

@@ -38,8 +38,13 @@ public struct MongoCollection<T: Codable> {
     /// The underlying asynchronous collection.
     internal let asyncColl: MongoSwift.MongoCollection<T>
 
+    /// The client this collection was derived from. We store this to ensure it remains open for as long as this object
+    /// is in scope.
+    private let client: MongoClient
+
     /// Initializes a new `MongoCollection` instance wrapping the provided async collection.
-    internal init(asyncCollection: MongoSwift.MongoCollection<T>) {
+    internal init(client: MongoClient, asyncCollection: MongoSwift.MongoCollection<T>) {
+        self.client = client
         self.asyncColl = asyncCollection
     }
 

Sources/MongoSwiftSync/MongoCollection.swift

@@ -24,8 +24,13 @@ public struct MongoDatabase {
     /// The underlying asynchronous database.
     private let asyncDB: MongoSwift.MongoDatabase
 
+    /// The client this database was derived from. We store this to ensure it remains open for as long as this object
+    /// is in scope.
+    private let client: MongoClient
+
     /// Initializes a new `MongoDatabase` instance wrapping the provided async database.
-    internal init(asyncDB: MongoSwift.MongoDatabase) {
+    internal init(client: MongoClient, asyncDB: MongoSwift.MongoDatabase) {
+        self.client = client
         self.asyncDB = asyncDB
     }
 
@@ -78,7 +83,7 @@ public struct MongoDatabase {
         options: CollectionOptions? = nil
     ) -> MongoCollection<T> {
         let asyncColl = self.asyncDB.collection(name, withType: T.self, options: options)
-        return MongoCollection(asyncCollection: asyncColl)
+        return MongoCollection(client: self.client, asyncCollection: asyncColl)
     }
 
     /**
@@ -135,7 +140,7 @@ public struct MongoDatabase {
                                                           options: options,
                                                           session: session?.asyncSession)
                                                           .wait()
-        return MongoCollection(asyncCollection: asyncColl)
+        return MongoCollection(client: self.client, asyncCollection: asyncColl)
     }
 
     /**