From 437040f0501879eae91634b2e66d9dcfb6e93e3a Mon Sep 17 00:00:00 2001 From: Ellie Huxtable Date: Sat, 25 May 2024 10:10:17 +0100 Subject: [PATCH] perf: only open the database for commands if strictly required The client commands would open sqlite, even if not 100% required Remove this for 1. history start/end 2. shell init Init seems to be around 2ms faster on my system, with this change. --- crates/atuin/src/command/client.rs | 14 ++- crates/atuin/src/command/client/history.rs | 125 +++++++++++++-------- 2 files changed, 92 insertions(+), 47 deletions(-) diff --git a/crates/atuin/src/command/client.rs b/crates/atuin/src/command/client.rs index 0f8f3ba0d2f..180ced102c8 100644 --- a/crates/atuin/src/command/client.rs +++ b/crates/atuin/src/command/client.rs @@ -111,6 +111,15 @@ impl Cmd { tracing::trace!(command = ?self, "client command"); + // Skip initializing any databases for history + // This is a pretty hot path, as it runs before and after every single command the user + // runs + match self { + Self::History(history) => return history.run(&settings).await, + Self::Init(init) => return init.run(&settings).await, + _ => {} + } + let db_path = PathBuf::from(settings.db_path.as_str()); let record_store_path = PathBuf::from(settings.record_store_path.as_str()); @@ -118,7 +127,6 @@ impl Cmd { let sqlite_store = SqliteStore::new(record_store_path, settings.local_timeout).await?; match self { - Self::History(history) => history.run(&settings, &db, sqlite_store).await, Self::Import(import) => import.run(&db).await, Self::Stats(stats) => stats.run(&db, &settings).await, Self::Search(search) => search.run(db, &mut settings, sqlite_store).await, @@ -135,8 +143,6 @@ impl Cmd { Self::Dotfiles(dotfiles) => dotfiles.run(&settings, sqlite_store).await, - Self::Init(init) => init.run(&settings).await, - Self::Info => { info::run(&settings); Ok(()) @@ -151,6 +157,8 @@ impl Cmd { #[cfg(feature = "daemon")] Self::Daemon => daemon::run(settings, sqlite_store, db).await, + + _ => unimplemented!(), } } } diff --git a/crates/atuin/src/command/client/history.rs b/crates/atuin/src/command/client/history.rs index 9d1453fc3e5..a9970559217 100644 --- a/crates/atuin/src/command/client/history.rs +++ b/crates/atuin/src/command/client/history.rs @@ -1,6 +1,7 @@ use std::{ fmt::{self, Display}, io::{self, IsTerminal, Write}, + path::PathBuf, time::Duration, }; @@ -10,7 +11,7 @@ use eyre::{Context, Result}; use runtime_format::{FormatKey, FormatKeyError, ParseSegment, ParsedFmt}; use atuin_client::{ - database::{current_context, Database}, + database::{current_context, Database, Sqlite}, encryption, history::{store::HistoryStore, History}, record::sqlite_store::SqliteStore, @@ -314,28 +315,45 @@ impl Cmd { return Ok(()); } - if settings.daemon.enabled { - let resp = atuin_daemon::client::HistoryClient::new( - #[cfg(not(unix))] - settings.daemon.tcp_port, - #[cfg(unix)] - settings.daemon.socket_path.clone(), - ) - .await? - .start_history(h) - .await?; + // print the ID + // we use this as the key for calling end + println!("{}", h.id); + db.save(&h).await?; - // print the ID - // we use this as the key for calling end - println!("{resp}"); + Ok(()) + } + + async fn handle_daemon_start(settings: &Settings, command: &[String]) -> Result<()> { + let command = command.join(" "); + + // It's better for atuin to silently fail here and attempt to + // store whatever is ran, than to throw an error to the terminal + let cwd = utils::get_current_dir(); + let h: History = History::capture() + .timestamp(OffsetDateTime::now_utc()) + .command(command) + .cwd(cwd) + .build() + .into(); + + if !h.should_save(settings) { return Ok(()); } + let resp = atuin_daemon::client::HistoryClient::new( + #[cfg(not(unix))] + settings.daemon.tcp_port, + #[cfg(unix)] + settings.daemon.socket_path.clone(), + ) + .await? + .start_history(h) + .await?; + // print the ID // we use this as the key for calling end - println!("{}", h.id); - db.save(&h).await?; + println!("{resp}"); Ok(()) } @@ -350,23 +368,6 @@ impl Cmd { exit: i64, duration: Option, ) -> Result<()> { - // If the daemon is enabled, use it. Ignore the rest. - // We will need to keep the old code around for a while. - // At the very least, while this is opt-in - if settings.daemon.enabled { - let resp = atuin_daemon::client::HistoryClient::new( - #[cfg(not(unix))] - settings.daemon.tcp_port, - #[cfg(unix)] - settings.daemon.socket_path.clone(), - ) - .await? - .end_history(id.to_string(), duration.unwrap_or(0), exit) - .await?; - - return Ok(()); - } - if id.trim() == "" { return Ok(()); } @@ -424,6 +425,26 @@ impl Cmd { Ok(()) } + #[allow(unused_variables)] + async fn handle_daemon_end( + settings: &Settings, + id: &str, + exit: i64, + duration: Option, + ) -> Result<()> { + let resp = atuin_daemon::client::HistoryClient::new( + #[cfg(not(unix))] + settings.daemon.tcp_port, + #[cfg(unix)] + settings.daemon.socket_path.clone(), + ) + .await? + .end_history(id.to_string(), duration.unwrap_or(0), exit) + .await?; + + Ok(()) + } + #[allow(clippy::too_many_arguments)] #[allow(clippy::fn_params_excessive_bools)] async fn handle_list( @@ -519,14 +540,30 @@ impl Cmd { Ok(()) } - pub async fn run( - self, - settings: &Settings, - db: &impl Database, - store: SqliteStore, - ) -> Result<()> { + pub async fn run(self, settings: &Settings) -> Result<()> { let context = current_context(); + // Skip initializing any databases for start/end, if the daemon is enabled + if settings.daemon.enabled { + match self { + Self::Start { command } => { + return Self::handle_daemon_start(settings, &command).await + } + + Self::End { id, exit, duration } => { + return Self::handle_daemon_end(settings, &id, exit, duration).await + } + + _ => {} + } + } + + 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, settings.local_timeout).await?; + let store = SqliteStore::new(record_store_path, settings.local_timeout).await?; + let encryption_key: [u8; 32] = encryption::load_key(settings) .context("could not load encryption key")? .into(); @@ -535,9 +572,9 @@ impl Cmd { let history_store = HistoryStore::new(store.clone(), host_id, encryption_key); match self { - Self::Start { command } => Self::handle_start(db, settings, &command).await, + Self::Start { command } => Self::handle_start(&db, settings, &command).await, Self::End { id, exit, duration } => { - Self::handle_end(db, store, history_store, settings, &id, exit, duration).await + Self::handle_end(&db, store, history_store, settings, &id, exit, duration).await } Self::List { session, @@ -552,7 +589,7 @@ impl Cmd { let mode = ListMode::from_flags(human, cmd_only); let tz = timezone.unwrap_or(settings.timezone); Self::handle_list( - db, settings, context, session, cwd, mode, format, false, print0, reverse, tz, + &db, settings, context, session, cwd, mode, format, false, print0, reverse, tz, ) .await } @@ -581,10 +618,10 @@ impl Cmd { Ok(()) } - Self::InitStore => history_store.init_store(db).await, + Self::InitStore => history_store.init_store(&db).await, Self::Prune { dry_run } => { - Self::handle_prune(db, settings, store, context, dry_run).await + Self::handle_prune(&db, settings, store, context, dry_run).await } } }