Skip to content

Instantly share code, notes, and snippets.

@markazmierczak
Created November 1, 2020 07:03
Show Gist options
  • Save markazmierczak/6ffad42239d9fe7d3bd6e9f0d698eeb9 to your computer and use it in GitHub Desktop.
Save markazmierczak/6ffad42239d9fe7d3bd6e9f0d698eeb9 to your computer and use it in GitHub Desktop.
Multiple writers with SQLite and diesel
[package]
name = "r2d2-locked"
version = "0.1.0"
edition = "2018"
[dependencies]
crossbeam = { version = "0.7.1" }
diesel = { version = "1.4.2", features = ["sqlite", "r2d2"] }
tokio = { version = "0.2", features = ["full"] }
// run with WAL env variable (zero or one):
//
// time WAL=1 target/release/r2d2-locked
#[macro_use]
extern crate diesel;
use diesel::connection::SimpleConnection;
use diesel::prelude::*;
use diesel::r2d2::{ConnectionManager, Pool};
use std::time::Duration;
table! {
users (name) {
name -> Text,
}
}
#[derive(Debug)]
pub struct ConnectionOptions {
pub enable_wal: bool,
pub enable_foreign_keys: bool,
pub busy_timeout: Option<Duration>,
}
impl diesel::r2d2::CustomizeConnection<SqliteConnection, diesel::r2d2::Error>
for ConnectionOptions
{
fn on_acquire(&self, conn: &mut SqliteConnection) -> Result<(), diesel::r2d2::Error> {
(|| {
if self.enable_wal {
conn.batch_execute("PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL;")?;
}
if self.enable_foreign_keys {
conn.batch_execute("PRAGMA foreign_keys = ON;")?;
}
if let Some(d) = self.busy_timeout {
conn.batch_execute(&format!("PRAGMA busy_timeout = {};", d.as_millis()))?;
}
Ok(())
})()
.map_err(diesel::r2d2::Error::QueryError)
}
}
fn main() {
let enable_wal = std::env::var("WAL").unwrap().parse::<i32>().unwrap() != 0;
let db_url = format!("/tmp/test-{}.db", if enable_wal { "wal" } else { "no-wal" });
let pool = Pool::builder()
.max_size(16)
.connection_customizer(Box::new(ConnectionOptions {
enable_wal,
enable_foreign_keys: true,
busy_timeout: Some(Duration::from_secs(30)),
}))
.build(ConnectionManager::<SqliteConnection>::new(db_url))
.unwrap();
let conn = pool.get().unwrap();
conn.batch_execute(
"CREATE TABLE IF NOT EXISTS users (
name TEXT PRIMARY KEY NOT NULL
);
DELETE FROM users;
",
)
.unwrap();
drop(conn);
crossbeam::scope(|scope| {
let mut guards = vec![];
let pool = &pool;
for i in 0..16 {
let guard = scope.spawn(move |_| {
for j in 0..10000 {
let conn = pool.get().unwrap();
drop::<Option<(String,)>>(users::table.first(&conn).optional().unwrap());
let name = format!("John{}a{}", i, j * 2);
diesel::insert_into(users::table)
.values(users::name.eq(name))
.execute(&conn)
.unwrap();
}
});
guards.push(guard);
}
for guard in guards {
guard.join().unwrap();
}
})
.unwrap();
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment