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.