Destroy open connections when resetting app DB; add `ReentrantReadWriteLock`...

Destroy open connections when resetting app DB; add `ReentrantReadWriteLock` around app DB `getConnection` (#21664)

* Destroy open connections when resetting app DB

* Use no-op dummy app DB

* Ok, don't shut down the app DB because it breaks in-memory app DBs

* Stop/start task scheduler

* Test fix ?

* Another test fix 

* Simplified impl

* Revert task changes

* Fix connection-pool destroying code

* Use ReentrantReadWriteLock for simpler implementation

* More doc

src/metabase/api/testing.clj

@@ -5,29 +5,30 @@
             [clojure.tools.logging :as log]
             [compojure.core :refer [POST]]
             [metabase.api.common :as api]
-            [metabase.models.setting.cache :as setting.cache]
-            [toucan.db :as db]))
+            [metabase.db.connection :as mdb.connection]
+            [metabase.util.files :as u.files])
+  (:import com.mchange.v2.c3p0.PoolBackedDataSource
+           java.util.concurrent.locks.ReentrantReadWriteLock))
+
+;; EVERYTHING BELOW IS FOR H2 ONLY.
+
+(defn- assert-h2 [app-db]
+  (assert (= (:db-type app-db) :h2)
+          "Snapshot/restore only works for :h2 application databases."))
 
 (defn- snapshot-path-for-name
   ^String [snapshot-name]
-  (str "frontend/test/snapshots/" (str/replace (name snapshot-name) #"\W" "_") ".sql"))
+  (let [path (u.files/get-path "frontend" "test" "snapshots"
+                               (str (str/replace (name snapshot-name) #"\W" "_") ".sql"))]
+    (str (.toAbsolutePath path))))
+
+;;;; SAVE
 
 (defn- save-snapshot! [snapshot-name]
+  (assert-h2 mdb.connection/*application-db*)
   (let [path (snapshot-path-for-name snapshot-name)]
     (log/infof "Saving snapshot to %s" path)
-    (jdbc/query (db/connection) ["SCRIPT TO ?" path]))
-  :ok)
-
-(defn- restore-snapshot! [snapshot-name]
-  (let [path (snapshot-path-for-name snapshot-name)]
-    (log/infof "Restoring snapshot from %s" path)
-    (api/check-404 (.exists (java.io.File. path)))
-    (with-open [conn (jdbc/get-connection (db/connection))]
-      (let [conn-spec {:connection conn}]
-        (jdbc/execute! conn-spec ["SET LOCK_TIMEOUT 180000"])
-        (jdbc/execute! conn-spec ["DROP ALL OBJECTS"])
-        (jdbc/execute! conn-spec ["RUNSCRIPT FROM ?" path]))))
-  (setting.cache/restore-cache!)
+    (jdbc/query {:datasource mdb.connection/*application-db*} ["SCRIPT TO ?" path]))
   :ok)
 
 (api/defendpoint POST "/snapshot/:name"
@@ -36,6 +37,49 @@
   (save-snapshot! name)
   nil)
 
+;;;; RESTORE
+
+(defn- reset-app-db-connection-pool!
+  "Immediately destroy all open connections in the app DB connection pool."
+  []
+  (let [{:keys [data-source]} mdb.connection/*application-db*]
+     (when (instance? PoolBackedDataSource data-source)
+       (log/info "Destroying application database connection pool")
+       (.hardReset ^PoolBackedDataSource data-source))))
+
+(defn- restore-app-db-from-snapshot!
+  "Drop all objects in the application DB, then reload everything from the SQL dump at `snapshot-path`."
+  [^String snapshot-path]
+  (log/infof "Restoring snapshot from %s" snapshot-path)
+  (api/check-404 (.exists (java.io.File. snapshot-path)))
+  (with-open [conn (.getConnection mdb.connection/*application-db*)]
+    (doseq [sql-args [["SET LOCK_TIMEOUT 180000"]
+                      ["DROP ALL OBJECTS"]
+                      ["RUNSCRIPT FROM ?" snapshot-path]]]
+      (jdbc/execute! {:connection conn} sql-args))))
+
+(defn- increment-app-db-unique-indentifier!
+  "Increment the [[mdb.connection/unique-identifier]] for the Metabase application DB. This effectively flushes all
+  caches using it as a key (including things using [[mdb.connection/memoize-for-application-db]]) such as the Settings
+  cache."
+  []
+  (alter-var-root #'mdb.connection/*application-db* assoc :id (swap! mdb.connection/application-db-counter inc)))
+
+(defn- restore-snapshot! [snapshot-name]
+  (assert-h2 mdb.connection/*application-db*)
+  (let [path                         (snapshot-path-for-name snapshot-name)
+        ^ReentrantReadWriteLock lock (:lock mdb.connection/*application-db*)]
+    ;; acquire the application DB WRITE LOCK which will prevent any other threads from getting any new connections until
+    ;; we release it.
+    (try
+      (.. lock writeLock lock)
+      (reset-app-db-connection-pool!)
+      (restore-app-db-from-snapshot! path)
+      (increment-app-db-unique-indentifier!)
+      (finally
+        (.. lock writeLock unlock))))
+  :ok)
+
 (api/defendpoint POST "/restore/:name"
   "Restore a database snapshot for testing purposes."
   [name]

src/metabase/db/connection.clj

@@ -3,11 +3,13 @@
    TODO - consider renaming this namespace `metabase.db.config`."
   (:require [metabase.db.connection-pool-setup :as connection-pool-setup]
             [metabase.db.env :as mdb.env]
-            [potemkin :as p]))
+            [potemkin :as p])
+  (:import java.util.concurrent.locks.ReentrantReadWriteLock))
 
-;; Counter for [[unique-identifier]] -- this is a simple counter rather that [[java.util.UUID/randomUUID]] so we don't
-;; waste precious entropy on launch generating something that doesn't need to be random (it just needs to be unique)
-(defonce ^:private application-db-counter
+(defonce ^{:doc "Counter for [[unique-identifier]] -- this is a simple counter rather that [[java.util.UUID/randomUUID]]
+  so we don't waste precious entropy on launch generating something that doesn't need to be random (it just needs to be
+  unique)"}
+  application-db-counter
   (atom 0))
 
 (p/defrecord+ ApplicationDB [^clojure.lang.Keyword db-type
@@ -18,13 +20,33 @@
                              ^clojure.lang.Atom    status
                              ;; A unique identifier generated for this specific application DB. Use this as a
                              ;; memoization/cache key. See [[unique-identifier]] for more information.
-                             id]
+                             id
+                             ;; Reentrant read-write lock for GETTING new connections. Lock doesn't track whether any
+                             ;; existing connections are open! Holding the write lock will however prevent any NEW
+                             ;; connections from being acquired.
+                             ;;
+                             ;; This is a reentrant read-write lock, which means any number of read locks are allowed at
+                             ;; the same time, but the write lock is exclusive. So if you want to prevent anyone from
+                             ;; getting new connections, lock the write lock.
+                             ;;
+                             ;; The main purpose of this is to power [[metabase.api.testing]] which allows you to reset
+                             ;; the application DB with data from a SQL dump -- during the restore process it is
+                             ;; important that we do not allow anyone to access the DB.
+                             ^ReentrantReadWriteLock lock]
   javax.sql.DataSource
   (getConnection [_]
-    (.getConnection data-source))
+    (try
+      (.. lock readLock lock)
+      (.getConnection data-source)
+      (finally
+        (.. lock readLock unlock))))
 
   (getConnection [_ user password]
-    (.getConnection data-source user password)))
+    (try
+      (.. lock readLock lock)
+      (.getConnection data-source user password)
+      (finally
+        (.. lock readLock unlock)))))
 
 (alter-meta! #'->ApplicationDB assoc :private true)
 (alter-meta! #'map->ApplicationDB assoc :private true)
@@ -52,7 +74,8 @@
                    data-source)
     :status      (atom nil)
     ;; for memoization purposes. See [[unique-identifier]] for more information.
-    :id          (swap! application-db-counter inc)}))
+    :id          (swap! application-db-counter inc)
+    :lock        (ReentrantReadWriteLock.)}))
 
 (def ^:dynamic ^ApplicationDB *application-db*
   "Type info and [[javax.sql.DataSource]] for the current Metabase application database. Create a new instance