db.rs

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)
}