Implement eager store backend.

Focuses on correctness over speed.

Author: Dragyx

src/store.rs

@@ -1,12 +1,15 @@
 #![allow(clippy::mem_forget)]
-
+//! Provides an interface for querying data from the nix store.
+//!
+//! - [LazyDBConnection] is a lazy connection the underlying sqlite database.
+mod db_common;
 mod db_eager;
 mod db_lazy;
 mod nix_command;
 mod queries;
 
 pub mod store {
-  pub use crate::store::db_lazy::DBConnection;
+  pub use crate::store::db_lazy::LazyDBConnection;
 }
 
 use std::{
@@ -26,8 +29,9 @@ use crate::{
   DerivationId,
   StorePath,
   store::{
+    db_eager::EagerDBConnection,
     nix_command::CommandBackend,
-    store::DBConnection,
+    store::LazyDBConnection,
   },
 };
 /// The normal database connection
@@ -83,6 +87,35 @@ impl<'a> CombinedStoreBackend<'a> {
   pub fn new(backends: Vec<Box<dyn StoreBackendPrintable<'a>>>) -> Self {
     Self { backends }
   }
+
+  /// Returns a backend that is focused on performance.
+  ///
+  /// The first choice is using direct sqlite queries that
+  /// return the rows lazily, but might skip rows should
+  /// a row conversion after the first row fail. Note that
+  /// this should be extremely unlikely / impossible since
+  /// the current row mappings perform only very basic conversion.
+  pub fn default_lazy() -> Self {
+    CombinedStoreBackend::new(vec![
+      Box::new(LazyDBConnection::new(DATABASE_PATH)),
+      Box::new(EagerDBConnection::new(DATABASE_PATH_IMMUTABLE)),
+      Box::new(CommandBackend),
+    ])
+  }
+
+  /// Returns a backend that is focused solely on absolutely guaranteeing
+  /// correct results at the cost of memory usage and database speed.
+  ///
+  /// Note that [DATABASE_PATH_IMMUTABLE] is not used here, since opening
+  /// the database can lead to undefined results (also silently with no errors)
+  /// if the database is actually modified while opened.
+  pub fn default_eager() -> Self {
+    CombinedStoreBackend::new(vec![
+      Box::new(EagerDBConnection::new(DATABASE_PATH)),
+      Box::new(CommandBackend),
+    ])
+  }
+
   // tries to execute a query until it succeeds or all connected backends have
   // been tried
   fn fallback_query<'b, F, Ret>(&'b self, query: F, path: &Path) -> Result<Ret>
@@ -121,11 +154,7 @@ impl<'a> CombinedStoreBackend<'a> {
 
 impl<'a> Default for CombinedStoreBackend<'a> {
   fn default() -> Self {
-    CombinedStoreBackend::new(vec![
-      Box::new(DBConnection::new(DATABASE_PATH)),
-      Box::new(DBConnection::new(DATABASE_PATH_IMMUTABLE)),
-      Box::new(CommandBackend),
-    ])
+    Self::default_lazy()
   }
 }
 

src/store/db_common.rs

@@ -0,0 +1,92 @@
+use std::path::Path;
+
+use anyhow::{
+  Context as _,
+  Result,
+  anyhow,
+};
+use rusqlite::{
+  Connection,
+  OpenFlags,
+};
+use size::Size;
+
+use crate::{
+  path_to_canonical_string,
+  store::queries,
+};
+
+pub(crate) fn default_sqlite_connection(path: &str) -> Result<Connection> {
+  let inner = rusqlite::Connection::open_with_flags(
+    path,
+    OpenFlags::SQLITE_OPEN_READ_ONLY // We only run queries, safeguard against corrupting the DB.
+      | OpenFlags::SQLITE_OPEN_NO_MUTEX // Part of the default flags, rusqlite takes care of locking anyways.
+      | OpenFlags::SQLITE_OPEN_URI,
+  )
+  .with_context(|| format!("failed to connect to Nix database at {}", path))?;
+
+  // Perform a batched query to set some settings using PRAGMA
+  // the main performance bottleneck when dix was run before
+  // was that the database file has to be brought from disk into
+  // memory.
+  //
+  // We read a large part of the DB anyways in each query,
+  // so it makes sense to set aside a large region of memory-mapped
+  // I/O prevent incurring page faults which can be done using
+  // `mmap_size`.
+  //
+  // This made a performance difference of about 500ms (but only
+  // when it was first run for a long time!).
+  //
+  // The file pages of the store can be evicted from main memory
+  // using:
+  //
+  // ```bash
+  // dd of=/nix/var/nix/db/db.sqlite oflag=nocache conv=notrunc,fdatasync count=0
+  // ```
+  //
+  // If you want to test this. Source: <https://unix.stackexchange.com/questions/36907/drop-a-specific-file-from-the-linux-filesystem-cache>.
+  //
+  // Documentation about the settings can be found here: <https://www.sqlite.org/pragma.html>
+  //
+  // [0]: 256MB, enough to fit the whole DB (at least on my system - Dragyx).
+  // [1]: Always store temporary tables in memory.
+  inner
+    .execute_batch(
+      "
+        PRAGMA mmap_size=268435456; -- See [0].
+        PRAGMA temp_store=2; -- See [1].
+        PRAGMA query_only;
+      ",
+    )
+    .with_context(|| format!("failed to cache Nix database at {}", path))?;
+  Ok(inner)
+}
+
+// FIXME: why is this marked as dead code? It is used by both the lazy
+// and eager backend implementation
+pub(crate) fn default_close_inner_connection(
+  path: &str,
+  maybe_conn: &mut Option<Connection>,
+) -> Result<()> {
+  let conn = maybe_conn.take().ok_or_else(|| {
+    anyhow!("Tried to close connection to {} that does not exist", path)
+  })?;
+  conn.close().map_err(|(conn_old, err)| {
+    *maybe_conn = Some(conn_old);
+    anyhow::Error::from(err).context("failed to close Nix database")
+  })
+}
+
+pub(crate) fn query_closure_size(
+  conn: &Connection,
+  path: &Path,
+) -> Result<Size> {
+  let path = path_to_canonical_string(path)?;
+
+  let closure_size = conn
+    .prepare_cached(queries::QUERY_CLOSURE_SIZE)?
+    .query_row([path], |row| Ok(Size::from_bytes(row.get::<_, i64>(0)?)))?;
+
+  Ok(closure_size)
+}

src/store/db_eager.rs

@@ -0,0 +1,132 @@
+use std::{
+  fmt::{
+    self,
+    Display,
+  },
+  path::Path,
+};
+
+use anyhow::{
+  Result,
+  anyhow,
+};
+use rusqlite::Row;
+
+use crate::{
+  DerivationId,
+  StorePath,
+  path_to_canonical_string,
+  store::{
+    StoreBackend,
+    db_common::{
+      self,
+    },
+    queries,
+  },
+};
+
+#[derive(Debug)]
+pub struct EagerDBConnection<'a> {
+  path: &'a str,
+  conn: Option<rusqlite::Connection>,
+}
+
+impl Display for EagerDBConnection<'_> {
+  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    write!(f, "DBConnection({})", self.path)
+  }
+}
+
+impl<'a> EagerDBConnection<'a> {
+  /// Create a new connection.
+  pub fn new(path: &'a str) -> EagerDBConnection<'a> {
+    EagerDBConnection { path, conn: None }
+  }
+  /// returns a reference to the inner connection
+  ///
+  /// raises an error if the connection has not been established
+  fn get_inner(&self) -> Result<&rusqlite::Connection> {
+    self
+      .conn
+      .as_ref()
+      .ok_or_else(|| anyhow!("Attempted to use database before connecting."))
+  }
+
+  /// Executes a query that returns multiple rows and returns
+  /// an iterator over them where the `map` is used to map
+  /// the rows to `T`.
+  ///
+  /// Note that this function collects all rows before returning
+  /// and raises the first error that is encountered (if any exist).
+  pub(crate) fn execute_row_query_with_path<T, M>(
+    &self,
+    query: &str,
+    path: &Path,
+    map: M,
+  ) -> Result<Box<dyn Iterator<Item = T> + '_>>
+  where
+    T: 'static,
+    M: Fn(&Row) -> rusqlite::Result<T> + 'static,
+  {
+    let path = path_to_canonical_string(path)?;
+    let mut results = Vec::new();
+    let mut query = self.get_inner()?.prepare_cached(query)?;
+    let queried_rows = query.query_map([path], map)?;
+    for row in queried_rows {
+      results.push(row?);
+    }
+    Ok(Box::new(results.into_iter()))
+  }
+}
+
+impl<'a> StoreBackend<'a> for EagerDBConnection<'_> {
+  fn connect(&mut self) -> Result<()> {
+    self.conn = Some(db_common::default_sqlite_connection(self.path)?);
+    Ok(())
+  }
+
+  fn connected(&self) -> bool {
+    self.conn.is_some()
+  }
+
+  fn close(&mut self) -> Result<()> {
+    db_common::default_close_inner_connection(self.path, &mut self.conn)
+  }
+
+  fn query_closure_size(&self, path: &std::path::Path) -> Result<size::Size> {
+    db_common::query_closure_size(self.get_inner()?, path)
+  }
+
+  fn query_system_derivations(
+    &self,
+    system: &std::path::Path,
+  ) -> Result<Box<dyn Iterator<Item = crate::StorePath> + '_>> {
+    self.execute_row_query_with_path(
+      queries::QUERY_SYSTEM_DERIVATIONS,
+      system,
+      |row| Ok(StorePath(row.get::<_, String>(0)?.into())),
+    )
+  }
+
+  fn query_dependents(
+    &self,
+    path: &std::path::Path,
+  ) -> Result<Box<dyn Iterator<Item = crate::StorePath> + '_>> {
+    self.execute_row_query_with_path(queries::QUERY_DEPENDENTS, path, |row| {
+      Ok(StorePath(row.get::<_, String>(0)?.into()))
+    })
+  }
+
+  fn query_dependency_graph(
+    &self,
+    path: &std::path::Path,
+  ) -> Result<
+    Box<dyn Iterator<Item = (crate::DerivationId, crate::DerivationId)> + '_>,
+  > {
+    self.execute_row_query_with_path(
+      queries::QUERY_DEPENDENCY_GRAPH,
+      path,
+      |row| Ok((DerivationId(row.get(0)?), DerivationId(row.get(1)?))),
+    )
+  }
+}