Refactor SQLiteFuture to allow it to be retried (#30)

Author: atdrendel

Package.resolved

@@ -15,7 +15,7 @@ let package = Package(
         .package(
             name: "CombineExtensions",
             url: "https://github.com/shareup/combine-extensions.git",
-            from: "2.9.0"
+            from: "3.1.0"
         ),
         .package(
             name: "Synchronized",
@@ -32,6 +32,7 @@ let package = Package(
         .testTarget(
             name: "SQLiteTests",
             dependencies: [
+                .product(name: "CombineExtensionsDynamic", package: "CombineExtensions"),
                 .product(name: "CombineTestExtensions", package: "CombineExtensions"),
                 "SQLite",
             ]),

Package.swift

@@ -179,7 +179,7 @@ extension SQLiteDatabase {
 
     public func readPublisher<T: SQLiteTransformable>(
         _ sql: SQL,
-        arguments: SQLiteArguments
+        arguments: SQLiteArguments = [:]
     ) -> AnyPublisher<[T], SQLiteError> {
         readPublisher(sql, arguments: arguments)
             .tryMap { try $0.map { try T.init(row: $0) } }
@@ -547,6 +547,7 @@ extension SQLiteDatabase {
 
     private func prepare(_ sql: SQL) throws -> SQLiteStatement {
         precondition(SQLiteQueue.isCurrentQueue)
+        guard _isOpen else { throw SQLiteError.databaseIsClosed }
         return try SQLiteStatement.prepare(sql, in: self)
     }
 }

Sources/SQLite/SQLiteDatabase.swift

@@ -89,6 +89,11 @@ extension SQLiteError {
         guard let code = code else { return false }
         return code == SQLITE_BUSY
     }
+
+    public var isClosed: Bool {
+        guard case .databaseIsClosed = self else { return false }
+        return true
+    }
 }
 
 extension SQLiteError: CustomStringConvertible {

Sources/SQLite/SQLiteError.swift

@@ -2,145 +2,53 @@ import Foundation
 import Combine
 import Synchronized
 
-class SQLiteFuture<Output>: Publisher {
+struct SQLiteFuture<Output>: Publisher {
     typealias Failure = SQLiteError
 
     typealias Promise = (Result<Output, Failure>) -> Void
 
-    private enum State {
-        case notStarted((@escaping SQLiteFuture.Promise) -> Void)
-        case inProgress
-        case failed(SQLiteError)
-        case finished(Output)
-
-        var result: Result<Output, SQLiteError>? {
-            switch self {
-            case .notStarted, .inProgress:
-                return nil
-
-            case let .failed(error):
-                return .failure(error)
-
-            case let .finished(output):
-                return .success(output)
-            }
-        }
-    }
-
-    private var state: State
-    private var subscriptions = Set<SQLiteFutureSubscription<Output>>()
-    private var lock = RecursiveLock()
+    private let attemptToFulfill: (@escaping Promise) -> Void
 
     init(
         _ attemptToFulfill: @escaping (@escaping Promise) -> Void
     ) {
-        self.state = .notStarted(attemptToFulfill)
+        self.attemptToFulfill = attemptToFulfill
     }
 
     func receive<S: Subscriber>(
         subscriber: S
     ) where Failure == S.Failure, Output == S.Input {
         let subscription = SQLiteFutureSubscription(
-            subscriber: subscriber,
-            resultProvider: resultProvider
+            attemptToFulfill: attemptToFulfill,
+            subscriber: subscriber
         )
-        lock.locked { let _ = subscriptions.insert(subscription) }
         subscriber.receive(subscription: subscription)
     }
+}
 
-    private func notifySubscribers() {
-        lock.locked {
-            guard let result = state.result else { return }
-
-            var notified = Set<SQLiteFutureSubscription<Output>>()
-            for sub in subscriptions {
-                if sub.sendResult(result) {
-                    notified.insert(sub)
-                }
-            }
-            subscriptions.subtract(notified)
-        }
-    }
-
-    enum NextAction {
-        case doNothing
-        case attemptToFulfill((@escaping Promise) -> Void)
-        case sendOutput(Output)
-        case sendFailure(SQLiteError)
+private final class SQLiteFutureSubscription<Output, S: Subscriber>: Subscription
+where
+    S.Input == Output,
+    S.Failure == SQLiteError
+{
+    private enum State {
+        case pending
+        case fulfilled(Result<S.Input, S.Failure>)
+        case finished
     }
 
-    private var resultProvider: () -> Result<Output, SQLiteError>? {
-        { [weak self] in
-            guard let self = self else { return nil }
-
-            let nextAction = self.lock.locked { () -> NextAction in
-                switch self.state {
-                case let .notStarted(attemptToFulfill):
-                    self.state = .inProgress
-                    return .attemptToFulfill(attemptToFulfill)
-
-                case .inProgress:
-                    return .doNothing
-
-                case let .finished(output):
-                    return .sendOutput(output)
-
-                case let .failed(error):
-                    return .sendFailure(error)
-                }
-            }
-
-            switch nextAction {
-            case let .attemptToFulfill(attemptToFulfill):
-                // We can't use `[weak self]` here or the block
-                // will be deallocated immediately after invoked.
-                // This should not create a retain cycle because
-                // `self` does not retain the block.
-                attemptToFulfill({ result in
-                    self.lock.locked {
-                        guard case .inProgress = self.state else { return }
-
-                        switch result {
-                        case let .success(output):
-                            self.state = .finished(output)
-
-                        case let .failure(error):
-                            self.state = .failed(error)
-                        }
-                    }
-
-                    self.notifySubscribers()
-                })
-                return nil
-
-            case .doNothing:
-                return nil
-
-            case let .sendOutput(output):
-                return .success(output)
-
-            case let .sendFailure(error):
-                return .failure(error)
-            }
-        }
-    }
-}
+    private var state: State = .pending
+    private var hasDemand = false
 
-private final class SQLiteFutureSubscription<Output>: Subscription, Hashable {
-    let id: CombineIdentifier
+    private var subscriber: S?
+    private let lock = Lock()
 
-    private var hasDemand = false
-    private var subscriber: AnySubscriber<Output, SQLiteError>?
-    private let resultProvider: () -> Result<Output, SQLiteError>?
-    private let lock = RecursiveLock()
-
-    init<S: Subscriber>(
-        subscriber: S,
-        resultProvider: @escaping () -> Result<Output, SQLiteError>?
-    ) where S.Input == Output, S.Failure == SQLiteError {
-        self.id = subscriber.combineIdentifier
-        self.subscriber = AnySubscriber(subscriber)
-        self.resultProvider = resultProvider
+    init(
+        attemptToFulfill: (@escaping SQLiteFuture<Output>.Promise) -> Void,
+        subscriber: S
+    ) {
+        self.subscriber = subscriber
+        attemptToFulfill({ result in self.fulfill(with: result) })
     }
 
     deinit {
@@ -149,50 +57,54 @@ private final class SQLiteFutureSubscription<Output>: Subscription, Hashable {
 
     func request(_ demand: Subscribers.Demand) {
         guard demand != .none else { return }
-        lock.locked { self.hasDemand = true }
-        guard let result = resultProvider() else { return }
-        let _ = sendResult(result)
-    }
 
-    func cancel() {
-        lock.locked { subscriber = nil }
-    }
-
-    func sendResult(_ result: Result<Output, SQLiteError>) -> Bool {
-        typealias Sub = AnySubscriber<Output, SQLiteError>
+        let subscriberAndResult: (S, Result<S.Input, S.Failure>)? = lock.locked {
+            hasDemand = true
+            guard let sub = subscriber, case let .fulfilled(result) = state
+            else { return nil }
+            state = .finished
+            subscriber = nil
+            return (sub, result)
+        }
 
-        let (didFinish, sub): (Bool, Sub?) = lock.locked {
-            // If we don't have any demand, we return `false` so that
-            // `SQLiteFuture` holds on to this subscription until it
-            // receives some demand and calls `resultProvider()`.
-            guard hasDemand else { return (false, nil) }
+        guard let (subscriber, result) = subscriberAndResult else { return }
+        notify(subscriber: subscriber, result: result)
+    }
 
-            // If we don't have a subscriber, we want to be removed from
-            // `SQLiteFuture`'s set of subscriptions. So, we need to
-            // return `true`.
-            guard let sub = subscriber else { return (true, nil) }
+    func cancel() {
+        lock.locked {
             subscriber = nil
-
-            return (true, sub)
+            state = .finished
         }
+    }
 
-        switch result {
-        case let .success(output):
-            let _ = sub?.receive(output)
-            sub?.receive(completion: .finished)
-
-        case let .failure(error):
-            sub?.receive(completion: .failure(error))
+    private func fulfill(with result: Result<S.Input, S.Failure>) {
+        let _subscriber: S? = lock.locked {
+            guard case .pending = state, let sub = subscriber
+            else { return nil }
+
+            if hasDemand {
+                state = .finished
+                subscriber = nil
+                return sub
+            } else {
+                state = .fulfilled(result)
+                return nil
+            }
         }
 
-        return didFinish
+        guard let subscriber = _subscriber else { return }
+        notify(subscriber: subscriber, result: result)
     }
 
-    func hash(into hasher: inout Hasher) {
-        hasher.combine(id)
-    }
+    private func notify(subscriber: S, result: Result<S.Input, S.Failure>) {
+        switch result {
+        case let .success(rows):
+            _ = subscriber.receive(rows)
+            subscriber.receive(completion: .finished)
 
-    static func == (lhs: SQLiteFutureSubscription, rhs: SQLiteFutureSubscription) -> Bool {
-        lhs.id == rhs.id
+        case let .failure(error):
+            subscriber.receive(completion: .failure(error))
+        }
     }
 }

Sources/SQLite/SQLiteFuture.swift

@@ -8,8 +8,8 @@ extension Publisher where Output == SQLiteDatabaseChange, Failure == Never {
         sql: SQL,
         arguments: SQLiteArguments = [:],
         database: SQLiteDatabase
-    ) -> SQLiteStatementResultsPublisher<Self> {
-        SQLiteStatementResultsPublisher(
+    ) -> SQLiteResultsPublisher<Self> {
+        SQLiteResultsPublisher(
             upstream: self,
             sql: sql,
             arguments: arguments,
@@ -18,7 +18,7 @@ extension Publisher where Output == SQLiteDatabaseChange, Failure == Never {
     }
 }
 
-struct SQLiteStatementResultsPublisher<Upstream: Publisher>: Publisher where
+struct SQLiteResultsPublisher<Upstream: Publisher>: Publisher where
     Upstream.Output == SQLiteDatabaseChange, Upstream.Failure == Never
 {
     typealias Output = [SQLiteRow]

…te/SQLiteStatementResultsPublisher.swift …rces/SQLite/SQLiteResultsPublisher.swiftSources/SQLite/SQLiteStatementResultsPublisher.swift renamed to Sources/SQLite/SQLiteResultsPublisher.swift Sources/SQLite/SQLiteStatementResultsPublisher.swift renamed to Sources/SQLite/SQLiteResultsPublisher.swift

@@ -20,6 +20,7 @@ class SQLiteCrossProcessMonitorTests: XCTestCase {
             let db = try SQLiteDatabase.makeShared(path: dbPath)
             try db.write(createTable)
             XCTAssertEqual(["test"], try db.tables())
+            try db.close()
         }
     }
 
@@ -48,6 +49,9 @@ class SQLiteCrossProcessMonitorTests: XCTestCase {
             )
 
             wait(for: [ex1, ex2], timeout: 4) // Coordinated writes can be very slow
+            
+            try db1.close()
+            try db2.close()
         }
     }
 
@@ -90,6 +94,8 @@ class SQLiteCrossProcessMonitorTests: XCTestCase {
 
             wait(for: [outputEx], timeout: 2)
             wait(for: [duplicateEx], timeout: 2) // Coordinated writes can be very slow
+
+            try db.close()
         }
     }
 }

Tests/SQLiteTests/SQLiteCrossProcessMonitorTests.swift

@@ -77,6 +77,8 @@ class SQLiteDatabaseTests: XCTestCase {
             XCTAssertEqual(2, rows.count)
             XCTAssertEqual(one, rows.first)
             XCTAssertEqual(three, rows.last)
+
+            try db.close()
         }
     }