Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle usize overflow on big cache capacity #28

Merged
merged 7 commits into from
Aug 29, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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/
Expand Down
50 changes: 37 additions & 13 deletions src/common/frequency_sketch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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()
}
}

Expand Down Expand Up @@ -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;
Expand Down
67 changes: 66 additions & 1 deletion src/sync/base_cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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");
};
}
}
48 changes: 47 additions & 1 deletion src/unsync/cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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");
};
}
}