Skip to content

Commit

Permalink
fix: add acquire timeout to sqlite database connection (#1590)
Browse files Browse the repository at this point in the history
* fix: add acquire timeout to sqlite database connection

This should fix #1503

I wasn't able to trigger enough IO pressure for the SQL connection to be
a problem.

This adds `local_timeout` to the client config. This is a float, and
represents the number of seconds (units in line with the other timeouts,
though those are ints). Users may well want to reduce this if they
regularly have issues, but by default I think 2s is fine and avoids a
non-responsive system in bad situations.

* tests
  • Loading branch information
ellie committed Jan 19, 2024
1 parent 10f465d commit 8899ce5
Show file tree
Hide file tree
Showing 7 changed files with 39 additions and 24 deletions.
6 changes: 6 additions & 0 deletions atuin-client/config.toml
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,12 @@ enter_accept = true
## the specified one.
# keymap_mode = "auto"

# network_connect_timeout = 5
# network_timeout = 5

## Timeout (in seconds) for acquiring a local database connection (sqlite)
# local_timeout = 5

#[stats]
# Set commands where we should consider the subcommand for statistics. Eg, kubectl get vs just kubectl
#common_subcommands = [
Expand Down
18 changes: 11 additions & 7 deletions atuin-client/src/database.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::{
env,
path::{Path, PathBuf},
str::FromStr,
time::Duration,
};

use async_trait::async_trait;
Expand Down Expand Up @@ -123,7 +124,7 @@ pub struct Sqlite {
}

impl Sqlite {
pub async fn new(path: impl AsRef<Path>) -> Result<Self> {
pub async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> {
let path = path.as_ref();
debug!("opening sqlite database at {:?}", path);

Expand All @@ -138,7 +139,10 @@ impl Sqlite {
.journal_mode(SqliteJournalMode::Wal)
.create_if_missing(true);

let pool = SqlitePoolOptions::new().connect_with(opts).await?;
let pool = SqlitePoolOptions::new()
.acquire_timeout(Duration::from_secs_f64(timeout))
.connect_with(opts)
.await?;

Self::setup_db(&pool).await?;

Expand Down Expand Up @@ -786,7 +790,7 @@ mod test {

#[tokio::test(flavor = "multi_thread")]
async fn test_search_prefix() {
let mut db = Sqlite::new("sqlite::memory:").await.unwrap();
let mut db = Sqlite::new("sqlite::memory:", 0.1).await.unwrap();
new_history_item(&mut db, "ls /home/ellie").await.unwrap();

assert_search_eq(&db, SearchMode::Prefix, FilterMode::Global, "ls", 1)
Expand All @@ -802,7 +806,7 @@ mod test {

#[tokio::test(flavor = "multi_thread")]
async fn test_search_fulltext() {
let mut db = Sqlite::new("sqlite::memory:").await.unwrap();
let mut db = Sqlite::new("sqlite::memory:", 0.1).await.unwrap();
new_history_item(&mut db, "ls /home/ellie").await.unwrap();

assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "ls", 1)
Expand All @@ -818,7 +822,7 @@ mod test {

#[tokio::test(flavor = "multi_thread")]
async fn test_search_fuzzy() {
let mut db = Sqlite::new("sqlite::memory:").await.unwrap();
let mut db = Sqlite::new("sqlite::memory:", 0.1).await.unwrap();
new_history_item(&mut db, "ls /home/ellie").await.unwrap();
new_history_item(&mut db, "ls /home/frank").await.unwrap();
new_history_item(&mut db, "cd /home/Ellie").await.unwrap();
Expand Down Expand Up @@ -908,7 +912,7 @@ mod test {

#[tokio::test(flavor = "multi_thread")]
async fn test_search_reordered_fuzzy() {
let mut db = Sqlite::new("sqlite::memory:").await.unwrap();
let mut db = Sqlite::new("sqlite::memory:", 0.1).await.unwrap();
// test ordering of results: we should choose the first, even though it happened longer ago.

new_history_item(&mut db, "curl").await.unwrap();
Expand Down Expand Up @@ -942,7 +946,7 @@ mod test {
git_root: None,
};

let mut db = Sqlite::new("sqlite::memory:").await.unwrap();
let mut db = Sqlite::new("sqlite::memory:", 0.1).await.unwrap();
for _i in 1..10000 {
new_history_item(&mut db, "i am a duplicated command")
.await
Expand Down
2 changes: 1 addition & 1 deletion atuin-client/src/kv.rs
Original file line number Diff line number Diff line change
Expand Up @@ -221,7 +221,7 @@ mod tests {

#[tokio::test]
async fn build_kv() {
let mut store = SqliteStore::new(":memory:").await.unwrap();
let mut store = SqliteStore::new(":memory:", 0.1).await.unwrap();
let kv = KvStore::new();
let key: [u8; 32] = XSalsa20Poly1305::generate_key(&mut OsRng).into();
let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7());
Expand Down
27 changes: 15 additions & 12 deletions atuin-client/src/record/sqlite_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@
// Multiple stores of multiple types are all stored in one chonky table (for now), and we just index
// by tag/host

use std::path::Path;
use std::str::FromStr;
use std::{path::Path, time::Duration};

use async_trait::async_trait;
use eyre::{eyre, Result};
Expand All @@ -27,7 +27,7 @@ pub struct SqliteStore {
}

impl SqliteStore {
pub async fn new(path: impl AsRef<Path>) -> Result<Self> {
pub async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> {
let path = path.as_ref();

debug!("opening sqlite database at {:?}", path);
Expand All @@ -44,7 +44,10 @@ impl SqliteStore {
.foreign_keys(true)
.create_if_missing(true);

let pool = SqlitePoolOptions::new().connect_with(opts).await?;
let pool = SqlitePoolOptions::new()
.acquire_timeout(Duration::from_secs_f64(timeout))
.connect_with(opts)
.await?;

Self::setup_db(&pool).await?;

Expand Down Expand Up @@ -261,7 +264,7 @@ mod tests {

#[tokio::test]
async fn create_db() {
let db = SqliteStore::new(":memory:").await;
let db = SqliteStore::new(":memory:", 0.1).await;

assert!(
db.is_ok(),
Expand All @@ -272,15 +275,15 @@ mod tests {

#[tokio::test]
async fn push_record() {
let db = SqliteStore::new(":memory:").await.unwrap();
let db = SqliteStore::new(":memory:", 0.1).await.unwrap();
let record = test_record();

db.push(&record).await.expect("failed to insert record");
}

#[tokio::test]
async fn get_record() {
let db = SqliteStore::new(":memory:").await.unwrap();
let db = SqliteStore::new(":memory:", 0.1).await.unwrap();
let record = test_record();
db.push(&record).await.unwrap();

Expand All @@ -291,7 +294,7 @@ mod tests {

#[tokio::test]
async fn last() {
let db = SqliteStore::new(":memory:").await.unwrap();
let db = SqliteStore::new(":memory:", 0.1).await.unwrap();
let record = test_record();
db.push(&record).await.unwrap();

Expand All @@ -309,7 +312,7 @@ mod tests {

#[tokio::test]
async fn first() {
let db = SqliteStore::new(":memory:").await.unwrap();
let db = SqliteStore::new(":memory:", 0.1).await.unwrap();
let record = test_record();
db.push(&record).await.unwrap();

Expand All @@ -327,7 +330,7 @@ mod tests {

#[tokio::test]
async fn len() {
let db = SqliteStore::new(":memory:").await.unwrap();
let db = SqliteStore::new(":memory:", 0.1).await.unwrap();
let record = test_record();
db.push(&record).await.unwrap();

Expand All @@ -341,7 +344,7 @@ mod tests {

#[tokio::test]
async fn len_different_tags() {
let db = SqliteStore::new(":memory:").await.unwrap();
let db = SqliteStore::new(":memory:", 0.1).await.unwrap();

// these have different tags, so the len should be the same
// we model multiple stores within one database
Expand All @@ -361,7 +364,7 @@ mod tests {

#[tokio::test]
async fn append_a_bunch() {
let db = SqliteStore::new(":memory:").await.unwrap();
let db = SqliteStore::new(":memory:", 0.1).await.unwrap();

let mut tail = test_record();
db.push(&tail).await.expect("failed to push record");
Expand All @@ -380,7 +383,7 @@ mod tests {

#[tokio::test]
async fn append_a_big_bunch() {
let db = SqliteStore::new(":memory:").await.unwrap();
let db = SqliteStore::new(":memory:", 0.1).await.unwrap();

let mut records: Vec<Record<EncryptedData>> = Vec::with_capacity(10000);

Expand Down
4 changes: 2 additions & 2 deletions atuin-client/src/record/sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -317,10 +317,10 @@ mod tests {
local_records: Vec<Record<EncryptedData>>,
remote_records: Vec<Record<EncryptedData>>,
) -> (SqliteStore, Vec<Diff>) {
let local_store = SqliteStore::new(":memory:")
let local_store = SqliteStore::new(":memory:", 0.1)
.await
.expect("failed to open in memory sqlite");
let remote_store = SqliteStore::new(":memory:")
let remote_store = SqliteStore::new(":memory:", 0.1)
.await
.expect("failed to open in memory sqlite"); // "remote"

Expand Down
2 changes: 2 additions & 0 deletions atuin-client/src/settings.rs
Original file line number Diff line number Diff line change
Expand Up @@ -244,6 +244,7 @@ pub struct Settings {

pub network_connect_timeout: u64,
pub network_timeout: u64,
pub local_timeout: f64,
pub enter_accept: bool,

#[serde(default)]
Expand Down Expand Up @@ -456,6 +457,7 @@ impl Settings {
.set_default("secrets_filter", true)?
.set_default("network_connect_timeout", 5)?
.set_default("network_timeout", 30)?
.set_default("local_timeout", 2.0)?
// enter_accept defaults to false here, but true in the default config file. The dissonance is
// intentional!
// Existing users will get the default "False", so we don't mess with any potential
Expand Down
4 changes: 2 additions & 2 deletions atuin/src/command/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,8 @@ impl Cmd {
let db_path = PathBuf::from(settings.db_path.as_str());
let record_store_path = PathBuf::from(settings.record_store_path.as_str());

let db = Sqlite::new(db_path).await?;
let sqlite_store = SqliteStore::new(record_store_path).await?;
let db = Sqlite::new(db_path, settings.local_timeout).await?;
let sqlite_store = SqliteStore::new(record_store_path, settings.local_timeout).await?;

match self {
Self::History(history) => history.run(&settings, &db, sqlite_store).await,
Expand Down

0 comments on commit 8899ce5

Please sign in to comment.