From 8dd16259287f58f9273002717ec4d27e97127719 Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Wed, 12 Jun 2024 07:43:14 +0200 Subject: Merging upstream version 127.0. Signed-off-by: Daniel Baumann --- third_party/rust/sql-support/src/lazy.rs | 151 +++++++++++++++++++++++++++++++ 1 file changed, 151 insertions(+) create mode 100644 third_party/rust/sql-support/src/lazy.rs (limited to 'third_party/rust/sql-support/src/lazy.rs') 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 { + 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>, + // 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>>, +} + +impl LazyDb { + /// 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, 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. 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 { + 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 + 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 { + 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(); + } +} -- cgit v1.2.3