Merge pull request #28 from moka-rs/fix-usize-overflow

Handle `usize` overflow on big cache capacity

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Moka — Change Log
 
+## Version 0.5.2
+
+### Fixed
+
+- `usize` overflow on big cache capacity. ([#28][gh-pull-0028])
+
+
 ## Version 0.5.1
 
 ### Changed
@@ -81,6 +88,7 @@
 
 [caffeine-git]: https://github.com/ben-manes/caffeine
 
+[gh-pull-0028]: https://github.com/moka-rs/moka/pull/28/
 [gh-pull-0022]: https://github.com/moka-rs/moka/pull/22/
 [gh-pull-0020]: https://github.com/moka-rs/moka/pull/20/
 [gh-pull-0019]: https://github.com/moka-rs/moka/pull/19/

src/common/frequency_sketch.rs

@@ -15,10 +15,10 @@
 /// a time window. The maximum frequency of an element is limited to 15 (4-bits)
 /// and an aging process periodically halves the popularity of all elements.
 pub(crate) struct FrequencySketch {
-    sample_size: usize,
-    table_mask: usize,
-    table: Vec<u64>,
-    size: usize,
+    sample_size: u32,
+    table_mask: u64,
+    table: Box<[u64]>,
+    size: u32,
 }
 
 // A mixture of seeds from FNV-1a, CityHash, and Murmur3. (Taken from Caffeine)
@@ -68,19 +68,35 @@ static ONE_MASK: u64 = 0x1111_1111_1111_1111;
 
 impl FrequencySketch {
     /// Creates a frequency sketch with the capacity.
-    pub(crate) fn with_capacity(cap: usize) -> Self {
-        let maximum = cap.min((i32::MAX >> 1) as usize);
+    pub(crate) fn with_capacity(cap: u32) -> Self {
+        // The max byte size of the table, Box<[u64; table_size]>
+        //
+        // | Pointer width    | Max size |
+        // |:-----------------|---------:|
+        // | 16 bit           |    8 KiB |
+        // | 32 bit           |  128 MiB |
+        // | 64 bit or bigger |    8 GiB |
+
+        let maximum = if cfg!(target_pointer_width = "16") {
+            cap.min(1024)
+        } else if cfg!(target_pointer_width = "32") {
+            cap.min(2u32.pow(24)) // about 16 millions
+        } else {
+            // Same to Caffeine's limit:
+            //   `Integer.MAX_VALUE >>> 1` with `ceilingPowerOfTwo()` applied.
+            cap.min(2u32.pow(30)) // about 1 billion
+        };
         let table_size = if maximum == 0 {
             1
         } else {
             maximum.next_power_of_two()
         };
-        let table = vec![0; table_size];
-        let table_mask = 0.max(table_size - 1);
+        let table = vec![0; table_size as usize].into_boxed_slice();
+        let table_mask = 0.max(table_size - 1) as u64;
         let sample_size = if cap == 0 {
             10
         } else {
-            maximum.saturating_mul(10).min(i32::MAX as usize)
+            maximum.saturating_mul(10).min(i32::MAX as u32)
         };
         Self {
             sample_size,
@@ -141,20 +157,28 @@ impl FrequencySketch {
     /// Reduces every counter by half of its original value.
     fn reset(&mut self) {
         let mut count = 0u32;
-        for entry in &mut self.table {
+        for entry in self.table.iter_mut() {
             // Count number of odd numbers.
             count += (*entry & ONE_MASK).count_ones();
             *entry = (*entry >> 1) & RESET_MASK;
         }
-        self.size = (self.size >> 1) - (count >> 2) as usize;
+        self.size = (self.size >> 1) - (count >> 2);
     }
 
     /// Returns the table index for the counter at the specified depth.
     fn index_of(&self, hash: u64, depth: u8) -> usize {
         let i = depth as usize;
         let mut hash = hash.wrapping_add(SEED[i]).wrapping_mul(SEED[i]);
         hash += hash >> 32;
-        hash as usize & self.table_mask
+        (hash & self.table_mask) as usize
+    }
+}
+
+// Methods only available for testing.
+#[cfg(test)]
+impl FrequencySketch {
+    pub(crate) fn table_len(&self) -> usize {
+        self.table.len()
     }
 }
 
@@ -229,7 +253,7 @@ mod tests {
         let mut sketch = FrequencySketch::with_capacity(64);
         let hasher = hasher();
 
-        for i in 1..(20 * sketch.table.len()) {
+        for i in 1..(20 * sketch.table.len() as u32) {
             sketch.increment(hasher(i));
             if sketch.size != i {
                 reset = true;

src/sync/base_cache.rs

@@ -19,6 +19,7 @@ use quanta::{Clock, Instant};
 use std::{
     borrow::Borrow,
     collections::hash_map::RandomState,
+    convert::TryInto,
     hash::{BuildHasher, Hash, Hasher},
     ptr::NonNull,
     rc::Rc,
@@ -392,7 +393,12 @@ where
             initial_capacity,
             build_hasher.clone(),
         );
-        let skt_capacity = usize::max(max_capacity * 32, 100);
+
+        // Ensure skt_capacity fits in a range of `128u32..=u32::MAX`.
+        let skt_capacity = max_capacity
+            .try_into() // Convert to u32.
+            .unwrap_or(u32::MAX)
+            .max(128);
         let frequency_sketch = FrequencySketch::with_capacity(skt_capacity);
 
         Self {
@@ -1063,3 +1069,62 @@ fn is_expired_entry_wo(
     }
     false
 }
+
+#[cfg(test)]
+mod tests {
+    use super::BaseCache;
+
+    #[cfg_attr(target_pointer_width = "16", ignore)]
+    #[test]
+    fn test_skt_capacity_will_not_overflow() {
+        use std::collections::hash_map::RandomState;
+
+        // power of two
+        let pot = |exp| 2_usize.pow(exp);
+
+        let ensure_sketch_len = |max_capacity, len, name| {
+            let cache = BaseCache::<u8, u8>::new(
+                max_capacity,
+                None,
+                RandomState::default(),
+                None,
+                None,
+                false,
+            );
+            assert_eq!(
+                cache.inner.frequency_sketch.read().table_len(),
+                len,
+                "{}",
+                name
+            );
+        };
+
+        if cfg!(target_pointer_width = "32") {
+            let pot24 = pot(24);
+            let pot16 = pot(16);
+            ensure_sketch_len(0, 128, "0");
+            ensure_sketch_len(128, 128, "128");
+            ensure_sketch_len(pot16, pot16, "pot16");
+            // due to ceiling to next_power_of_two
+            ensure_sketch_len(pot16 + 1, pot(17), "pot16 + 1");
+            // due to ceiling to next_power_of_two
+            ensure_sketch_len(pot24 - 1, pot24, "pot24 - 1");
+            ensure_sketch_len(pot24, pot24, "pot24");
+            ensure_sketch_len(pot(27), pot24, "pot(27)");
+            ensure_sketch_len(usize::MAX, pot24, "usize::MAX");
+        } else {
+            // target_pointer_width: 64 or larger.
+            let pot30 = pot(30);
+            let pot16 = pot(16);
+            ensure_sketch_len(0, 128, "0");
+            ensure_sketch_len(128, 128, "128");
+            ensure_sketch_len(pot16, pot16, "pot16");
+            // due to ceiling to next_power_of_two
+            ensure_sketch_len(pot16 + 1, pot(17), "pot16 + 1");
+            // due to ceiling to next_power_of_two
+            ensure_sketch_len(pot30 - 1, pot30, "pot30- 1");
+            ensure_sketch_len(pot30, pot30, "pot30");
+            ensure_sketch_len(usize::MAX, pot30, "usize::MAX");
+        };
+    }
+}

src/unsync/cache.rs

@@ -9,6 +9,7 @@ use quanta::{Clock, Instant};
 use std::{
     borrow::Borrow,
     collections::{hash_map::RandomState, HashMap},
+    convert::TryInto,
     hash::{BuildHasher, Hash, Hasher},
     ptr::NonNull,
     rc::Rc,
@@ -153,7 +154,12 @@ where
             initial_capacity.unwrap_or_default(),
             build_hasher.clone(),
         );
-        let skt_capacity = usize::max(max_capacity * 32, 100);
+
+        // Ensure skt_capacity fits in a range of `128u32..=u32::MAX`.
+        let skt_capacity = max_capacity
+            .try_into() // Convert to u32.
+            .unwrap_or(u32::MAX)
+            .max(128);
         let frequency_sketch = FrequencySketch::with_capacity(skt_capacity);
         Self {
             max_capacity,
@@ -766,4 +772,44 @@ mod tests {
         assert_eq!(cache.get(&"b"), None);
         assert!(cache.cache.is_empty());
     }
+
+    #[cfg_attr(target_pointer_width = "16", ignore)]
+    #[test]
+    fn test_skt_capacity_will_not_overflow() {
+        // power of two
+        let pot = |exp| 2_usize.pow(exp);
+
+        let ensure_sketch_len = |max_capacity, len, name| {
+            let cache = Cache::<u8, u8>::new(max_capacity);
+            assert_eq!(cache.frequency_sketch.table_len(), len, "{}", name);
+        };
+
+        if cfg!(target_pointer_width = "32") {
+            let pot24 = pot(24);
+            let pot16 = pot(16);
+            ensure_sketch_len(0, 128, "0");
+            ensure_sketch_len(128, 128, "128");
+            ensure_sketch_len(pot16, pot16, "pot16");
+            // due to ceiling to next_power_of_two
+            ensure_sketch_len(pot16 + 1, pot(17), "pot16 + 1");
+            // due to ceiling to next_power_of_two
+            ensure_sketch_len(pot24 - 1, pot24, "pot24 - 1");
+            ensure_sketch_len(pot24, pot24, "pot24");
+            ensure_sketch_len(pot(27), pot24, "pot(27)");
+            ensure_sketch_len(usize::MAX, pot24, "usize::MAX");
+        } else {
+            // target_pointer_width: 64 or larger.
+            let pot30 = pot(30);
+            let pot16 = pot(16);
+            ensure_sketch_len(0, 128, "0");
+            ensure_sketch_len(128, 128, "128");
+            ensure_sketch_len(pot16, pot16, "pot16");
+            // due to ceiling to next_power_of_two
+            ensure_sketch_len(pot16 + 1, pot(17), "pot16 + 1");
+            // due to ceiling to next_power_of_two
+            ensure_sketch_len(pot30 - 1, pot30, "pot30- 1");
+            ensure_sketch_len(pot30, pot30, "pot30");
+            ensure_sketch_len(usize::MAX, pot30, "usize::MAX");
+        };
+    }
 }