Avoid opening the Auth environment multiple times

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,

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(),
         }
     }

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);

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.

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<Path>, dst: impl AsRef<Path>) -> 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<Path>) -> 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<Path>, dst: impl AsRef<Path>) -> Result<()> {
-        let store = HeedAuthStore::new(&dst)?;
+    pub fn load_dump(src: impl AsRef<Path>, auth_env: Env) -> Result<()> {
+        let store = HeedAuthStore::new(auth_env)?;
 
         let keys_file_path = src.as_ref().join(KEYS_PATH);
 

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<HeedAuthStore>,
+    store: HeedAuthStore,
     master_key: Option<String>,
 }
 
 impl AuthController {
-    pub fn new(db_path: impl AsRef<Path>, master_key: &Option<String>) -> Result<Self> {
-        let store = HeedAuthStore::new(db_path)?;
+    pub fn new(auth_env: Env, master_key: &Option<String>) -> Result<Self> {
+        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.

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<WithTls>>,
+    env: Env,
     keys: Database<Bytes, SerdeJson<Key>>,
     action_keyid_index_expiration: Database<KeyIdActionCodec, SerdeJson<Option<OffsetDateTime>>>,
-    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<milli::heed::Env<WithTls>> {
+pub fn open_auth_store_env(path: &Path) -> milli::heed::Result<milli::heed::Env> {
     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<milli::heed::Env<
 }
 
 impl HeedAuthStore {
-    pub fn new(path: impl AsRef<Path>) -> Result<Self> {
-        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<Self> {
         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<bool> {
         let rtxn = self.env.read_txn()?;
 

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(),

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;