Repository: quinn-rs/quinn
Pull request #2137 (2025-01-26 20:42:11)

Add `TokenMemoryCache`
Implements `TokenMemoryCache`, and makes it the default `TokenStore` for clients.

gretchenfrage: Switched from bespoke LRU cache implementation to Ralith's `lru_slab` crate. Saves a lot of code for sure, although I don't wanna merge until we decide what we wanna do about https://github.com/Ralith/lru-slab/issues/1

Switched from bespoke in-place queue implementation to just using `VecDeque`.

Together, these changes cut the size of this PR roughly in half. Nice!
gretchenfrage: @djc do you want to review this as well? It's fine either way.
djc: Yes, will try to take a look soon.

Patch:
diff --git a/Cargo.lock b/Cargo.lock
index eda3fdee..96c00226 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1283,6 +1283,12 @@ dependencies = [
  "value-bag",
 ]
 
+[[package]]
+name = "lru-slab"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
+
 [[package]]
 name = "matchers"
 version = "0.1.0"
@@ -1590,7 +1596,9 @@ dependencies = [
  "getrandom 0.3.2",
  "hex-literal",
  "lazy_static",
+ "lru-slab",
  "rand",
+ "rand_pcg",
  "rcgen",
  "ring",
  "rustc-hash 2.1.1",
@@ -1665,6 +1673,15 @@ dependencies = [
  "getrandom 0.3.2",
 ]
 
+[[package]]
+name = "rand_pcg"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b48ac3f7ffaab7fac4d2376632268aa5f89abdb55f7ebf8f4d11fffccb2320f7"
+dependencies = [
+ "rand_core",
+]
+
 [[package]]
 name = "rcgen"
 version = "0.13.2"
diff --git a/Cargo.toml b/Cargo.toml
index 1928245e..c058ac5e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -27,6 +27,7 @@ futures-io = "0.3.19"
 getrandom = { version = "0.3", default-features = false }
 hdrhistogram = { version = "7.2", default-features = false }
 hex-literal = "0.4"
+lru-slab = "0.1.2"
 lazy_static = "1"
 log = "0.4"
 once_cell = "1.19"
diff --git a/quinn-proto/Cargo.toml b/quinn-proto/Cargo.toml
index e59ea393..ff0fb301 100644
--- a/quinn-proto/Cargo.toml
+++ b/quinn-proto/Cargo.toml
@@ -36,6 +36,7 @@ rustls-log = ["rustls?/logging"]
 arbitrary = { workspace = true, optional = true }
 aws-lc-rs = { workspace = true, optional = true }
 bytes = { workspace = true }
+lru-slab = { workspace = true }
 rustc-hash = { workspace = true }
 rand = { workspace = true }
 ring = { workspace = true, optional = true }
@@ -56,6 +57,7 @@ web-time = { workspace = true }
 
 [dev-dependencies]
 assert_matches = { workspace = true }
+rand_pcg = "0.9"
 hex-literal = { workspace = true  }
 rcgen = { workspace = true }
 tracing-subscriber = { workspace = true }
diff --git a/quinn-proto/src/config/mod.rs b/quinn-proto/src/config/mod.rs
index 8d9364a6..69fe2a09 100644
--- a/quinn-proto/src/config/mod.rs
+++ b/quinn-proto/src/config/mod.rs
@@ -14,8 +14,8 @@ use thiserror::Error;
 #[cfg(any(feature = "rustls-aws-lc-rs", feature = "rustls-ring"))]
 use crate::crypto::rustls::{QuicServerConfig, configured_provider};
 use crate::{
-    DEFAULT_SUPPORTED_VERSIONS, Duration, MAX_CID_SIZE, NoneTokenLog, NoneTokenStore,
-    RandomConnectionIdGenerator, SystemTime, TokenLog, TokenStore, VarInt, VarIntBoundsExceeded,
+    DEFAULT_SUPPORTED_VERSIONS, Duration, MAX_CID_SIZE, NoneTokenLog, RandomConnectionIdGenerator,
+    SystemTime, TokenLog, TokenMemoryCache, TokenStore, VarInt, VarIntBoundsExceeded,
     cid_generator::{ConnectionIdGenerator, HashedConnectionIdGenerator},
     crypto::{self, HandshakeTokenKey, HmacKey},
     shared::ConnectionId,
@@ -553,7 +553,7 @@ impl ClientConfig {
         Self {
             transport: Default::default(),
             crypto,
-            token_store: Arc::new(NoneTokenStore),
+            token_store: Arc::new(TokenMemoryCache::default()),
             initial_dst_cid_provider: Arc::new(|| {
                 RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid()
             }),
@@ -585,8 +585,7 @@ impl ClientConfig {
 
     /// Set a custom [`TokenStore`]
     ///
-    /// Defaults to [`NoneTokenStore`], which disables the use of tokens from NEW_TOKEN frames as a
-    /// client.
+    /// Defaults to [`TokenMemoryCache`], which is suitable for most internet applications.
     pub fn token_store(&mut self, store: Arc<dyn TokenStore>) -> &mut Self {
         self.token_store = store;
         self
diff --git a/quinn-proto/src/lib.rs b/quinn-proto/src/lib.rs
index dbc5187d..796b24b4 100644
--- a/quinn-proto/src/lib.rs
+++ b/quinn-proto/src/lib.rs
@@ -88,6 +88,9 @@ mod token;
 use token::ResetToken;
 pub use token::{NoneTokenLog, NoneTokenStore, TokenLog, TokenReuseError, TokenStore};
 
+mod token_memory_cache;
+pub use token_memory_cache::TokenMemoryCache;
+
 #[cfg(feature = "arbitrary")]
 use arbitrary::Arbitrary;
 
diff --git a/quinn-proto/src/token_memory_cache.rs b/quinn-proto/src/token_memory_cache.rs
new file mode 100644
index 00000000..0e015873
--- /dev/null
+++ b/quinn-proto/src/token_memory_cache.rs
@@ -0,0 +1,246 @@
+//! Storing tokens sent from servers in NEW_TOKEN frames and using them in subsequent connections
+
+use std::{
+    collections::{HashMap, VecDeque, hash_map},
+    sync::{Arc, Mutex},
+};
+
+use bytes::Bytes;
+use lru_slab::LruSlab;
+use tracing::trace;
+
+use crate::token::TokenStore;
+
+/// `TokenStore` implementation that stores up to `N` tokens per server name for up to a
+/// limited number of server names, in-memory
+#[derive(Debug)]
+pub struct TokenMemoryCache(Mutex<State>);
+
+impl TokenMemoryCache {
+    /// Construct empty
+    pub fn new(max_server_names: u32, max_tokens_per_server: usize) -> Self {
+        Self(Mutex::new(State::new(
+            max_server_names,
+            max_tokens_per_server,
+        )))
+    }
+}
+
+impl TokenStore for TokenMemoryCache {
+    fn insert(&self, server_name: &str, token: Bytes) {
+        trace!(%server_name, "storing token");
+        self.0.lock().unwrap().store(server_name, token)
+    }
+
+    fn take(&self, server_name: &str) -> Option<Bytes> {
+        let token = self.0.lock().unwrap().take(server_name);
+        trace!(%server_name, found=%token.is_some(), "taking token");
+        token
+    }
+}
+
+/// Defaults to a maximum of 256 servers and 2 tokens per server
+impl Default for TokenMemoryCache {
+    fn default() -> Self {
+        Self::new(256, 2)
+    }
+}
+
+/// Lockable inner state of `TokenMemoryCache`
+#[derive(Debug)]
+struct State {
+    max_server_names: u32,
+    max_tokens_per_server: usize,
+    // map from server name to index in lru
+    lookup: HashMap<Arc<str>, u32>,
+    lru: LruSlab<CacheEntry>,
+}
+
+impl State {
+    fn new(max_server_names: u32, max_tokens_per_server: usize) -> Self {
+        Self {
+            max_server_names,
+            max_tokens_per_server,
+            lookup: HashMap::new(),
+            lru: LruSlab::default(),
+        }
+    }
+
+    fn store(&mut self, server_name: &str, token: Bytes) {
+        if self.max_server_names == 0 {
+            // the rest of this method assumes that we can always insert a new entry so long as
+            // we're willing to evict a pre-existing entry. thus, an entry limit of 0 is an edge
+            // case we must short-circuit on now.
+            return;
+        }
+        if self.max_tokens_per_server == 0 {
+            // similarly to above, the rest of this method assumes that we can always push a new
+            // token to a queue so long as we're willing to evict a pre-existing token, so we
+            // short-circuit on the edge case of a token limit of 0.
+            return;
+        }
+
+        let server_name = Arc::<str>::from(server_name);
+        match self.lookup.entry(server_name.clone()) {
+            hash_map::Entry::Occupied(hmap_entry) => {
+                // key already exists, push the new token to its token queue
+                let tokens = &mut self.lru.get_mut(*hmap_entry.get()).tokens;
+                if tokens.len() >= self.max_tokens_per_server {
+                    debug_assert!(tokens.len() == self.max_tokens_per_server);
+                    tokens.pop_front().unwrap();
+                }
+                tokens.push_back(token);
+            }
+            hash_map::Entry::Vacant(hmap_entry) => {
+                // key does not yet exist, create a new one, evicting the oldest if necessary
+                let removed_key = if self.lru.len() >= self.max_server_names {
+                    // unwrap safety: max_server_names is > 0, so there's at least one entry, so
+                    //                lru() is some
+                    Some(self.lru.remove(self.lru.lru().unwrap()).server_name)
+                } else {
+                    None
+                };
+
+                hmap_entry.insert(self.lru.insert(CacheEntry::new(server_name, token)));
+
+                // for borrowing reasons, we must defer removing the evicted hmap entry to here
+                if let Some(removed_slot) = removed_key {
+                    let removed = self.lookup.remove(&removed_slot);
+                    debug_assert!(removed.is_some());
+                }
+            }
+        };
+    }
+
+    fn take(&mut self, server_name: &str) -> Option<Bytes> {
+        let slab_key = *self.lookup.get(server_name)?;
+
+        // pop from entry's token queue
+        let entry = self.lru.get_mut(slab_key);
+        // unwrap safety: we never leave tokens empty
+        let token = entry.tokens.pop_front().unwrap();
+
+        if entry.tokens.is_empty() {
+            // token stack emptied, remove entry
+            self.lru.remove(slab_key);
+            self.lookup.remove(server_name);
+        }
+
+        Some(token)
+    }
+}
+
+/// Cache entry within `TokenMemoryCache`'s LRU slab
+#[derive(Debug)]
+struct CacheEntry {
+    server_name: Arc<str>,
+    // invariant: tokens is never empty
+    tokens: VecDeque<Bytes>,
+}
+
+impl CacheEntry {
+    /// Construct with a single token
+    fn new(server_name: Arc<str>, token: Bytes) -> Self {
+        let mut tokens = VecDeque::new();
+        tokens.push_back(token);
+        Self {
+            server_name,
+            tokens,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::VecDeque;
+
+    use super::*;
+    use rand::prelude::*;
+    use rand_pcg::Pcg32;
+
+    fn new_rng() -> impl Rng {
+        Pcg32::from_seed(0xdeadbeefdeadbeefdeadbeefdeadbeefu128.to_le_bytes())
+    }
+
+    #[test]
+    fn cache_test() {
+        let mut rng = new_rng();
+        const N: usize = 2;
+
+        for _ in 0..10 {
+            let mut cache_1: Vec<(u32, VecDeque<Bytes>)> = Vec::new(); // keep it sorted oldest to newest
+            let cache_2 = TokenMemoryCache::new(20, 2);
+
+            for i in 0..200 {
+                let server_name = rng.random::<u32>() % 10;
+                if rng.random_bool(0.666) {
+                    // store
+                    let token = Bytes::from(vec![i]);
+                    println!("STORE {} {:?}", server_name, token);
+                    if let Some((j, _)) = cache_1
+                        .iter()
+                        .enumerate()
+                        .find(|&(_, &(server_name_2, _))| server_name_2 == server_name)
+                    {
+                        let (_, mut queue) = cache_1.remove(j);
+                        queue.push_back(token.clone());
+                        if queue.len() > N {
+                            queue.pop_front();
+                        }
+                        cache_1.push((server_name, queue));
+                    } else {
+                        let mut queue = VecDeque::new();
+                        queue.push_back(token.clone());
+                        cache_1.push((server_name, queue));
+                        if cache_1.len() > 20 {
+                            cache_1.remove(0);
+                        }
+                    }
+                    cache_2.insert(&server_name.to_string(), token);
+                } else {
+                    // take
+                    println!("TAKE {}", server_name);
+                    let expecting = cache_1
+                        .iter()
+                        .enumerate()
+                        .find(|&(_, &(server_name_2, _))| server_name_2 == server_name)
+                        .map(|(j, _)| j)
+                        .map(|j| {
+                            let (_, mut queue) = cache_1.remove(j);
+                            let token = queue.pop_front().unwrap();
+                            if !queue.is_empty() {
+                                cache_1.push((server_name, queue));
+                            }
+                            token
+                        });
+                    println!("EXPECTING {:?}", expecting);
+                    assert_eq!(cache_2.take(&server_name.to_string()), expecting);
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn zero_max_server_names() {
+        // test that this edge case doesn't panic
+        let cache = TokenMemoryCache::new(0, 2);
+        for i in 0..10 {
+            cache.insert(&i.to_string(), Bytes::from(vec![i]));
+            for j in 0..10 {
+                assert!(cache.take(&j.to_string()).is_none());
+            }
+        }
+    }
+
+    #[test]
+    fn zero_queue_length() {
+        // test that this edge case doesn't panic
+        let cache = TokenMemoryCache::new(256, 0);
+        for i in 0..10 {
+            cache.insert(&i.to_string(), Bytes::from(vec![i]));
+            for j in 0..10 {
+                assert!(cache.take(&j.to_string()).is_none());
+            }
+        }
+    }
+}
diff --git a/quinn/src/lib.rs b/quinn/src/lib.rs
index a3c6f575..1c550470 100644
--- a/quinn/src/lib.rs
+++ b/quinn/src/lib.rs
@@ -62,9 +62,9 @@ pub use proto::{
     ConnectError, ConnectionClose, ConnectionError, ConnectionId, ConnectionIdGenerator,
     ConnectionStats, Dir, EcnCodepoint, EndpointConfig, FrameStats, FrameType, IdleTimeout,
     MtuDiscoveryConfig, NoneTokenLog, NoneTokenStore, PathStats, ServerConfig, Side, StdSystemTime,
-    StreamId, TimeSource, TokenLog, TokenReuseError, TokenStore, Transmit, TransportConfig,
-    TransportErrorCode, UdpStats, ValidationTokenConfig, VarInt, VarIntBoundsExceeded, Written,
-    congestion, crypto,
+    StreamId, TimeSource, TokenLog, TokenMemoryCache, TokenReuseError, TokenStore, Transmit,
+    TransportConfig, TransportErrorCode, UdpStats, ValidationTokenConfig, VarInt,
+    VarIntBoundsExceeded, Written, congestion, crypto,
 };
 #[cfg(any(feature = "rustls-aws-lc-rs", feature = "rustls-ring"))]
 pub use rustls;


Test patch:
diff --git a/quinn-proto/src/tests/util.rs b/quinn-proto/src/tests/util.rs
index 2e934ba6..433c2820 100644
--- a/quinn-proto/src/tests/util.rs
+++ b/quinn-proto/src/tests/util.rs
@@ -615,9 +615,7 @@ fn server_crypto_inner(
 }
 
 pub(super) fn client_config() -> ClientConfig {
-    let mut config = ClientConfig::new(Arc::new(client_crypto()));
-    config.token_store(Arc::new(SimpleTokenStore::default()));
-    config
+    ClientConfig::new(Arc::new(client_crypto()))
 }
 
 pub(super) fn client_config_with_deterministic_pns() -> ClientConfig {
@@ -743,25 +741,3 @@ impl TokenLog for SimpleTokenLog {
         }
     }
 }
-
-#[derive(Default)]
-struct SimpleTokenStore(Mutex<HashMap<String, VecDeque<Bytes>>>);
-
-impl TokenStore for SimpleTokenStore {
-    fn insert(&self, server_name: &str, token: Bytes) {
-        self.0
-            .lock()
-            .unwrap()
-            .entry(server_name.into())
-            .or_default()
-            .push_back(token);
-    }
-
-    fn take(&self, server_name: &str) -> Option<Bytes> {
-        self.0
-            .lock()
-            .unwrap()
-            .get_mut(server_name)
-            .and_then(|queue| queue.pop_front())
-    }
-}