Skip to content

Commit

Permalink
feat(s2n-quic-dc): Use a new fixed-size map for path secret storage (#…
Browse files Browse the repository at this point in the history
…2337)

This tightly bounds the maximum memory usage of the path secret storage.
  • Loading branch information
Mark-Simulacrum authored Sep 28, 2024
1 parent 246ba90 commit 09a33a8
Show file tree
Hide file tree
Showing 6 changed files with 285 additions and 116 deletions.
1 change: 1 addition & 0 deletions dc/s2n-quic-dc/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ tokio = { version = "1", default-features = false, features = ["sync"] }
tracing = "0.1"
zerocopy = { version = "0.7", features = ["derive"] }
zeroize = "1"
parking_lot = "0.12"

[dev-dependencies]
bolero = "0.11"
Expand Down
170 changes: 170 additions & 0 deletions dc/s2n-quic-dc/src/fixed_map.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,170 @@
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! A fixed-allocation concurrent HashMap.
//!
//! This implements a concurrent map backed by a fixed-size allocation created at construction
//! time, with a fixed memory footprint. The expectation is that all storage is inline (to the
//! extent possible) reducing the likelihood.

use core::{
hash::Hash,
sync::atomic::{AtomicU8, Ordering},
};
use parking_lot::{MappedRwLockReadGuard, RwLock, RwLockReadGuard, RwLockUpgradableReadGuard};
use std::{collections::hash_map::RandomState, hash::BuildHasher};

pub struct Map<K, V, S = RandomState> {
slots: Box<[Slot<K, V>]>,
hash_builder: S,
}

impl<K, V, S> Map<K, V, S>
where
K: Hash + Eq,
S: BuildHasher,
{
pub fn with_capacity(entries: usize, hasher: S) -> Self {
let map = Map {
slots: (0..std::cmp::min(1, (entries + SLOT_CAPACITY) / SLOT_CAPACITY))
.map(|_| Slot::new())
.collect::<Vec<_>>()
.into_boxed_slice(),
hash_builder: hasher,
};
assert!(map.slots.len().is_power_of_two());
assert!(u32::try_from(map.slots.len()).is_ok());
map
}

pub fn clear(&self) {
for slot in self.slots.iter() {
slot.clear();
}
}

pub fn len(&self) -> usize {
self.slots.iter().map(|s| s.len()).sum()
}

// can't lend references to values outside of a lock, so Iterator interface doesn't work
#[allow(unused)]
pub fn iter(&self, mut f: impl FnMut(&K, &V)) {
for slot in self.slots.iter() {
// this feels more readable than flatten
#[allow(clippy::manual_flatten)]
for entry in slot.values.read().iter() {
if let Some(v) = entry {
f(&v.0, &v.1);
}
}
}
}

pub fn retain(&self, mut f: impl FnMut(&K, &V) -> bool) {
for slot in self.slots.iter() {
// this feels more readable than flatten
#[allow(clippy::manual_flatten)]
for entry in slot.values.write().iter_mut() {
if let Some(v) = entry {
if !f(&v.0, &v.1) {
*entry = None;
}
}
}
}
}

fn slot_by_hash(&self, key: &K) -> &Slot<K, V> {
let hash = self.hash_builder.hash_one(key);
// needed for bit-and modulus, checked in new as a non-debug assert!.
debug_assert!(self.slots.len().is_power_of_two());
let slot_idx = hash as usize & (self.slots.len() - 1);
&self.slots[slot_idx]
}

/// Returns Some(v) if overwriting a previous value for the same key.
pub fn insert(&self, key: K, value: V) -> Option<V> {
self.slot_by_hash(&key).put(key, value)
}

pub fn contains_key(&self, key: &K) -> bool {
self.get_by_key(key).is_some()
}

pub fn get_by_key(&self, key: &K) -> Option<MappedRwLockReadGuard<'_, V>> {
self.slot_by_hash(key).get_by_key(key)
}
}

// Balance of speed of access (put or get) and likelihood of false positive eviction.
const SLOT_CAPACITY: usize = 32;

struct Slot<K, V> {
next_write: AtomicU8,
values: RwLock<[Option<(K, V)>; SLOT_CAPACITY]>,
}

impl<K, V> Slot<K, V>
where
K: Hash + Eq,
{
fn new() -> Self {
Slot {
next_write: AtomicU8::new(0),
values: RwLock::new(std::array::from_fn(|_| None)),
}
}

fn clear(&self) {
*self.values.write() = std::array::from_fn(|_| None);
}

/// Returns Some(v) if overwriting a previous value for the same key.
fn put(&self, new_key: K, new_value: V) -> Option<V> {
let values = self.values.upgradable_read();
for (value_idx, value) in values.iter().enumerate() {
// overwrite if same key or if no key/value pair yet
if value.as_ref().map_or(true, |(k, _)| *k == new_key) {
let mut values = RwLockUpgradableReadGuard::upgrade(values);
let old = values[value_idx].take().map(|v| v.1);
values[value_idx] = Some((new_key, new_value));
return old;
}
}

let mut values = RwLockUpgradableReadGuard::upgrade(values);

// If `new_key` isn't already in this slot, replace one of the existing entries with the
// new key. For now we rotate through based on `next_write`.
let replacement = self.next_write.fetch_add(1, Ordering::Relaxed) as usize % SLOT_CAPACITY;
values[replacement] = Some((new_key, new_value));
None
}

fn get_by_key(&self, needle: &K) -> Option<MappedRwLockReadGuard<'_, V>> {
// Scan each value and check if our requested needle is present.
let values = self.values.read();
for (value_idx, value) in values.iter().enumerate() {
if value.as_ref().map_or(false, |(k, _)| *k == *needle) {
return Some(RwLockReadGuard::map(values, |values| {
&values[value_idx].as_ref().unwrap().1
}));
}
}

None
}

fn len(&self) -> usize {
let values = self.values.read();
let mut len = 0;
for value in values.iter().enumerate() {
len += value.1.is_some() as usize;
}
len
}
}

#[cfg(test)]
mod test;
33 changes: 33 additions & 0 deletions dc/s2n-quic-dc/src/fixed_map/test.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;

#[test]
fn slot_insert_and_get() {
let slot = Slot::new();
assert!(slot.get_by_key(&3).is_none());
assert_eq!(slot.put(3, "key 1"), None);
// still same slot, but new generation
assert_eq!(slot.put(3, "key 2"), Some("key 1"));
// still same slot, but new generation
assert_eq!(slot.put(3, "key 3"), Some("key 2"));

// new slot
assert_eq!(slot.put(5, "key 4"), None);
assert_eq!(slot.put(6, "key 4"), None);
}

#[test]
fn slot_clear() {
let slot = Slot::new();
assert_eq!(slot.put(3, "key 1"), None);
// still same slot, but new generation
assert_eq!(slot.put(3, "key 2"), Some("key 1"));
// still same slot, but new generation
assert_eq!(slot.put(3, "key 3"), Some("key 2"));

slot.clear();

assert_eq!(slot.len(), 0);
}
1 change: 1 addition & 0 deletions dc/s2n-quic-dc/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ pub mod control;
pub mod credentials;
pub mod crypto;
pub mod datagram;
mod fixed_map;
pub mod msg;
pub mod packet;
pub mod path;
Expand Down
Loading

0 comments on commit 09a33a8

Please sign in to comment.