perf: reuse a chain's sector list across opens instead of rewalking the FAT

benches/benchmark.rs

@@ -4,6 +4,8 @@ use std::fs::OpenOptions;
 use std::hint::black_box;
 use std::io::Cursor;
 use std::io::Read;
+use std::io::Seek;
+use std::io::SeekFrom;
 use std::io::Write;
 use tempfile::NamedTempFile;
 
@@ -119,5 +121,45 @@ fn criterion_benchmark(c: &mut Criterion) {
     read_group.finish();
 }
 
-criterion_group!(benches, criterion_benchmark);
+/// Reads one large stream the way a caller that streams it out would: in
+/// pieces, so the stream's buffer is refilled many times.
+fn large_stream_benchmark(c: &mut Criterion) {
+    let size = 256 * 1024 * 1024usize;
+    let buff = write_many_streams(1, size);
+
+    let mut group = c.benchmark_group("read_large_stream_memory");
+    group.sample_size(10);
+    group.throughput(Throughput::Bytes(size as u64));
+    group.bench_function("copy 256MiB", |b| {
+        b.iter(|| {
+            let mut comp = CompoundFile::open(Cursor::new(&buff)).unwrap();
+            let mut stream = comp.open_stream("test0").unwrap();
+            let n = std::io::copy(&mut stream, &mut std::io::sink()).unwrap();
+            black_box(n);
+        })
+    });
+    group.finish();
+
+    let reads = 1000u64;
+    let mut group = c.benchmark_group("seek_large_stream_memory");
+    group.sample_size(10);
+    group.throughput(Throughput::Elements(reads));
+    group.bench_function("1000 x seek+read 16B in 256MiB", |b| {
+        b.iter(|| {
+            let mut comp = CompoundFile::open(Cursor::new(&buff)).unwrap();
+            let mut stream = comp.open_stream("test0").unwrap();
+            let len = stream.len();
+            let mut chunk = [0u8; 16];
+            for i in 0..reads {
+                let offset = (i * 7919 * 4099) % (len - 16);
+                stream.seek(SeekFrom::Start(offset)).unwrap();
+                stream.read_exact(&mut chunk).unwrap();
+                black_box(chunk);
+            }
+        })
+    });
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark, large_stream_benchmark);
 criterion_main!(benches);

src/internal/alloc.rs

@@ -25,6 +25,13 @@ pub struct Allocator<F> {
     difat: Vec<u32>,
     fat: Vec<u32>,
     free_sectors: Vec<u32>,
+    /// The sector IDs of the most recently used chain, keyed by its start
+    /// sector.  A `Stream` opens its chain anew for every buffer refill,
+    /// and walking a long FAT chain from the start each time made reading
+    /// or writing a stream of `n` sectors cost `O(n^2)`.  The list is
+    /// handed to the next `Chain` opened for the same start sector, which
+    /// gives it back when dropped; any FAT change discards it.
+    chain_cache: Option<(u32, Vec<u32>)>,
 }
 
 impl<F> Allocator<F> {
@@ -41,6 +48,7 @@ impl<F> Allocator<F> {
             difat,
             fat,
             free_sectors: Vec::new(),
+            chain_cache: None,
         };
         alloc.validate(validation)?;
         Ok(alloc)
@@ -86,7 +94,26 @@ impl<F> Allocator<F> {
         start_sector_id: u32,
         init: SectorInit,
     ) -> io::Result<Chain<'_, F>> {
-        Chain::new(self, start_sector_id, init)
+        match self.chain_cache.take() {
+            Some((cached_start, sector_ids))
+                if cached_start == start_sector_id =>
+            {
+                Ok(Chain::from_sector_ids(self, sector_ids, init))
+            }
+            _ => Chain::new(self, start_sector_id, init),
+        }
+    }
+
+    /// Remembers the sector IDs of the chain starting at `start_sector_id`
+    /// so that the next `open_chain` for it need not walk the FAT again.
+    /// The IDs must reflect the current FAT.
+    pub(crate) fn cache_chain(
+        &mut self,
+        start_sector_id: u32,
+        sector_ids: Vec<u32>,
+    ) {
+        debug_assert_eq!(sector_ids.first(), Some(&start_sector_id));
+        self.chain_cache = Some((start_sector_id, sector_ids));
     }
 
     fn validate(&mut self, validation: Validation) -> io::Result<()> {
@@ -354,6 +381,7 @@ impl<F: Write + Seek> Allocator<F> {
     /// Sets `self.fat[index] = value`, and also writes that change to the
     /// underlying file.  The `index` must be <= `self.fat.len()`.
     fn set_fat(&mut self, index: u32, value: u32) -> io::Result<()> {
+        self.chain_cache = None;
         let index = index as usize;
         debug_assert!(index <= self.fat.len());
         let fat_entries_per_sector =
@@ -383,7 +411,9 @@ impl<F: Write + Seek> Allocator<F> {
 #[cfg(test)]
 mod tests {
     use super::Allocator;
-    use crate::internal::{consts, Sectors, Validation, Version};
+    use crate::internal::{
+        consts, Chain, SectorInit, Sectors, Validation, Version,
+    };
     use std::io::Cursor;
 
     fn make_sectors(
@@ -519,6 +549,57 @@ mod tests {
         make_allocator(difat, fat, Validation::Permissive);
     }
 
+    #[test]
+    fn chain_cache_follows_the_fat() {
+        let difat = vec![0];
+        let fat = vec![consts::FAT_SECTOR, 2, 3, consts::END_OF_CHAIN];
+        let mut alloc = make_allocator(difat, fat, Validation::Strict);
+        assert_eq!(alloc.chain_cache, None);
+
+        // Dropping a chain leaves its sector list behind for the next open.
+        let chain = alloc.open_chain(1, SectorInit::Zero).unwrap();
+        assert_eq!(chain.sector_ids(), &[1, 2, 3]);
+        drop(chain);
+        assert_eq!(alloc.chain_cache, Some((1, vec![1, 2, 3])));
+
+        // A chain that grew hands back the grown list.
+        let mut chain = alloc.open_chain(1, SectorInit::Zero).unwrap();
+        chain.set_len(4 * 512).unwrap();
+        assert_eq!(chain.sector_ids(), &[1, 2, 3, 4]);
+        drop(chain);
+        assert_eq!(alloc.chain_cache, Some((1, vec![1, 2, 3, 4])));
+        let chain = Chain::new(&mut alloc, 1, SectorInit::Zero).unwrap();
+        assert_eq!(chain.sector_ids(), &[1, 2, 3, 4]);
+        drop(chain);
+
+        // A chain that shrank hands back the shortened list.
+        let mut chain = alloc.open_chain(1, SectorInit::Zero).unwrap();
+        chain.set_len(2 * 512).unwrap();
+        assert_eq!(chain.sector_ids(), &[1, 2]);
+        drop(chain);
+        assert_eq!(alloc.chain_cache, Some((1, vec![1, 2])));
+        let chain = Chain::new(&mut alloc, 1, SectorInit::Zero).unwrap();
+        assert_eq!(chain.sector_ids(), &[1, 2]);
+        drop(chain);
+
+        // Opening a different chain does not use the cached one.
+        let start = alloc.begin_chain(SectorInit::Zero).unwrap();
+        let chain = alloc.open_chain(start, SectorInit::Zero).unwrap();
+        assert_eq!(chain.sector_ids(), &[start]);
+        drop(chain);
+        assert_eq!(alloc.chain_cache, Some((start, vec![start])));
+
+        // Any FAT change from outside a chain discards the cache.
+        alloc.free_chain(1).unwrap();
+        assert_eq!(alloc.chain_cache, None);
+
+        // A freed chain leaves nothing behind.
+        let chain = alloc.open_chain(start, SectorInit::Zero).unwrap();
+        chain.free().unwrap();
+        assert_eq!(alloc.chain_cache, None);
+        assert_eq!(alloc.fat[start as usize], consts::FREE_SECTOR);
+    }
+
     #[test]
     #[should_panic(expected = "Malformed FAT (sector 3 pointed to twice)")]
     fn double_pointee() {

src/internal/chain.rs

@@ -1,6 +1,7 @@
 use crate::internal::{consts, Allocator, SectorInit};
 use std::cmp;
 use std::io::{self, Read, Seek, SeekFrom, Write};
+use std::mem;
 
 //===========================================================================//
 
@@ -33,6 +34,15 @@ impl<'a, F> Chain<'a, F> {
         Ok(Chain { allocator, init, sector_ids, offset_from_start: 0 })
     }
 
+    /// Creates a chain from sector IDs that are known to match the FAT.
+    pub(crate) fn from_sector_ids(
+        allocator: &'a mut Allocator<F>,
+        sector_ids: Vec<u32>,
+        init: SectorInit,
+    ) -> Chain<'a, F> {
+        Chain { allocator, init, sector_ids, offset_from_start: 0 }
+    }
+
     pub fn start_sector_id(&self) -> u32 {
         self.sector_ids.first().copied().unwrap_or(consts::END_OF_CHAIN)
     }
@@ -61,11 +71,13 @@ impl<'a, F: Write + Seek> Chain<'a, F> {
         if new_num_sectors == 0 {
             if let Some(&start_sector) = self.sector_ids.first() {
                 self.allocator.free_chain(start_sector)?;
+                self.sector_ids.clear();
             }
         } else if new_num_sectors <= self.sector_ids.len() {
             if new_num_sectors < self.sector_ids.len() {
                 self.allocator
                     .free_chain_after(self.sector_ids[new_num_sectors - 1])?;
+                self.sector_ids.truncate(new_num_sectors);
             }
             // TODO: init remainder of final sector
         } else {
@@ -83,8 +95,23 @@ impl<'a, F: Write + Seek> Chain<'a, F> {
         Ok(())
     }
 
-    pub fn free(self) -> io::Result<()> {
-        self.allocator.free_chain(self.start_sector_id())
+    pub fn free(mut self) -> io::Result<()> {
+        let start_sector_id = self.start_sector_id();
+        self.sector_ids.clear();
+        self.allocator.free_chain(start_sector_id)
+    }
+}
+
+impl<'a, F> Drop for Chain<'a, F> {
+    fn drop(&mut self) {
+        // The list is kept in step with every change this chain made to the
+        // FAT, so the next chain opened for the same start sector can reuse
+        // it instead of walking the FAT again.
+        if !self.sector_ids.is_empty() {
+            let start_sector_id = self.sector_ids[0];
+            self.allocator
+                .cache_chain(start_sector_id, mem::take(&mut self.sector_ids));
+        }
     }
 }