diff --git a/crates/index-scheduler/src/lib.rs b/crates/index-scheduler/src/lib.rs index 4605be6eb..636615858 100644 --- a/crates/index-scheduler/src/lib.rs +++ b/crates/index-scheduler/src/lib.rs @@ -209,6 +209,7 @@ impl IndexScheduler { #[allow(private_interfaces)] // because test_utils is private pub fn new( options: IndexSchedulerOptions, + auth_env: Env, from_db_version: (u32, u32, u32), #[cfg(test)] test_breakpoint_sdr: crossbeam_channel::Sender<(test_utils::Breakpoint, bool)>, #[cfg(test)] planned_failures: Vec<(usize, test_utils::FailureLocation)>, @@ -262,7 +263,7 @@ impl IndexScheduler { processing_tasks: Arc::new(RwLock::new(ProcessingTasks::new())), version, queue, - scheduler: Scheduler::new(&options), + scheduler: Scheduler::new(&options, auth_env), index_mapper, env, diff --git a/crates/index-scheduler/src/scheduler/mod.rs b/crates/index-scheduler/src/scheduler/mod.rs index 41ff7f809..1237b22c4 100644 --- a/crates/index-scheduler/src/scheduler/mod.rs +++ b/crates/index-scheduler/src/scheduler/mod.rs @@ -21,6 +21,7 @@ use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use std::sync::Arc; use meilisearch_types::error::ResponseError; +use meilisearch_types::heed::Env; use meilisearch_types::milli; use meilisearch_types::tasks::Status; use rayon::current_num_threads; @@ -71,7 +72,7 @@ pub struct Scheduler { pub(crate) snapshots_path: PathBuf, /// The path to the folder containing the auth LMDB env. - pub(crate) auth_path: PathBuf, + pub(crate) auth_env: Env, /// The path to the version file of Meilisearch. pub(crate) version_file_path: PathBuf, @@ -87,12 +88,12 @@ impl Scheduler { batched_tasks_size_limit: self.batched_tasks_size_limit, dumps_path: self.dumps_path.clone(), snapshots_path: self.snapshots_path.clone(), - auth_path: self.auth_path.clone(), + auth_env: self.auth_env.clone(), version_file_path: self.version_file_path.clone(), } } - pub fn new(options: &IndexSchedulerOptions) -> Scheduler { + pub fn new(options: &IndexSchedulerOptions, auth_env: Env) -> Scheduler { Scheduler { must_stop_processing: MustStopProcessing::default(), // we want to start the loop right away in case meilisearch was ctrl+Ced while processing things @@ -102,7 +103,7 @@ impl Scheduler { batched_tasks_size_limit: options.batched_tasks_size_limit, dumps_path: options.dumps_path.clone(), snapshots_path: options.snapshots_path.clone(), - auth_path: options.auth_path.clone(), + auth_env, version_file_path: options.version_file_path.clone(), } } diff --git a/crates/index-scheduler/src/scheduler/process_snapshot_creation.rs b/crates/index-scheduler/src/scheduler/process_snapshot_creation.rs index 85522c160..f1bafed01 100644 --- a/crates/index-scheduler/src/scheduler/process_snapshot_creation.rs +++ b/crates/index-scheduler/src/scheduler/process_snapshot_creation.rs @@ -4,7 +4,6 @@ use std::sync::atomic::Ordering; use meilisearch_types::heed::CompactionOption; use meilisearch_types::milli::progress::{Progress, VariableNameStep}; -use meilisearch_types::milli::{self}; use meilisearch_types::tasks::{Status, Task}; use meilisearch_types::{compression, VERSION_FILE_NAME}; @@ -91,14 +90,7 @@ impl IndexScheduler { progress.update_progress(SnapshotCreationProgress::SnapshotTheApiKeys); let dst = temp_snapshot_dir.path().join("auth"); fs::create_dir_all(&dst)?; - // TODO We can't use the open_auth_store_env function here but we should - let auth = unsafe { - milli::heed::EnvOpenOptions::new() - .map_size(1024 * 1024 * 1024) // 1 GiB - .max_dbs(2) - .open(&self.scheduler.auth_path) - }?; - auth.copy_to_path(dst.join("data.mdb"), CompactionOption::Enabled)?; + self.scheduler.auth_env.copy_to_path(dst.join("data.mdb"), CompactionOption::Enabled)?; // 5. Copy and tarball the flat snapshot progress.update_progress(SnapshotCreationProgress::CreateTheTarball); diff --git a/crates/index-scheduler/src/test_utils.rs b/crates/index-scheduler/src/test_utils.rs index 072220b6c..3efcc523a 100644 --- a/crates/index-scheduler/src/test_utils.rs +++ b/crates/index-scheduler/src/test_utils.rs @@ -5,6 +5,7 @@ use std::time::Duration; use big_s::S; use crossbeam_channel::RecvTimeoutError; use file_store::File; +use meilisearch_auth::open_auth_store_env; use meilisearch_types::document_formats::DocumentFormatError; use meilisearch_types::milli::update::IndexDocumentsMethod::ReplaceDocuments; use meilisearch_types::milli::update::IndexerConfig; @@ -120,7 +121,10 @@ impl IndexScheduler { ) }); - let index_scheduler = Self::new(options, version, sender, planned_failures).unwrap(); + std::fs::create_dir_all(&options.auth_path).unwrap(); + let auth_env = open_auth_store_env(&options.auth_path).unwrap(); + let index_scheduler = + Self::new(options, auth_env, version, sender, planned_failures).unwrap(); // To be 100% consistent between all test we're going to start the scheduler right now // and ensure it's in the expected starting state. diff --git a/crates/meilisearch-auth/src/dump.rs b/crates/meilisearch-auth/src/dump.rs index 8a215d8ae..d73b5cf09 100644 --- a/crates/meilisearch-auth/src/dump.rs +++ b/crates/meilisearch-auth/src/dump.rs @@ -2,6 +2,7 @@ use std::fs::File; use std::io::{BufReader, Write}; use std::path::Path; +use meilisearch_types::heed::Env; use serde_json::Deserializer; use crate::{AuthController, HeedAuthStore, Result}; @@ -9,11 +10,8 @@ use crate::{AuthController, HeedAuthStore, Result}; const KEYS_PATH: &str = "keys"; impl AuthController { - pub fn dump(src: impl AsRef, dst: impl AsRef) -> Result<()> { - let mut store = HeedAuthStore::new(&src)?; - - // do not attempt to close the database on drop! - store.set_drop_on_close(false); + pub fn dump(auth_env: Env, dst: impl AsRef) -> Result<()> { + let store = HeedAuthStore::new(auth_env)?; let keys_file_path = dst.as_ref().join(KEYS_PATH); @@ -27,8 +25,8 @@ impl AuthController { Ok(()) } - pub fn load_dump(src: impl AsRef, dst: impl AsRef) -> Result<()> { - let store = HeedAuthStore::new(&dst)?; + pub fn load_dump(src: impl AsRef, auth_env: Env) -> Result<()> { + let store = HeedAuthStore::new(auth_env)?; let keys_file_path = src.as_ref().join(KEYS_PATH); diff --git a/crates/meilisearch-auth/src/lib.rs b/crates/meilisearch-auth/src/lib.rs index 4dbf1bf6f..8e2445703 100644 --- a/crates/meilisearch-auth/src/lib.rs +++ b/crates/meilisearch-auth/src/lib.rs @@ -3,11 +3,10 @@ pub mod error; mod store; use std::collections::{HashMap, HashSet}; -use std::path::Path; -use std::sync::Arc; use error::{AuthControllerError, Result}; use maplit::hashset; +use meilisearch_types::heed::Env; use meilisearch_types::index_uid_pattern::IndexUidPattern; use meilisearch_types::keys::{Action, CreateApiKey, Key, PatchApiKey}; use meilisearch_types::milli::update::Setting; @@ -19,19 +18,19 @@ use uuid::Uuid; #[derive(Clone)] pub struct AuthController { - store: Arc, + store: HeedAuthStore, master_key: Option, } impl AuthController { - pub fn new(db_path: impl AsRef, master_key: &Option) -> Result { - let store = HeedAuthStore::new(db_path)?; + pub fn new(auth_env: Env, master_key: &Option) -> Result { + let store = HeedAuthStore::new(auth_env)?; if store.is_empty()? { generate_default_keys(&store)?; } - Ok(Self { store: Arc::new(store), master_key: master_key.clone() }) + Ok(Self { store, master_key: master_key.clone() }) } /// Return `Ok(())` if the auth controller is able to access one of its database. diff --git a/crates/meilisearch-auth/src/store.rs b/crates/meilisearch-auth/src/store.rs index 67059f854..6fcd43cb6 100644 --- a/crates/meilisearch-auth/src/store.rs +++ b/crates/meilisearch-auth/src/store.rs @@ -1,15 +1,13 @@ use std::borrow::Cow; use std::cmp::Reverse; use std::collections::HashSet; -use std::fs::create_dir_all; use std::path::Path; use std::result::Result as StdResult; use std::str; use std::str::FromStr; -use std::sync::Arc; use hmac::{Hmac, Mac}; -use meilisearch_types::heed::{BoxedError, WithTls}; +use meilisearch_types::heed::BoxedError; use meilisearch_types::index_uid_pattern::IndexUidPattern; use meilisearch_types::keys::KeyId; use meilisearch_types::milli; @@ -25,27 +23,17 @@ use super::error::{AuthControllerError, Result}; use super::{Action, Key}; const AUTH_STORE_SIZE: usize = 1_073_741_824; //1GiB -const AUTH_DB_PATH: &str = "auth"; const KEY_DB_NAME: &str = "api-keys"; const KEY_ID_ACTION_INDEX_EXPIRATION_DB_NAME: &str = "keyid-action-index-expiration"; #[derive(Clone)] pub struct HeedAuthStore { - env: Arc>, + env: Env, keys: Database>, action_keyid_index_expiration: Database>>, - should_close_on_drop: bool, } -impl Drop for HeedAuthStore { - fn drop(&mut self) { - if self.should_close_on_drop && Arc::strong_count(&self.env) == 1 { - self.env.as_ref().clone().prepare_for_closing(); - } - } -} - -pub fn open_auth_store_env(path: &Path) -> milli::heed::Result> { +pub fn open_auth_store_env(path: &Path) -> milli::heed::Result { let mut options = EnvOpenOptions::new(); options.map_size(AUTH_STORE_SIZE); // 1GB options.max_dbs(2); @@ -53,16 +41,13 @@ pub fn open_auth_store_env(path: &Path) -> milli::heed::Result) -> Result { - let path = path.as_ref().join(AUTH_DB_PATH); - create_dir_all(&path)?; - let env = Arc::new(open_auth_store_env(path.as_ref())?); + pub fn new(env: Env) -> Result { let mut wtxn = env.write_txn()?; let keys = env.create_database(&mut wtxn, Some(KEY_DB_NAME))?; let action_keyid_index_expiration = env.create_database(&mut wtxn, Some(KEY_ID_ACTION_INDEX_EXPIRATION_DB_NAME))?; wtxn.commit()?; - Ok(Self { env, keys, action_keyid_index_expiration, should_close_on_drop: true }) + Ok(Self { env, keys, action_keyid_index_expiration }) } /// Return `Ok(())` if the auth store is able to access one of its database. @@ -82,10 +67,6 @@ impl HeedAuthStore { Ok(self.env.non_free_pages_size()?) } - pub fn set_drop_on_close(&mut self, v: bool) { - self.should_close_on_drop = v; - } - pub fn is_empty(&self) -> Result { let rtxn = self.env.read_txn()?; diff --git a/crates/meilisearch/src/lib.rs b/crates/meilisearch/src/lib.rs index 4f4a9c1e0..1841d5556 100644 --- a/crates/meilisearch/src/lib.rs +++ b/crates/meilisearch/src/lib.rs @@ -34,7 +34,7 @@ use error::PayloadError; use extractors::payload::PayloadConfig; use index_scheduler::versioning::Versioning; use index_scheduler::{IndexScheduler, IndexSchedulerOptions}; -use meilisearch_auth::AuthController; +use meilisearch_auth::{open_auth_store_env, AuthController}; use meilisearch_types::milli::constants::VERSION_MAJOR; use meilisearch_types::milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader}; use meilisearch_types::milli::update::{IndexDocumentsConfig, IndexDocumentsMethod}; @@ -335,9 +335,12 @@ fn open_or_create_database_unchecked( ) -> anyhow::Result<(IndexScheduler, AuthController)> { // we don't want to create anything in the data.ms yet, thus we // wrap our two builders in a closure that'll be executed later. - let auth_controller = AuthController::new(&opt.db_path, &opt.master_key); - let index_scheduler_builder = - || -> anyhow::Result<_> { Ok(IndexScheduler::new(index_scheduler_opt, version)?) }; + std::fs::create_dir_all(&index_scheduler_opt.auth_path)?; + let auth_env = open_auth_store_env(&index_scheduler_opt.auth_path).unwrap(); + let auth_controller = AuthController::new(auth_env.clone(), &opt.master_key); + let index_scheduler_builder = || -> anyhow::Result<_> { + Ok(IndexScheduler::new(index_scheduler_opt, auth_env, version)?) + }; match ( index_scheduler_builder(), diff --git a/crates/meilitool/src/main.rs b/crates/meilitool/src/main.rs index c32dadd5c..1da846ac3 100644 --- a/crates/meilitool/src/main.rs +++ b/crates/meilitool/src/main.rs @@ -7,7 +7,7 @@ use anyhow::{bail, Context}; use clap::{Parser, Subcommand, ValueEnum}; use dump::{DumpWriter, IndexMetadata}; use file_store::FileStore; -use meilisearch_auth::AuthController; +use meilisearch_auth::{open_auth_store_env, AuthController}; use meilisearch_types::batches::Batch; use meilisearch_types::heed::types::{Bytes, SerdeJson, Str}; use meilisearch_types::heed::{ @@ -290,7 +290,10 @@ fn export_a_dump( eprintln!("Dumping the keys..."); // 2. dump the keys - let auth_store = AuthController::new(&db_path, &None) + let auth_path = db_path.join("auth"); + std::fs::create_dir_all(&auth_path).context("While creating the auth directory")?; + let auth_env = open_auth_store_env(&auth_path).context("While opening the auth store")?; + let auth_store = AuthController::new(auth_env, &None) .with_context(|| format!("While opening the auth store at {}", db_path.display()))?; let mut dump_keys = dump.create_keys()?; let mut count = 0;