introduce config to have page lock in different phases

worldcoin · Jan 28, 2025 · 18c109f · 18c109f
1 parent 78e8598
commit 18c109f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/iris-mpc-common/src/config/mod.rs b/iris-mpc-common/src/config/mod.rs
@@ -116,7 +116,10 @@ pub struct Config {
     pub load_chunks_buffer_size: usize,
 
     #[serde(default)]
-    pub page_lock_after: bool,
+    pub page_lock_sleep: bool,
+
+    #[serde(default)]
+    pub page_lock_before: bool,
 }
 
 fn default_load_chunks_parallelism() -> usize {

diff --git a/iris-mpc/src/bin/server.rs b/iris-mpc/src/bin/server.rs
@@ -1094,10 +1094,10 @@ async fn server_main(config: Config) -> eyre::Result<()> {
 
                         let device_manager_clone = actor.device_manager.clone();
 
-                        let page_lock_after = config.page_lock_after;
+                        let page_lock_sleep = config.page_lock_sleep;
                         // prepare the handle for the rest of the page locks
                         let page_lock_handle = spawn_blocking(move || {
-                            if page_lock_after {
+                            if page_lock_sleep {
                                 tracing::info!("Sleeping before page lock");
                                 sleep(Duration::from_secs(120));
                             }
@@ -1129,6 +1129,11 @@ async fn server_main(config: Config) -> eyre::Result<()> {
                                 now.elapsed()
                             );
                         });
+                        let mut page_lock_handle = Some(page_lock_handle);
+
+                        if config.page_lock_before {
+                            page_lock_handle.take().unwrap().await?;
+                        }
 
                         let now = Instant::now();
                         let mut record_counter = 0;
@@ -1268,7 +1273,9 @@ async fn server_main(config: Config) -> eyre::Result<()> {
                         actor.preprocess_db();
 
                         tracing::info!("Waiting for all page-locks to finish");
-                        page_lock_handle.await.expect("Error while page-locking");
+                        if !config.page_lock_before {
+                            page_lock_handle.take().unwrap().await?;
+                        }
 
                         tracing::info!(
                             "Loaded {} records from db into memory in {:?} [DB sizes: {:?}]",