Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Reworked thread safe connection be threadsafer,,,, again

Browse source 
Co-Authored-By: kay@zed.dev
This commit is contained in:
Mikayla Maki 2022-12-01 18:31:05 -08:00
parent 189a820113
commit 5e240f98f0
12 changed files with 741 additions and 584 deletions
Download patch file Download diff file Expand all files Collapse all files

164
crates/sqlez/src/bindable.rs
View file

@ -137,13 +137,6 @@ impl Column for usize {
}
}
impl Bind for () {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
statement.bind_null(start_index)?;
Ok(start_index + 1)
}
}
impl Bind for &str {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
statement.bind_text(start_index, self)?;
@ -179,78 +172,6 @@ impl Column for String {
}
}
impl<T1: Bind, T2: Bind> Bind for (T1, T2) {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
let next_index = self.0.bind(statement, start_index)?;
self.1.bind(statement, next_index)
}
}
impl<T1: Column, T2: Column> Column for (T1, T2) {
fn column<'a>(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
let (first, next_index) = T1::column(statement, start_index)?;
let (second, next_index) = T2::column(statement, next_index)?;
Ok(((first, second), next_index))
}
}
impl<T1: Bind, T2: Bind, T3: Bind> Bind for (T1, T2, T3) {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
let next_index = self.0.bind(statement, start_index)?;
let next_index = self.1.bind(statement, next_index)?;
self.2.bind(statement, next_index)
}
}
impl<T1: Column, T2: Column, T3: Column> Column for (T1, T2, T3) {
fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
let (first, next_index) = T1::column(statement, start_index)?;
let (second, next_index) = T2::column(statement, next_index)?;
let (third, next_index) = T3::column(statement, next_index)?;
Ok(((first, second, third), next_index))
}
}
impl<T1: Bind, T2: Bind, T3: Bind, T4: Bind> Bind for (T1, T2, T3, T4) {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
let next_index = self.0.bind(statement, start_index)?;
let next_index = self.1.bind(statement, next_index)?;
let next_index = self.2.bind(statement, next_index)?;
self.3.bind(statement, next_index)
}
}
impl<T1: Column, T2: Column, T3: Column, T4: Column> Column for (T1, T2, T3, T4) {
fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
let (first, next_index) = T1::column(statement, start_index)?;
let (second, next_index) = T2::column(statement, next_index)?;
let (third, next_index) = T3::column(statement, next_index)?;
let (fourth, next_index) = T4::column(statement, next_index)?;
Ok(((first, second, third, fourth), next_index))
}
}
impl<T1: Bind, T2: Bind, T3: Bind, T4: Bind, T5: Bind> Bind for (T1, T2, T3, T4, T5) {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
let next_index = self.0.bind(statement, start_index)?;
let next_index = self.1.bind(statement, next_index)?;
let next_index = self.2.bind(statement, next_index)?;
let next_index = self.3.bind(statement, next_index)?;
self.4.bind(statement, next_index)
}
}
impl<T1: Column, T2: Column, T3: Column, T4: Column, T5: Column> Column for (T1, T2, T3, T4, T5) {
fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
let (first, next_index) = T1::column(statement, start_index)?;
let (second, next_index) = T2::column(statement, next_index)?;
let (third, next_index) = T3::column(statement, next_index)?;
let (fourth, next_index) = T4::column(statement, next_index)?;
let (fifth, next_index) = T5::column(statement, next_index)?;
Ok(((first, second, third, fourth, fifth), next_index))
}
}
impl<T: Bind> Bind for Option<T> {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
if let Some(this) = self {
@ -344,3 +265,88 @@ impl Column for PathBuf {
))
}
}
/// Unit impls do nothing. This simplifies query macros
impl Bind for () {
fn bind(&self, _statement: &Statement, start_index: i32) -> Result<i32> {
Ok(start_index)
}
}
impl Column for () {
fn column(_statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
Ok(((), start_index))
}
}
impl<T1: Bind, T2: Bind> Bind for (T1, T2) {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
let next_index = self.0.bind(statement, start_index)?;
self.1.bind(statement, next_index)
}
}
impl<T1: Column, T2: Column> Column for (T1, T2) {
fn column<'a>(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
let (first, next_index) = T1::column(statement, start_index)?;
let (second, next_index) = T2::column(statement, next_index)?;
Ok(((first, second), next_index))
}
}
impl<T1: Bind, T2: Bind, T3: Bind> Bind for (T1, T2, T3) {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
let next_index = self.0.bind(statement, start_index)?;
let next_index = self.1.bind(statement, next_index)?;
self.2.bind(statement, next_index)
}
}
impl<T1: Column, T2: Column, T3: Column> Column for (T1, T2, T3) {
fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
let (first, next_index) = T1::column(statement, start_index)?;
let (second, next_index) = T2::column(statement, next_index)?;
let (third, next_index) = T3::column(statement, next_index)?;
Ok(((first, second, third), next_index))
}
}
impl<T1: Bind, T2: Bind, T3: Bind, T4: Bind> Bind for (T1, T2, T3, T4) {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
let next_index = self.0.bind(statement, start_index)?;
let next_index = self.1.bind(statement, next_index)?;
let next_index = self.2.bind(statement, next_index)?;
self.3.bind(statement, next_index)
}
}
impl<T1: Column, T2: Column, T3: Column, T4: Column> Column for (T1, T2, T3, T4) {
fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
let (first, next_index) = T1::column(statement, start_index)?;
let (second, next_index) = T2::column(statement, next_index)?;
let (third, next_index) = T3::column(statement, next_index)?;
let (fourth, next_index) = T4::column(statement, next_index)?;
Ok(((first, second, third, fourth), next_index))
}
}
impl<T1: Bind, T2: Bind, T3: Bind, T4: Bind, T5: Bind> Bind for (T1, T2, T3, T4, T5) {
fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
let next_index = self.0.bind(statement, start_index)?;
let next_index = self.1.bind(statement, next_index)?;
let next_index = self.2.bind(statement, next_index)?;
let next_index = self.3.bind(statement, next_index)?;
self.4.bind(statement, next_index)
}
}
impl<T1: Column, T2: Column, T3: Column, T4: Column, T5: Column> Column for (T1, T2, T3, T4, T5) {
fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
let (first, next_index) = T1::column(statement, start_index)?;
let (second, next_index) = T2::column(statement, next_index)?;
let (third, next_index) = T3::column(statement, next_index)?;
let (fourth, next_index) = T4::column(statement, next_index)?;
let (fifth, next_index) = T5::column(statement, next_index)?;
Ok(((first, second, third, fourth, fifth), next_index))
}
}

14
crates/sqlez/src/connection.rs
View file

@ -1,4 +1,5 @@
use std::{
cell::RefCell,
ffi::{CStr, CString},
marker::PhantomData,
path::Path,
@ -11,7 +12,7 @@ use libsqlite3_sys::*;
pub struct Connection {
pub(crate) sqlite3: *mut sqlite3,
persistent: bool,
pub(crate) write: bool,
pub(crate) write: RefCell<bool>,
_sqlite: PhantomData<sqlite3>,
}
unsafe impl Send for Connection {}
@ -21,7 +22,7 @@ impl Connection {
let mut connection = Self {
sqlite3: 0 as *mut _,
persistent,
write: true,
write: RefCell::new(true),
_sqlite: PhantomData,
};
@ -64,7 +65,7 @@ impl Connection {
}
pub fn can_write(&self) -> bool {
self.write
*self.write.borrow()
}
pub fn backup_main(&self, destination: &Connection) -> Result<()> {
@ -152,6 +153,13 @@ impl Connection {
))
}
}
pub(crate) fn with_write<T>(&self, callback: impl FnOnce(&Connection) -> T) -> T {
*self.write.borrow_mut() = true;
let result = callback(self);
*self.write.borrow_mut() = false;
result
}
}
impl Drop for Connection {

4
crates/sqlez/src/domain.rs
View file

@ -1,11 +1,11 @@
use crate::connection::Connection;
pub trait Domain {
pub trait Domain: 'static {
fn name() -> &'static str;
fn migrations() -> &'static [&'static str];
}
pub trait Migrator {
pub trait Migrator: 'static {
fn migrate(connection: &Connection) -> anyhow::Result<()>;
}

3
crates/sqlez/src/migrations.rs
View file

@ -12,6 +12,7 @@ use crate::connection::Connection;
impl Connection {
pub fn migrate(&self, domain: &'static str, migrations: &[&'static str]) -> Result<()> {
self.with_savepoint("migrating", || {
println!("Processing domain");
// Setup the migrations table unconditionally
self.exec(indoc! {"
CREATE TABLE IF NOT EXISTS migrations (
@ -43,11 +44,13 @@ impl Connection {
{}", domain, index, completed_migration, migration}));
} else {
// Migration already run. Continue
println!("Migration already run");
continue;
}
}
self.exec(migration)?()?;
println!("Ran migration");
store_completed_migration((domain, index, *migration))?;
}

143
crates/sqlez/src/thread_safe_connection.rs
View file

@ -5,17 +5,13 @@ use parking_lot::{Mutex, RwLock};
use std::{collections::HashMap, marker::PhantomData, ops::Deref, sync::Arc, thread};
use thread_local::ThreadLocal;
use crate::{
connection::Connection,
domain::{Domain, Migrator},
util::UnboundedSyncSender,
};
use crate::{connection::Connection, domain::Migrator, util::UnboundedSyncSender};
const MIGRATION_RETRIES: usize = 10;
type QueuedWrite = Box<dyn 'static + Send + FnOnce(&Connection)>;
type QueuedWrite = Box<dyn 'static + Send + FnOnce()>;
type WriteQueueConstructor =
Box<dyn 'static + Send + FnMut(Connection) -> Box<dyn 'static + Send + Sync + Fn(QueuedWrite)>>;
Box<dyn 'static + Send + FnMut() -> Box<dyn 'static + Send + Sync + Fn(QueuedWrite)>>;
lazy_static! {
/// List of queues of tasks by database uri. This lets us serialize writes to the database
/// and have a single worker thread per db file. This means many thread safe connections
@ -28,18 +24,18 @@ lazy_static! {
/// Thread safe connection to a given database file or in memory db. This can be cloned, shared, static,
/// whatever. It derefs to a synchronous connection by thread that is read only. A write capable connection
/// may be accessed by passing a callback to the `write` function which will queue the callback
pub struct ThreadSafeConnection<M: Migrator = ()> {
pub struct ThreadSafeConnection<M: Migrator + 'static = ()> {
uri: Arc<str>,
persistent: bool,
connection_initialize_query: Option<&'static str>,
connections: Arc<ThreadLocal<Connection>>,
_migrator: PhantomData<M>,
_migrator: PhantomData<*mut M>,
}
unsafe impl<T: Migrator> Send for ThreadSafeConnection<T> {}
unsafe impl<T: Migrator> Sync for ThreadSafeConnection<T> {}
unsafe impl<M: Migrator> Send for ThreadSafeConnection<M> {}
unsafe impl<M: Migrator> Sync for ThreadSafeConnection<M> {}
pub struct ThreadSafeConnectionBuilder<M: Migrator = ()> {
pub struct ThreadSafeConnectionBuilder<M: Migrator + 'static = ()> {
db_initialize_query: Option<&'static str>,
write_queue_constructor: Option<WriteQueueConstructor>,
connection: ThreadSafeConnection<M>,
@ -54,6 +50,13 @@ impl<M: Migrator> ThreadSafeConnectionBuilder<M> {
self
}
/// Queues an initialization query for the database file. This must be infallible
/// but may cause changes to the database file such as with `PRAGMA journal_mode`
pub fn with_db_initialization_query(mut self, initialize_query: &'static str) -> Self {
self.db_initialize_query = Some(initialize_query);
self
}
/// Specifies how the thread safe connection should serialize writes. If provided
/// the connection will call the write_queue_constructor for each database file in
/// this process. The constructor is responsible for setting up a background thread or
@ -66,13 +69,6 @@ impl<M: Migrator> ThreadSafeConnectionBuilder<M> {
self
}
/// Queues an initialization query for the database file. This must be infallible
/// but may cause changes to the database file such as with `PRAGMA journal_mode`
pub fn with_db_initialization_query(mut self, initialize_query: &'static str) -> Self {
self.db_initialize_query = Some(initialize_query);
self
}
pub async fn build(self) -> anyhow::Result<ThreadSafeConnection<M>> {
self.connection
.initialize_queues(self.write_queue_constructor);
@ -100,6 +96,7 @@ impl<M: Migrator> ThreadSafeConnectionBuilder<M> {
.with_savepoint("thread_safe_multi_migration", || M::migrate(connection));
if migration_result.is_ok() {
println!("Migration succeded");
break;
}
}
@ -113,38 +110,17 @@ impl<M: Migrator> ThreadSafeConnectionBuilder<M> {
}
impl<M: Migrator> ThreadSafeConnection<M> {
fn initialize_queues(&self, write_queue_constructor: Option<WriteQueueConstructor>) {
fn initialize_queues(&self, write_queue_constructor: Option<WriteQueueConstructor>) -> bool {
if !QUEUES.read().contains_key(&self.uri) {
let mut queues = QUEUES.write();
if !queues.contains_key(&self.uri) {
let mut write_connection = self.create_connection();
// Enable writes for this connection
write_connection.write = true;
if let Some(mut write_queue_constructor) = write_queue_constructor {
let write_channel = write_queue_constructor(write_connection);
queues.insert(self.uri.clone(), write_channel);
} else {
use std::sync::mpsc::channel;
let (sender, reciever) = channel::<QueuedWrite>();
thread::spawn(move || {
while let Ok(write) = reciever.recv() {
write(&write_connection)
}
});
let sender = UnboundedSyncSender::new(sender);
queues.insert(
self.uri.clone(),
Box::new(move |queued_write| {
sender
.send(queued_write)
.expect("Could not send write action to backgorund thread");
}),
);
}
let mut write_queue_constructor =
write_queue_constructor.unwrap_or(background_thread_queue());
queues.insert(self.uri.clone(), write_queue_constructor());
return true;
}
}
return false;
}
pub fn builder(uri: &str, persistent: bool) -> ThreadSafeConnectionBuilder<M> {
@ -163,20 +139,21 @@ impl<M: Migrator> ThreadSafeConnection<M> {
/// Opens a new db connection with the initialized file path. This is internal and only
/// called from the deref function.
fn open_file(&self) -> Connection {
Connection::open_file(self.uri.as_ref())
fn open_file(uri: &str) -> Connection {
Connection::open_file(uri)
}
/// Opens a shared memory connection using the file path as the identifier. This is internal
/// and only called from the deref function.
fn open_shared_memory(&self) -> Connection {
Connection::open_memory(Some(self.uri.as_ref()))
fn open_shared_memory(uri: &str) -> Connection {
Connection::open_memory(Some(uri))
}
pub fn write<T: 'static + Send + Sync>(
&self,
callback: impl 'static + Send + FnOnce(&Connection) -> T,
) -> impl Future<Output = T> {
// Check and invalidate queue and maybe recreate queue
let queues = QUEUES.read();
let write_channel = queues
.get(&self.uri)
@ -185,24 +162,32 @@ impl<M: Migrator> ThreadSafeConnection<M> {
// Create a one shot channel for the result of the queued write
// so we can await on the result
let (sender, reciever) = oneshot::channel();
write_channel(Box::new(move |connection| {
sender.send(callback(connection)).ok();
let thread_safe_connection = (*self).clone();
write_channel(Box::new(move || {
let connection = thread_safe_connection.deref();
let result = connection.with_write(|connection| callback(connection));
sender.send(result).ok();
}));
reciever.map(|response| response.expect("Background writer thread unexpectedly closed"))
}
pub(crate) fn create_connection(&self) -> Connection {
let mut connection = if self.persistent {
self.open_file()
pub(crate) fn create_connection(
persistent: bool,
uri: &str,
connection_initialize_query: Option<&'static str>,
) -> Connection {
let mut connection = if persistent {
Self::open_file(uri)
} else {
self.open_shared_memory()
Self::open_shared_memory(uri)
};
// Disallow writes on the connection. The only writes allowed for thread safe connections
// are from the background thread that can serialize them.
connection.write = false;
*connection.write.get_mut() = false;
if let Some(initialize_query) = self.connection_initialize_query {
if let Some(initialize_query) = connection_initialize_query {
connection.exec(initialize_query).expect(&format!(
"Initialize query failed to execute: {}",
initialize_query
@ -236,7 +221,7 @@ impl ThreadSafeConnection<()> {
}
}
impl<D: Domain> Clone for ThreadSafeConnection<D> {
impl<M: Migrator> Clone for ThreadSafeConnection<M> {
fn clone(&self) -> Self {
Self {
uri: self.uri.clone(),
@ -252,16 +237,41 @@ impl<M: Migrator> Deref for ThreadSafeConnection<M> {
type Target = Connection;
fn deref(&self) -> &Self::Target {
self.connections.get_or(|| self.create_connection())
self.connections.get_or(|| {
Self::create_connection(self.persistent, &self.uri, self.connection_initialize_query)
})
}
}
pub fn locking_queue() -> WriteQueueConstructor {
Box::new(|connection| {
let connection = Mutex::new(connection);
pub fn background_thread_queue() -> WriteQueueConstructor {
use std::sync::mpsc::channel;
Box::new(|| {
let (sender, reciever) = channel::<QueuedWrite>();
thread::spawn(move || {
while let Ok(write) = reciever.recv() {
write()
}
});
let sender = UnboundedSyncSender::new(sender);
Box::new(move |queued_write| {
let connection = connection.lock();
queued_write(&connection)
sender
.send(queued_write)
.expect("Could not send write action to background thread");
})
})
}
pub fn locking_queue() -> WriteQueueConstructor {
Box::new(|| {
let mutex = Mutex::new(());
Box::new(move |queued_write| {
eprintln!("Write started");
let _ = mutex.lock();
queued_write();
eprintln!("Write finished");
})
})
}
@ -269,7 +279,8 @@ pub fn locking_queue() -> WriteQueueConstructor {
#[cfg(test)]
mod test {
use indoc::indoc;
use std::ops::Deref;
use lazy_static::__Deref;
use std::thread;
use crate::{domain::Domain, thread_safe_connection::ThreadSafeConnection};