improve archive index reading performance

syphar · syphar · commit c173e278a94c · 2021-11-07T09:03:57.000+01:00
* use a streaming parser to find a file in the local index
* switch to memmap as reader
* use zero-copy deserialization for the key, change key data type
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -77,6 +77,7 @@ rusoto_credential = "0.46.0"
 # Data serialization and deserialization
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+memmap = "0.7.0"
 
 # iron dependencies
 iron = "0.6"
diff --git a/src/storage/archive_index.rs b/src/storage/archive_index.rs
@@ -1,10 +1,14 @@
 use crate::error::Result;
 use crate::storage::{compression::CompressionAlgorithm, FileRange};
 use anyhow::{bail, Context as _};
-use serde::{Deserialize, Serialize};
+use memmap::MmapOptions;
+use serde::de::DeserializeSeed;
+use serde::de::{IgnoredAny, MapAccess, Visitor};
+use serde::{Deserialize, Deserializer, Serialize};
 use std::collections::HashMap;
-use std::io;
-use std::path::{Path, PathBuf};
+use std::fmt;
+use std::path::Path;
+use std::{fs, io};
 
 #[derive(Deserialize, Serialize)]
 pub(crate) struct FileInfo {
@@ -21,51 +25,155 @@ impl FileInfo {
     }
 }
 
-#[derive(Deserialize, Serialize)]
-pub(crate) struct Index {
-    files: HashMap<PathBuf, FileInfo>,
+#[derive(Serialize)]
+struct Index {
+    files: HashMap<String, FileInfo>,
 }
 
-impl Index {
-    pub(crate) fn load(reader: impl io::Read) -> Result<Index> {
-        serde_cbor::from_reader(reader).context("deserialization error")
+pub(crate) fn create<R: io::Read + io::Seek, W: io::Write>(
+    zipfile: &mut R,
+    writer: &mut W,
+) -> Result<()> {
+    let mut archive = zip::ZipArchive::new(zipfile)?;
+
+    // get file locations
+    let mut files: HashMap<String, FileInfo> = HashMap::with_capacity(archive.len());
+    for i in 0..archive.len() {
+        let zf = archive.by_index(i)?;
+
+        files.insert(
+            zf.name().to_string(),
+            FileInfo {
+                range: FileRange::new(zf.data_start(), zf.data_start() + zf.compressed_size() - 1),
+                compression: match zf.compression() {
+                    zip::CompressionMethod::Bzip2 => CompressionAlgorithm::Bzip2,
+                    c => bail!("unsupported compression algorithm {} in zip-file", c),
+                },
+            },
+        );
     }
 
-    pub(crate) fn save(&self, writer: impl io::Write) -> Result<()> {
-        serde_cbor::to_writer(writer, self).context("serialization error")
+    serde_cbor::to_writer(writer, &Index { files }).context("serialization error")
+}
+
+pub(crate) fn find_in_slice(bytes: &[u8], search_for: &str) -> Result<Option<FileInfo>> {
+    let mut deserializer = serde_cbor::Deserializer::from_slice(bytes);
+
+    /// This visitor will just find the `files` element in the top-level map.
+    /// Then it will call the `FindFileVisitor` that should find the actual
+    /// FileInfo for the path we are searching for.
+    struct FindFileListVisitor {
+        search_for: String,
     }
 
-    pub(crate) fn new_from_zip<R: io::Read + io::Seek>(zipfile: &mut R) -> Result<Index> {
-        let mut archive = zip::ZipArchive::new(zipfile)?;
-
-        // get file locations
-        let mut files: HashMap<PathBuf, FileInfo> = HashMap::with_capacity(archive.len());
-        for i in 0..archive.len() {
-            let zf = archive.by_index(i)?;
-
-            files.insert(
-                PathBuf::from(zf.name()),
-                FileInfo {
-                    range: FileRange::new(
-                        zf.data_start(),
-                        zf.data_start() + zf.compressed_size() - 1,
-                    ),
-                    compression: match zf.compression() {
-                        zip::CompressionMethod::Bzip2 => CompressionAlgorithm::Bzip2,
-                        c => bail!("unsupported compression algorithm {} in zip-file", c),
-                    },
-                },
-            );
+    impl FindFileListVisitor {
+        pub fn new(path: String) -> Self {
+            FindFileListVisitor { search_for: path }
         }
+    }
 
-        Ok(Index { files })
+    impl<'de> Visitor<'de> for FindFileListVisitor {
+        type Value = Option<FileInfo>;
+
+        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+            write!(formatter, "a map with a 'files' key")
+        }
+
+        fn visit_map<V>(self, mut map: V) -> Result<Self::Value, V::Error>
+        where
+            V: MapAccess<'de>,
+        {
+            /// This visitor will walk the full `files` map and search for
+            /// the path we want to have.
+            /// Return value is just the `FileInfo` we want to have, or
+            /// `None`.
+            struct FindFileVisitor {
+                search_for: String,
+            }
+
+            impl FindFileVisitor {
+                pub fn new(search_for: String) -> Self {
+                    FindFileVisitor { search_for }
+                }
+            }
+
+            impl<'de> DeserializeSeed<'de> for FindFileVisitor {
+                type Value = Option<FileInfo>;
+                fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
+                where
+                    D: Deserializer<'de>,
+                {
+                    deserializer.deserialize_map(self)
+                }
+            }
+
+            impl<'de> Visitor<'de> for FindFileVisitor {
+                type Value = Option<FileInfo>;
+                fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                    write!(
+                        formatter,
+                        "a map with path => FileInfo, searching for path {:?}",
+                        self.search_for
+                    )
+                }
+                fn visit_map<V>(self, mut map: V) -> Result<Self::Value, V::Error>
+                where
+                    V: MapAccess<'de>,
+                {
+                    while let Some(key) = map.next_key::<&str>()? {
+                        if key == self.search_for {
+                            let value = map.next_value::<FileInfo>()?;
+                            // skip over the rest of the data without really parsing it.
+                            // If we don't do this the serde_cbor deserializer fails because not
+                            // the whole map is consumed.
+                            while map.next_entry::<IgnoredAny, IgnoredAny>()?.is_some() {}
+                            return Ok(Some(value));
+                        } else {
+                            // skip parsing the FileInfo structure when the key doesn't match.
+                            map.next_value::<IgnoredAny>()?;
+                        }
+                    }
+
+                    Ok(None)
+                }
+            }
+
+            while let Some(key) = map.next_key::<&str>()? {
+                if key == "files" {
+                    return map.next_value_seed(FindFileVisitor::new(self.search_for));
+                }
+            }
+
+            Ok(None)
+        }
     }
 
-    pub(crate) fn find_file<P: AsRef<Path>>(&self, path: P) -> Result<&FileInfo> {
-        self.files
-            .get(path.as_ref())
-            .ok_or_else(|| super::PathNotFoundError.into())
+    impl<'de> DeserializeSeed<'de> for FindFileListVisitor {
+        type Value = Option<FileInfo>;
+
+        fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
+        where
+            D: Deserializer<'de>,
+        {
+            deserializer.deserialize_map(self)
+        }
     }
+
+    Ok(FindFileListVisitor::new(search_for.to_string()).deserialize(&mut deserializer)?)
+}
+
+pub(crate) fn find_in_file<P: AsRef<Path>>(
+    archive_index_path: P,
+    search_for: &str,
+) -> Result<Option<FileInfo>> {
+    let file = fs::File::open(archive_index_path).context("could not open file")?;
+    let mmap = unsafe {
+        MmapOptions::new()
+            .map(&file)
+            .context("could not create memory map")?
+    };
+
+    find_in_slice(&mmap, search_for)
 }
 
 #[cfg(test)]
@@ -74,14 +182,6 @@ mod tests {
     use std::io::Write;
     use zip::write::FileOptions;
 
-    fn validate_index(index: &Index) {
-        assert_eq!(index.files.len(), 1);
-
-        let fi = index.files.get(&PathBuf::from("testfile1")).unwrap();
-        assert_eq!(fi.range, FileRange::new(39, 459));
-        assert_eq!(fi.compression, CompressionAlgorithm::Bzip2);
-    }
-
     #[test]
     fn index_create_save_load() {
         let mut tf = tempfile::tempfile().unwrap();
@@ -98,13 +198,13 @@ mod tests {
         archive.write_all(&objectcontent).unwrap();
         tf = archive.finish().unwrap();
 
-        let index = Index::new_from_zip(&mut tf).unwrap();
-        validate_index(&index);
-
         let mut buf = Vec::new();
-        index.save(&mut buf).unwrap();
+        create(&mut tf, &mut buf).unwrap();
+
+        let fi = find_in_slice(&buf, "testfile1").unwrap().unwrap();
+        assert_eq!(fi.range, FileRange::new(39, 459));
+        assert_eq!(fi.compression, CompressionAlgorithm::Bzip2);
 
-        let new_index = Index::load(io::Cursor::new(&buf)).unwrap();
-        validate_index(&new_index);
+        assert!(find_in_slice(&buf, "some_other_file").unwrap().is_none());
     }
 }
diff --git a/src/storage/mod.rs b/src/storage/mod.rs
@@ -200,8 +200,8 @@ impl Storage {
     }
 
     pub(crate) fn exists_in_archive(&self, archive_path: &str, path: &str) -> Result<bool> {
-        match self.get_index_for(archive_path) {
-            Ok(index) => Ok(index.find_file(path).is_ok()),
+        match self.get_index_filename(archive_path) {
+            Ok(index_filename) => Ok(archive_index::find_in_file(index_filename, path)?.is_some()),
             Err(err) => {
                 if err.downcast_ref::<PathNotFoundError>().is_some() {
                     Ok(false)
@@ -245,17 +245,15 @@ impl Storage {
         Ok(blob)
     }
 
-    fn get_index_for(&self, archive_path: &str) -> Result<archive_index::Index> {
+    fn get_index_filename(&self, archive_path: &str) -> Result<PathBuf> {
         // remote/folder/and/x.zip.index
         let remote_index_path = format!("{}.index", archive_path);
         let local_index_path = self
             .config
             .local_archive_cache_path
             .join(&remote_index_path);
 
-        if local_index_path.exists() {
-            archive_index::Index::load(io::BufReader::new(fs::File::open(local_index_path)?))
-        } else {
+        if !local_index_path.exists() {
             let index_content = self.get(&remote_index_path, std::usize::MAX)?.content;
 
             fs::create_dir_all(
@@ -265,9 +263,9 @@ impl Storage {
             )?;
             let mut file = fs::File::create(&local_index_path)?;
             file.write_all(&index_content)?;
-
-            archive_index::Index::load(&mut &index_content[..])
         }
+
+        Ok(local_index_path)
     }
 
     pub(crate) fn get_from_archive(
@@ -280,8 +278,8 @@ impl Storage {
         if let Some(ref mut t) = fetch_time {
             t.step("find path in index");
         }
-        let index = self.get_index_for(archive_path)?;
-        let info = index.find_file(path)?;
+        let info = archive_index::find_in_file(self.get_index_filename(archive_path)?, path)?
+            .ok_or(PathNotFoundError)?;
 
         if let Some(t) = fetch_time {
             t.step("range request");
@@ -336,9 +334,8 @@ impl Storage {
         }
 
         let mut zip_content = zip.finish()?.into_inner();
-        let index = archive_index::Index::new_from_zip(&mut io::Cursor::new(&mut zip_content))?;
         let mut index_content = vec![];
-        index.save(&mut index_content)?;
+        archive_index::create(&mut io::Cursor::new(&mut zip_content), &mut index_content)?;
         let alg = CompressionAlgorithm::default();
         let compressed_index_content = compress(&index_content[..], alg)?;
 
