Pass tokio handle to index-scheduler

crates/index-scheduler/src/insta_snapshot.rs

@@ -36,6 +36,7 @@ pub fn snapshot_index_scheduler(scheduler: &IndexScheduler) -> String {
         run_loop_iteration: _,
         embedders: _,
         chat_settings: _,
+        runtime: _,
     } = scheduler;
 
     let rtxn = env.read_txn().unwrap();

crates/index-scheduler/src/lib.rs

@@ -216,6 +216,8 @@ pub struct IndexScheduler {
     /// A counter that is incremented before every call to [`tick`](IndexScheduler::tick)
     #[cfg(test)]
     run_loop_iteration: Arc<RwLock<usize>>,
+
+    runtime: Option<tokio::runtime::Handle>,
 }
 
 impl IndexScheduler {
@@ -242,6 +244,7 @@ impl IndexScheduler {
             run_loop_iteration: self.run_loop_iteration.clone(),
             features: self.features.clone(),
             chat_settings: self.chat_settings,
+            runtime: self.runtime.clone(),
         }
     }
 
@@ -260,6 +263,7 @@ impl IndexScheduler {
         options: IndexSchedulerOptions,
         auth_env: Env<WithoutTls>,
         from_db_version: (u32, u32, u32),
+        runtime: Option<tokio::runtime::Handle>,
         #[cfg(test)] test_breakpoint_sdr: crossbeam_channel::Sender<(test_utils::Breakpoint, bool)>,
         #[cfg(test)] planned_failures: Vec<(usize, test_utils::FailureLocation)>,
     ) -> Result<Self> {
@@ -341,6 +345,7 @@ impl IndexScheduler {
             run_loop_iteration: Arc::new(RwLock::new(0)),
             features,
             chat_settings,
+            runtime,
         };
 
         this.run();

crates/index-scheduler/src/test_utils.rs

@@ -126,7 +126,7 @@ impl IndexScheduler {
         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();
+            Self::new(options, auth_env, version, None, 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.