Skip to content

Instantly share code, notes, and snippets.

@xacrimon
Created September 2, 2022 15:42
Show Gist options
  • Save xacrimon/a1fac9bfdf7da6d06123730aea548582 to your computer and use it in GitHub Desktop.
Save xacrimon/a1fac9bfdf7da6d06123730aea548582 to your computer and use it in GitHub Desktop.
trait QueryMode {
type Handle<'conn>;
fn handle(connection: &mut rusqlite::Connection) -> rusqlite::Result<Self::Handle<'_>>;
}
struct NoTx;
impl QueryMode for NoTx {
type Handle<'conn> = &'conn rusqlite::Connection;
fn handle(connection: &mut rusqlite::Connection) -> rusqlite::Result<Self::Handle<'_>> {
Ok(connection)
}
}
struct Tx;
impl QueryMode for Tx {
type Handle<'conn> = rusqlite::Transaction<'conn>;
fn handle(connection: &mut rusqlite::Connection) -> rusqlite::Result<Self::Handle<'_>> {
connection.transaction()
}
}
#[instrument(skip(db, query), err)]
fn query_inner<M, F, T>(db: &Database, mut query: F) -> Result<T>
where
M: QueryMode,
F: FnMut(M::Handle<'_>) -> Result<T>,
{
task::block_in_place(|| {
let start = Instant::now();
let end = start + DB_TIMEOUT;
let conn = &mut db.acquire_conn();
let ret = loop {
let handle = M::handle(conn)?;
match query(handle) {
Ok(item) => break Ok(item),
Err(ref err) if let Some(sq_err) = err.downcast_ref::<rusqlite::Error>() => {
if let Some(code) = sq_err.sqlite_error_code() {
if code == rusqlite::ErrorCode::DatabaseBusy {
warn!("database is busy, retrying");
if Instant::now() > end {
bail!("database busy, timed out");
} else {
thread::yield_now();
continue;
}
}
}
}
Err(err) => break Err(err),
}
};
debug!("transaction took {:?}", start.elapsed());
ret
})
}
pub fn query<F, T>(db: &Database, query: F) -> Result<T>
where
F: FnMut(&rusqlite::Connection) -> Result<T>,
{
query_inner::<NoTx, _, _>(db, query)
}
pub fn query_tx<F, T>(db: &Database, query: F) -> Result<T>
where
F: FnMut(rusqlite::Transaction) -> Result<T>,
{
query_inner::<Tx, _, _>(db, query)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment