1 files changed, 151 insertions, 0 deletions
diff --git a/third_party/rust/sql-support/src/lazy.rs b/third_party/rust/sql-support/src/lazy.rs
new file mode 100644
index 0000000000..b22d9c39e3
--- /dev/null
+++ b/third_party/rust/sql-support/src/lazy.rs
@@ -0,0 +1,151 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+use crate::open_database::{open_database_with_flags, ConnectionInitializer, Error};
+use interrupt_support::{register_interrupt, SqlInterruptHandle, SqlInterruptScope};
+use parking_lot::{MappedMutexGuard, Mutex, MutexGuard};
+use rusqlite::{Connection, OpenFlags};
+use std::{
+    path::{Path, PathBuf},
+    sync::{Arc, Weak},
+};
+
+/// Lazily-loaded database with interruption support
+///
+/// In addition to the [Self::interrupt] method, LazyDb also calls
+/// [interrupt_support::register_interrupt] on any opened database.  This means that if
+/// [interrupt_support::shutdown] is called it will interrupt this database if it's open and
+/// in-use.
+pub struct LazyDb<CI> {
+    path: PathBuf,
+    open_flags: OpenFlags,
+    connection_initializer: CI,
+    // Note: if you're going to lock both mutexes at once, make sure to lock the connection mutex
+    // first.  Otherwise, you risk creating a deadlock where two threads each hold one of the locks
+    // and is waiting for the other.
+    connection: Mutex<Option<Connection>>,
+    // It's important to use a separate mutex for the interrupt handle, since the whole point is to
+    // interrupt while another thread holds the connection mutex.  Since the only mutation is
+    // setting/unsetting the Arc, maybe this could be sped up by using something like
+    // `arc_swap::ArcSwap`, but that seems like overkill for our purposes. This mutex should rarely
+    // be contested and interrupt operations execute quickly.
+    interrupt_handle: Mutex<Option<Arc<SqlInterruptHandle>>>,
+}
+
+impl<CI: ConnectionInitializer> LazyDb<CI> {
+    /// Create a new LazyDb
+    ///
+    /// This does not open the connection and is non-blocking
+    pub fn new(path: &Path, open_flags: OpenFlags, connection_initializer: CI) -> Self {
+        Self {
+            path: path.to_owned(),
+            open_flags,
+            connection_initializer,
+            connection: Mutex::new(None),
+            interrupt_handle: Mutex::new(None),
+        }
+    }
+
+    /// Lock the database mutex and get a connection and interrupt scope.
+    ///
+    /// If the connection is closed, it will be opened.
+    ///
+    /// Calling `lock` again, or calling `close`, from the same thread while the mutex guard is
+    /// still alive will cause a deadlock.
+    pub fn lock(&self) -> Result<(MappedMutexGuard<'_, Connection>, SqlInterruptScope), Error> {
+        // Call get_conn first, then get_scope to ensure we acquire the locks in the correct order
+        let conn = self.get_conn()?;
+        let scope = self.get_scope(&conn)?;
+        Ok((conn, scope))
+    }
+
+    fn get_conn(&self) -> Result<MappedMutexGuard<'_, Connection>, Error> {
+        let mut guard = self.connection.lock();
+        // Open the database if it wasn't opened before.  Do this outside of the MutexGuard::map call to simplify the error handling
+        if guard.is_none() {
+            *guard = Some(open_database_with_flags(
+                &self.path,
+                self.open_flags,
+                &self.connection_initializer,
+            )?);
+        };
+        // Use MutexGuard::map to get a Connection rather than Option<Connection>.  The unwrap()
+        // call can't fail because of the previous code.
+        Ok(MutexGuard::map(guard, |conn_option| {
+            conn_option.as_mut().unwrap()
+        }))
+    }
+
+    fn get_scope(&self, conn: &Connection) -> Result<SqlInterruptScope, Error> {
+        let mut handle_guard = self.interrupt_handle.lock();
+        let result = match handle_guard.as_ref() {
+            Some(handle) => handle.begin_interrupt_scope(),
+            None => {
+                let handle = Arc::new(SqlInterruptHandle::new(conn));
+                register_interrupt(
+                    Arc::downgrade(&handle) as Weak<dyn AsRef<SqlInterruptHandle> + Send + Sync>
+                );
+                handle_guard.insert(handle).begin_interrupt_scope()
+            }
+        };
+        // If we see an Interrupted error when beginning the scope, it means that we're in shutdown
+        // mode.
+        result.map_err(|_| Error::Shutdown)
+    }
+
+    /// Close the database if it's open
+    ///
+    /// Pass interrupt=true to interrupt any in-progress queries before closing the database.
+    ///
+    /// Do not call `close` if you already have a lock on the database in the current thread, as
+    /// this will cause a deadlock.
+    pub fn close(&self, interrupt: bool) {
+        let mut interrupt_handle = self.interrupt_handle.lock();
+        if let Some(handle) = interrupt_handle.as_ref() {
+            if interrupt {
+                handle.interrupt();
+            }
+            *interrupt_handle = None;
+        }
+        // Drop the interrupt handle lock to avoid holding both locks at once.
+        drop(interrupt_handle);
+        *self.connection.lock() = None;
+    }
+
+    /// Interrupt any in-progress queries
+    pub fn interrupt(&self) {
+        if let Some(handle) = self.interrupt_handle.lock().as_ref() {
+            handle.interrupt();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::open_database::test_utils::TestConnectionInitializer;
+
+    fn open_test_db() -> LazyDb<TestConnectionInitializer> {
+        LazyDb::new(
+            Path::new(":memory:"),
+            OpenFlags::default(),
+            TestConnectionInitializer::new(),
+        )
+    }
+
+    #[test]
+    fn test_interrupt() {
+        let lazy_db = open_test_db();
+        let (_, scope) = lazy_db.lock().unwrap();
+        assert!(!scope.was_interrupted());
+        lazy_db.interrupt();
+        assert!(scope.was_interrupted());
+    }
+
+    #[test]
+    fn interrupt_before_db_is_opened_should_not_fail() {
+        let lazy_db = open_test_db();
+        lazy_db.interrupt();
+    }
+}