refactor: mmap query and preprocessing PPs (#237)

* feat: mmap query and preprocessing param files

* fix: inverted logic

Author: hackaugusto

Cargo.lock

@@ -25,6 +25,7 @@ verifiable-db = { workspace = true }
 
 lgn-messages = { path = "../lgn-messages" }
 exponential-backoff = "2.0.1"
+memmap2 = "0.9.7"
 
 [features]
 dummy-prover = []

lgn-provers/Cargo.toml

@@ -1,18 +1,15 @@
 use std::collections::HashMap;
-use std::path::Path;
 use std::path::PathBuf;
 use std::str::FromStr;
 use std::time::Duration;
 
 use anyhow::bail;
 use anyhow::ensure;
 use anyhow::Context;
-use bytes::Bytes;
 use futures::StreamExt;
 use reqwest::StatusCode;
 use tokio::fs::File;
 use tokio::io::AsyncWriteExt;
-use tracing::error;
 use tracing::info;
 use tracing::warn;
 
@@ -39,7 +36,6 @@ const DOWNLOAD_BACKOFF_MIN_MILLIS: u64 = 100;
 const DOWNLOAD_BACKOFF_MAX_MILLIS: u64 = 10_000;
 
 /// Ratio to convert byte to megabytes.
-const BYTES_TO_MEGABYTES_USIZE: usize = 1024 * 1024;
 const BYTES_TO_MEGABYTES_U64: u64 = 1024 * 1024;
 
 /// Download and verify `file_name`.
@@ -55,7 +51,7 @@ pub async fn download_and_checksum(
     param_dir: &str,
     file_name: &str,
     checksums: &HashMap<String, blake3::Hash>,
-) -> anyhow::Result<Bytes> {
+) -> anyhow::Result<PathBuf> {
     let mut filepath = PathBuf::from(param_dir);
     filepath.push(file_name);
 
@@ -84,16 +80,14 @@ pub async fn download_and_checksum(
         .with_context(|| format!("Failed to create file. filepath: {}", filepath.display()))?;
 
     let mut hasher = blake3::Hasher::new();
-    let mut buf = std::fs::read(&filepath)
-        .with_context(|| format!("Reading file failed. filepath: {}", filepath.display()))?;
-    hasher.update_rayon(&buf);
+    hasher.update_mmap_rayon(&filepath)?;
 
     if *expected_checksum == hasher.finalize() {
         info!(
             "Found file matching checksum, skipping download. filepath: {}",
             filepath.display()
         );
-        return Ok(Bytes::from(buf));
+        return Ok(filepath);
     }
 
     let fileurl = format!("{base_url}/{file_name}");
@@ -111,10 +105,7 @@ pub async fn download_and_checksum(
 
     for (retry, duration) in backoff.iter().enumerate() {
         let result = resume_download(
-            &mut buf,
-            base_url,
             &mut file,
-            &filepath,
             &fileurl,
             retry,
             &client,
@@ -125,12 +116,8 @@ pub async fn download_and_checksum(
 
         match result {
             Ok(()) => {
-                info!(
-                    "Params loaded. filepath: {} size: {}MiB",
-                    filepath.display(),
-                    buf.len() / BYTES_TO_MEGABYTES_USIZE,
-                );
-                return Ok(Bytes::from(buf));
+                info!("Downloaded file. filepath: {}", filepath.display());
+                return Ok(filepath);
             },
             Err(err) => {
                 if let Some(duration) = duration {
@@ -154,10 +141,7 @@ pub async fn download_and_checksum(
 
 #[allow(clippy::too_many_arguments)]
 async fn resume_download(
-    buf: &mut Vec<u8>,
-    base_url: &str,
     file: &mut File,
-    filepath: &Path,
     fileurl: &str,
     retry: usize,
     client: &reqwest::Client,
@@ -182,66 +166,44 @@ async fn resume_download(
                 Ok(parsed)
             })?;
 
+    if response.status() == StatusCode::RANGE_NOT_SATISFIABLE {
+        hasher.reset();
+        file.set_len(0).await?;
+        bail!("Local file is bigger than remote, reset file and restart download");
+    }
+
+    ensure!(
+        response.status().is_success(),
+        "Requesting params failed. status: {}",
+        response.status(),
+    );
+
     info!(
-        "Downloading params. base_url: {} filepath: {} present: {}MiB download: {}MiB retry: {}",
-        base_url,
-        filepath.display(),
+        "Downloading params. fileurl: {} present: {}MiB download: {}MiB retry: {}",
+        fileurl,
         metadata.len() / BYTES_TO_MEGABYTES_U64,
         length / BYTES_TO_MEGABYTES_U64,
         retry,
     );
 
-    if response.status() == StatusCode::RANGE_NOT_SATISFIABLE {
-        warn!(
-            "Local file is bigger than remote, resetting length and checking checksum. filepath: {}",
-            filepath.display(),
-        );
-
-        hasher.reset();
-        buf.resize(
-            length.try_into().expect("File size should fit in a usize"),
-            0,
-        );
-        file.set_len(length).await?;
-        hasher.update_rayon(buf);
-    } else {
-        ensure!(
-            response.status().is_success(),
-            "Requesting params failed. status: {} filepath: {}",
-            response.status(),
-            filepath.display(),
-        );
-
-        let mut stream = response.bytes_stream();
-        while let Some(data) = stream.next().await {
-            let data = data?;
-            file.write_all(&data).await?;
-            hasher.update_rayon(&data);
-            buf.extend(data);
-        }
+    let mut stream = response.bytes_stream();
+    while let Some(data) = stream.next().await {
+        let data = data?;
+        file.write_all(&data).await?;
+        hasher.update_rayon(&data);
     }
 
     let found_checksum = hasher.finalize();
     if found_checksum != *expected_checksum {
-        error!(
-            "Checksum failed, restarting download. checksum: {} expected: {} filepath: {}",
-            found_checksum.to_hex(),
-            expected_checksum.to_hex(),
-            filepath.display(),
-        );
-
         hasher.reset();
-        buf.clear();
         file.set_len(0).await?;
 
         bail!(
-            "Checksum failed, restarting download. checksum: {} expected: {} filepath: {}",
+            "Checksum failed, restarting download. checksum: {} expected: {}",
             found_checksum.to_hex(),
             expected_checksum.to_hex(),
-            filepath.display(),
         );
     } else {
-        info!("Downloaded file. filepath: {}", filepath.display());
         Ok(())
     }
 }

lgn-provers/src/params/mod.rs

@@ -1,5 +1,7 @@
 use std::collections::HashMap;
+use std::fs::read;
 
+use anyhow::Context;
 use anyhow::Result;
 use groth16_framework::Groth16Prover;
 use tracing::debug;
@@ -21,17 +23,20 @@ impl Groth16EuclidProver {
         pk_file: &str,
         checksums: &HashMap<String, blake3::Hash>,
     ) -> Result<Self> {
-        let circuit_bytes =
+        let circuit_bytes_path =
             params::download_and_checksum(url, dir, circuit_file, checksums).await?;
-        let r1cs_bytes = params::download_and_checksum(url, dir, r1cs_file, checksums).await?;
-        let pk_bytes = params::download_and_checksum(url, dir, pk_file, checksums).await?;
+        let r1cs_bytes_path = params::download_and_checksum(url, dir, r1cs_file, checksums).await?;
+        let pk_bytes_path = params::download_and_checksum(url, dir, pk_file, checksums).await?;
+
+        let r1cs = read(&r1cs_bytes_path)
+            .with_context(|| format!("while reading {}", r1cs_bytes_path.display()))?;
+        let pk = read(&pk_bytes_path)
+            .with_context(|| format!("while reading {}", pk_bytes_path.display()))?;
+        let circuit = read(&circuit_bytes_path)
+            .with_context(|| format!("while reading {}", circuit_bytes_path.display()))?;
 
         debug!("Creating Groth16 prover");
-        let inner = Groth16Prover::from_bytes(
-            r1cs_bytes.to_vec(),
-            pk_bytes.to_vec(),
-            circuit_bytes.to_vec(),
-        )?;
+        let inner = Groth16Prover::from_bytes(r1cs, pk, circuit)?;
 
         debug!("Groth16 prover created");
         Ok(Self { inner })

lgn-provers/src/provers/v1/groth16/euclid_prover.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::fs::File;
 
 use alloy::primitives::Address;
 use alloy::primitives::U256;
@@ -46,9 +47,10 @@ impl PreprocessingEuclidProver {
         checksums: &HashMap<String, blake3::Hash>,
         with_tracing: bool,
     ) -> anyhow::Result<Self> {
-        let params = params::download_and_checksum(url, dir, file, checksums).await?;
-        let reader = std::io::BufReader::new(params.as_ref());
-        let params = bincode::deserialize_from(reader)?;
+        let params_path = params::download_and_checksum(url, dir, file, checksums).await?;
+        let file = File::open(params_path)?;
+        let reader = unsafe { memmap2::Mmap::map(&file)? };
+        let params = bincode::deserialize(&reader)?;
         Ok(Self {
             params,
             with_tracing,

lgn-provers/src/provers/v1/preprocessing/euclid_prover.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::fs::File;
 
 use anyhow::Context;
 use lgn_messages::types::v1::query::tasks::MatchingRowInput;
@@ -51,9 +52,10 @@ impl QueryEuclidProver {
         checksums: &HashMap<String, blake3::Hash>,
         with_tracing: bool,
     ) -> anyhow::Result<Self> {
-        let params = params::download_and_checksum(url, dir, file, checksums).await?;
-        let reader = std::io::BufReader::new(params.as_ref());
-        let params = bincode::deserialize_from(reader)?;
+        let params_path = params::download_and_checksum(url, dir, file, checksums).await?;
+        let file = File::open(params_path)?;
+        let reader = unsafe { memmap2::Mmap::map(&file)? };
+        let params = bincode::deserialize(&reader)?;
         Ok(Self {
             params,
             with_tracing,

lgn-provers/src/provers/v1/query/euclid_prover.rs

@@ -20,7 +20,7 @@ alloy = { workspace = true, features = ["signers", "signer-local", "signer-keyst
 anyhow = { workspace = true }
 backtrace = { workspace = true }
 bincode.workspace = true
-blake3 = { workspace = true }
+blake3 = { workspace = true, features = ["mmap", "rayon"] }
 clap = { workspace = true, features = ["derive", "env", "help", "std", "suggestions"] }
 config = { workspace = true, features = ["toml"] }
 elliptic-curve = { workspace = true }