From 3c83b1d13ade89d5eff86a7941eaf1aee8d3ece6 Mon Sep 17 00:00:00 2001
From: Dapeng Gao <peter@dpgao.cc>
Date: Sun, 6 Jan 2019 18:18:02 +0800
Subject: [PATCH] Allow separate connections to share the same in-memory
 database

---
 .../SQLite/Database/SQLiteConnection.swift    | 12 +-----
 Sources/SQLite/Database/SQLiteDatabase.swift  | 37 +++++++++++++++++++
 Sources/SQLite/Database/SQLiteStorage.swift   |  7 ----
 3 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/Sources/SQLite/Database/SQLiteConnection.swift b/Sources/SQLite/Database/SQLiteConnection.swift
index 277df51..2804bc0 100644
--- a/Sources/SQLite/Database/SQLiteConnection.swift
+++ b/Sources/SQLite/Database/SQLiteConnection.swift
@@ -33,7 +33,7 @@ public final class SQLiteConnection: BasicWorker, DatabaseConnection, DatabaseQu
     
     /// Reference to parent `SQLiteDatabase` that created this connection.
     /// This reference will ensure the DB stays alive since this connection uses
-    /// it's dispatch queue.
+    /// it's thread pool.
     private let database: SQLiteDatabase
     
     /// Internal SQLite database handle.
@@ -45,15 +45,7 @@ public final class SQLiteConnection: BasicWorker, DatabaseConnection, DatabaseQu
     /// Create a new SQLite conncetion.
     internal init(database: SQLiteDatabase, on worker: Worker) throws {
         self.database = database
-        // Make database connection
-        var handle: OpaquePointer?
-        let options = SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX
-        guard sqlite3_open_v2(database.storage.path, &handle, options, nil) == SQLITE_OK,
-            let c = handle,
-            sqlite3_busy_handler(c, { _, _ in 1 }, nil) == SQLITE_OK else {
-                throw SQLiteError(problem: .error, reason: "Could not open database.", source: .capture())
-        }
-        self.handle = c
+        self.handle = try database.openConnection()
         self.eventLoop = worker.eventLoop
     }
     
diff --git a/Sources/SQLite/Database/SQLiteDatabase.swift b/Sources/SQLite/Database/SQLiteDatabase.swift
index e500443..12fcd66 100644
--- a/Sources/SQLite/Database/SQLiteDatabase.swift
+++ b/Sources/SQLite/Database/SQLiteDatabase.swift
@@ -20,6 +20,10 @@ public final class SQLiteDatabase: Database, LogSupporting {
     /// Thread pool for performing blocking IO work. See `BlockingIOThreadPool`.
     internal let blockingIO: BlockingIOThreadPool
     
+    /// If the database uses in-memory storage, this property will be set to
+    /// keep the database alive when there is no `SQLiteConnection` to it.
+    private var handle: OpaquePointer?
+    
     /// Create a new SQLite database.
     ///
     ///     let sqliteDB = SQLiteDatabase(storage: .memory)
@@ -31,6 +35,28 @@ public final class SQLiteDatabase: Database, LogSupporting {
     public init(storage: SQLiteStorage = .memory, threadPool: BlockingIOThreadPool? = nil) throws {
         self.storage = storage
         self.blockingIO = threadPool ?? BlockingIOThreadPool(numberOfThreads: 2)
+        if case .memory = storage {
+            self.handle = try openConnection()
+        }
+    }
+    
+    // Make database connection
+    internal func openConnection() throws -> OpaquePointer {
+        let path: String
+        switch storage {
+        case .memory:
+            path = "file:\(ObjectIdentifier(self))?mode=memory&cache=shared"
+        case .file(let file):
+            path = file
+        }
+        var handle: OpaquePointer?
+        let options = SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX
+        guard sqlite3_open_v2(path, &handle, options, nil) == SQLITE_OK,
+            let c = handle,
+            sqlite3_busy_handler(c, { _, _ in 1 }, nil) == SQLITE_OK else {
+                throw SQLiteError(problem: .error, reason: "Could not open database.", source: .capture())
+        }
+        return c
     }
     
     /// See `Database`.
@@ -47,4 +73,15 @@ public final class SQLiteDatabase: Database, LogSupporting {
     public static func enableLogging(_ logger: DatabaseLogger, on conn: SQLiteConnection) {
         conn.logger = logger
     }
+    
+    deinit {
+        self.blockingIO.shutdownGracefully { [handle] error in
+            if let error = error {
+                print("[SQLite] [ERROR] Could not shutdown BlockingIOThreadPool: \(error)")
+            }
+            if let handle = handle, sqlite3_close(handle) != SQLITE_OK {
+                print("[SQLite] [ERROR] Could not close database.")
+            }
+        }
+    }
 }
diff --git a/Sources/SQLite/Database/SQLiteStorage.swift b/Sources/SQLite/Database/SQLiteStorage.swift
index ff26161..7ee818f 100644
--- a/Sources/SQLite/Database/SQLiteStorage.swift
+++ b/Sources/SQLite/Database/SQLiteStorage.swift
@@ -6,11 +6,4 @@ public enum SQLiteStorage {
     
     /// File-based storage, persisted between application launches.
     case file(path: String)
-
-    internal var path: String {
-        switch self {
-        case .memory: return ":memory:"
-        case .file(let path): return path
-        }
-    }
 }