diff --git a/Cargo.lock b/Cargo.lock index 04b194fbec9..e919d234b24 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -337,6 +337,7 @@ dependencies = [ "crossterm", "crypto_secretbox", "directories", + "divan", "eyre", "fs-err", "futures", @@ -369,6 +370,7 @@ dependencies = [ "sqlx", "strum", "strum_macros", + "tempfile", "testing_logger", "thiserror 2.0.18", "time", diff --git a/crates/atuin-client/Cargo.toml b/crates/atuin-client/Cargo.toml index 1b007f2e2d9..e33f96aff69 100644 --- a/crates/atuin-client/Cargo.toml +++ b/crates/atuin-client/Cargo.toml @@ -82,3 +82,9 @@ strum = { version = "0.27", features = ["strum_macros"] } tokio = { version = "1", features = ["full"] } pretty_assertions = { workspace = true } testing_logger = "0.1.1" +divan = "0.1.14" +tempfile = "3" + +[[bench]] +name = "record_store" +harness = false diff --git a/crates/atuin-client/benches/record/mod.rs b/crates/atuin-client/benches/record/mod.rs new file mode 100644 index 00000000000..ed91421c1f7 --- /dev/null +++ b/crates/atuin-client/benches/record/mod.rs @@ -0,0 +1 @@ +mod sqlite_store; diff --git a/crates/atuin-client/benches/record/sqlite_store.rs b/crates/atuin-client/benches/record/sqlite_store.rs new file mode 100644 index 00000000000..c5041ab6e05 --- /dev/null +++ b/crates/atuin-client/benches/record/sqlite_store.rs @@ -0,0 +1,102 @@ +use atuin_client::record::sqlite_store::SqliteStore; +use atuin_client::record::store::Store; +use atuin_common::record::{EncryptedData, Host, HostId, Record}; +use atuin_common::utils::uuid_v7; +use rand::{Rng, distributions::Alphanumeric}; +use tempfile::TempDir; + +struct BenchRecordBuilder; + +impl BenchRecordBuilder { + /// Controls how large the record payload is. Roughly, this is between 200 and 400 bytes for + /// a typical history record. + /// + /// Breakdown: + /// - id (UUID string, 36 bytes) + /// - timestamp (u64, 8 bytes) + /// - duration (i64, 8 bytes) + /// - exit code (i64, 8 bytes) + /// - command (variable — average shell command is ~20-50 bytes, but can be much longer) + /// - cwd (path string, ~20-60 bytes) + /// - session (string, ~36 bytes) + /// - hostname (string, ~10-30 bytes) + /// - deleted_at (optional u64) + /// - author (string) + const PAYLOAD_SIZE: usize = 300; + + /// Rough size of the PASETO PIE-wrapped key. + const KEY_SIZE: usize = 150; + + fn chain(n: usize) -> Vec> { + let host = Host::new(HostId(uuid_v7())); + let version: String = "v1".into(); + let tag = uuid_v7().simple().to_string(); + let data: String = rand::thread_rng() + .sample_iter(&Alphanumeric) + .take(Self::PAYLOAD_SIZE) + .map(char::from) + .collect(); + let key: String = rand::thread_rng() + .sample_iter(&Alphanumeric) + .take(Self::KEY_SIZE) + .map(char::from) + .collect(); + + (0..n as u64) + .map(|idx| { + Record::builder() + .host(host.clone()) + .version(version.clone()) + .tag(tag.clone()) + .data(EncryptedData { + data: data.clone(), + content_encryption_key: key.clone(), + }) + .idx(idx) + .build() + }) + .collect() + } +} + +struct BenchSqliteStore { + _temp_dir: TempDir, + sqlite_store: SqliteStore, +} + +impl BenchSqliteStore { + const SQL_TIMEOUT_S: f64 = 5.0; + + async fn new() -> Self { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("bench.db"); + let store = SqliteStore::new(db_path, Self::SQL_TIMEOUT_S) + .await + .unwrap(); + + Self { + _temp_dir: dir, + sqlite_store: store, + } + } +} + +/// Benchmark to exercise the latency of pushing a batch of varying records. +/// The parameters are: +/// - 1 proves out the case of adding one shell entry via `push_record` (history/store.rs). +/// - 100 is the page size used by `sync_remote` (record/sync.rs). +#[divan::bench(args = [1, 10, 100], sample_count = 500, min_time = 5)] +fn push_batch(bencher: divan::Bencher, n: usize) { + let rt = tokio::runtime::Runtime::new().unwrap(); + + bencher + .with_inputs(|| { + let db = rt.block_on(BenchSqliteStore::new()); + let records = BenchRecordBuilder::chain(n); + (db, records) + }) + .bench_values(|(db, records)| { + rt.block_on(db.sqlite_store.push_batch(records.iter())) + .unwrap(); + }); +} diff --git a/crates/atuin-client/benches/record_store.rs b/crates/atuin-client/benches/record_store.rs new file mode 100644 index 00000000000..4b502701d34 --- /dev/null +++ b/crates/atuin-client/benches/record_store.rs @@ -0,0 +1,5 @@ +mod record; + +fn main() { + divan::main(); +} diff --git a/crates/atuin-client/src/record/sqlite_store.rs b/crates/atuin-client/src/record/sqlite_store.rs index ed51f3fdc07..04220631081 100644 --- a/crates/atuin-client/src/record/sqlite_store.rs +++ b/crates/atuin-client/src/record/sqlite_store.rs @@ -11,7 +11,10 @@ use fs_err as fs; use sqlx::{ Row, - sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions, SqliteRow}, + sqlite::{ + SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions, SqliteRow, + SqliteSynchronous, + }, }; use atuin_common::record::{ @@ -49,6 +52,8 @@ impl SqliteStore { let opts = SqliteConnectOptions::from_str(path.as_os_str().to_str().unwrap())? .journal_mode(SqliteJournalMode::Wal) + .optimize_on_close(true, None) + .synchronous(SqliteSynchronous::Normal) .foreign_keys(true) .create_if_missing(true);