From 58de0b2d1016c964d6f18f042f10e2a744415b08 Mon Sep 17 00:00:00 2001
From: Michal Trybus <komar007@gmail.com>
Date: Thu, 18 May 2023 00:45:46 +0200
Subject: [PATCH] feat(mmap): new feature for mmap-based optimization (#49)

Ref: https://github.com/zkat/cacache-rs/issues/48
---
 Cargo.toml           |  5 +--
 src/content/write.rs | 77 +++++++++++++++++++++++---------------------
 2 files changed, 44 insertions(+), 38 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 6e49c4e..db34ec5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,7 +16,7 @@ digest = "0.10.6"
 either = "1.6.1"
 futures = "0.3.17"
 hex = "0.4.3"
-memmap2 = "0.5.8"
+memmap2 = { version = "0.5.8", optional = true }
 miette = "5.7.0"
 reflink = "0.1.3"
 serde = "1.0.130"
@@ -53,6 +53,7 @@ name = "benchmarks"
 harness = false
 
 [features]
-default = ["async-std"]
+default = ["async-std", "mmap"]
+mmap = ["memmap2"]
 link_to = []
 tokio-runtime = ["tokio", "tokio-stream"]
diff --git a/src/content/write.rs b/src/content/write.rs
index a5b3cd6..41b7db3 100644
--- a/src/content/write.rs
+++ b/src/content/write.rs
@@ -6,6 +6,7 @@ use std::sync::Mutex;
 use std::task::{Context, Poll};
 
 use futures::prelude::*;
+#[cfg(feature = "mmap")]
 use memmap2::MmapMut;
 use ssri::{Algorithm, Integrity, IntegrityOpts};
 use tempfile::NamedTempFile;
@@ -14,8 +15,23 @@ use crate::async_lib::{AsyncWrite, JoinHandle};
 use crate::content::path;
 use crate::errors::{IoErrorExt, Result};
 
+#[cfg(feature = "mmap")]
 pub const MAX_MMAP_SIZE: usize = 1024 * 1024;
 
+#[cfg(not(feature = "mmap"))]
+struct MmapMut;
+
+#[cfg(not(feature = "mmap"))]
+impl MmapMut {
+    fn flush_async(&self) -> std::io::Result<()> {
+        panic!()
+    }
+
+    fn copy_from_slice(&self, _: &[u8]) {
+        panic!()
+    }
+}
+
 pub struct Writer {
     cache: PathBuf,
     builder: IntegrityOpts,
@@ -44,24 +60,7 @@ impl Writer {
                 tmp_path_clone.display()
             )
         })?;
-        let mmap = if let Some(size) = size {
-            if size <= MAX_MMAP_SIZE {
-                tmpfile
-                    .as_file_mut()
-                    .set_len(size as u64)
-                    .with_context(|| {
-                        format!(
-                            "Failed to configure file length for temp file at {}",
-                            tmpfile.path().display()
-                        )
-                    })?;
-                unsafe { MmapMut::map_mut(tmpfile.as_file()).ok() }
-            } else {
-                None
-            }
-        } else {
-            None
-        };
+        let mmap = make_mmap(&mut tmpfile, size)?;
         Ok(Writer {
             cache: cache_path,
             builder: IntegrityOpts::new().algorithm(algo),
@@ -162,24 +161,7 @@ impl AsyncWriter {
                 )
             })?;
         let mut tmpfile = crate::async_lib::create_named_tempfile(tmp_path).await?;
-        let mmap = if let Some(size) = size {
-            if size <= MAX_MMAP_SIZE {
-                tmpfile
-                    .as_file_mut()
-                    .set_len(size as u64)
-                    .with_context(|| {
-                        format!(
-                            "Failed to configure file length for temp file at {}",
-                            tmpfile.path().display()
-                        )
-                    })?;
-                unsafe { MmapMut::map_mut(tmpfile.as_file()).ok() }
-            } else {
-                None
-            }
-        } else {
-            None
-        };
+        let mmap = make_mmap(&mut tmpfile, size)?;
         Ok(AsyncWriter(Mutex::new(State::Idle(Some(Inner {
             cache: cache_path,
             builder: IntegrityOpts::new().algorithm(algo),
@@ -428,6 +410,29 @@ impl AsyncWriter {
     }
 }
 
+#[cfg(feature = "mmap")]
+fn make_mmap(tmpfile: &mut NamedTempFile, size: Option<usize>) -> Result<Option<MmapMut>> {
+    if let Some(size @ 0..=MAX_MMAP_SIZE) = size {
+        tmpfile
+            .as_file_mut()
+            .set_len(size as u64)
+            .with_context(|| {
+                format!(
+                    "Failed to configure file length for temp file at {}",
+                    tmpfile.path().display()
+                )
+            })?;
+        Ok(unsafe { MmapMut::map_mut(tmpfile.as_file()).ok() })
+    } else {
+        Ok(None)
+    }
+}
+
+#[cfg(not(feature = "mmap"))]
+fn make_mmap(_: &mut NamedTempFile, _: Option<usize>) -> Result<Option<MmapMut>> {
+    Ok(None)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;