Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

perf: only open the database for commands if strictly required #2043

Merged
merged 1 commit into from
May 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 11 additions & 3 deletions crates/atuin/src/command/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,14 +111,22 @@ 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());

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,
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,
Expand All @@ -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(())
Expand All @@ -151,6 +157,8 @@ impl Cmd {

#[cfg(feature = "daemon")]
Self::Daemon => daemon::run(settings, sqlite_store, db).await,

_ => unimplemented!(),
}
}
}
125 changes: 81 additions & 44 deletions crates/atuin/src/command/client/history.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use std::{
fmt::{self, Display},
io::{self, IsTerminal, Write},
path::PathBuf,
time::Duration,
};

Expand All @@ -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,
Expand Down Expand Up @@ -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(())
}
Expand All @@ -350,23 +368,6 @@ impl Cmd {
exit: i64,
duration: Option<u64>,
) -> 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(());
}
Expand Down Expand Up @@ -424,6 +425,26 @@ impl Cmd {
Ok(())
}

#[allow(unused_variables)]
async fn handle_daemon_end(
settings: &Settings,
id: &str,
exit: i64,
duration: Option<u64>,
) -> 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(
Expand Down Expand Up @@ -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();
Expand All @@ -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,
Expand All @@ -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
}
Expand Down Expand Up @@ -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
}
}
}
Expand Down
Loading