feat(iroh-relay)!: use explicit key cache (#3053)

## Description

This wires up an explicit key cache to replace the implicit one that was
removed in #3051.

~~The default for a key cache is Disabled. A disabled key cache has a
size of 1 pointer and otherwise zero performance overhead.~~

I have removed the Default instance for both KeyCache and DerpProtocol
so you don't accidentally pass the default despite having a cache
available.

We use the lru crate for the cache for now. Please comment if it should
be something else.

Benchmarks have shown that conversion from a [u8;32] to a VerifyingKey
is relatively cheap, so the purpose of the cache is solely to validate
incoming public keys.

We add a Borrow instance to PublicKey so we can use it as a cache key.

Some performance measurements:

```
Benchmarking validate valid ed25519 key: 3.7674 µs
Benchmarking validate invalid ed25519 key: 3.6637 µs
Benchmarking validate valid iroh ed25519 key: 67.089 ns
Benchmarking validate invalid iroh ed25519 key: 64.004 ns
```

So just from validating incoming keys, without cache you would be
limited to ~250 000 msgs/s per thread. At a message size of 1KiB that
would be 250MB/s, which is not great.

With the cache deserialization can do 14 000 000 msgs/s, which means
that this is no longer a bottleneck.

## Breaking Changes

- RelayConfig has new public field key_cache_capacity

## Notes & open questions

- Size of the cache
- source_and_box has a PublicKey::try_from. Is that perf critical?

## Change checklist

- [ ] Self-review.
- [ ] Documentation updates following the [style
guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text),
if relevant.
- [ ] Tests if relevant.
- [ ] All breaking changes documented.

---------

Co-authored-by: dignifiedquire <me@dignifiedquire.com>

Author: rklaehn

Cargo.lock

@@ -1,6 +1,7 @@
 //! Cryptographic key handling for `iroh`.
 
 use std::{
+    borrow::Borrow,
     cmp::{Ord, PartialOrd},
     fmt::{Debug, Display},
     hash::Hash,
@@ -21,6 +22,12 @@ use serde::{Deserialize, Serialize};
 #[repr(transparent)]
 pub struct PublicKey(CompressedEdwardsY);
 
+impl Borrow<[u8; 32]> for PublicKey {
+    fn borrow(&self) -> &[u8; 32] {
+        self.as_bytes()
+    }
+}
+
 impl PartialOrd for PublicKey {
     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
         Some(self.cmp(other))

iroh-base/src/key.rs

@@ -96,6 +96,7 @@ url = { version = "2.5", features = ["serde"] }
 webpki = { package = "rustls-webpki", version = "0.102" }
 webpki-roots = "0.26"
 data-encoding = "2.6.0"
+lru = "0.12"
 
 [dev-dependencies]
 clap = { version = "4", features = ["derive"] }

iroh-relay/Cargo.toml

@@ -45,6 +45,7 @@ use crate::{
     defaults::timeouts::*,
     http::{Protocol, RELAY_PATH},
     protos::relay::DerpCodec,
+    KeyCache,
 };
 
 pub(crate) mod conn;
@@ -159,6 +160,7 @@ struct Actor {
     ping_tasks: JoinSet<()>,
     dns_resolver: DnsResolver,
     proxy_url: Option<Url>,
+    key_cache: KeyCache,
 }
 
 #[derive(Default, Debug)]
@@ -208,6 +210,8 @@ pub struct ClientBuilder {
     insecure_skip_cert_verify: bool,
     /// HTTP Proxy
     proxy_url: Option<Url>,
+    /// Capacity of the key cache
+    key_cache_capacity: usize,
 }
 
 impl ClientBuilder {
@@ -223,6 +227,7 @@ impl ClientBuilder {
             #[cfg(any(test, feature = "test-utils"))]
             insecure_skip_cert_verify: false,
             proxy_url: None,
+            key_cache_capacity: 128,
         }
     }
 
@@ -281,6 +286,12 @@ impl ClientBuilder {
         self
     }
 
+    /// Set the capacity of the cache for public keys.
+    pub fn key_cache_capacity(mut self, capacity: usize) -> Self {
+        self.key_cache_capacity = capacity;
+        self
+    }
+
     /// Build the [`Client`]
     pub fn build(self, key: SecretKey, dns_resolver: DnsResolver) -> (Client, ClientReceiver) {
         // TODO: review TLS config
@@ -320,6 +331,7 @@ impl ClientBuilder {
             tls_connector,
             dns_resolver,
             proxy_url: self.proxy_url,
+            key_cache: KeyCache::new(self.key_cache_capacity),
         };
 
         let (msg_sender, inbox) = mpsc::channel(64);
@@ -629,7 +641,9 @@ impl Actor {
 
         let (writer, reader) = tokio_tungstenite_wasm::connect(dial_url).await?.split();
 
-        let reader = ConnReader::Ws(reader);
+        let cache = self.key_cache.clone();
+
+        let reader = ConnReader::Ws(reader, cache);
         let writer = ConnWriter::Ws(writer);
 
         Ok((reader, writer))
@@ -683,8 +697,10 @@ impl Actor {
         let (reader, writer) =
             downcast_upgrade(upgraded).map_err(|e| ClientError::Upgrade(e.to_string()))?;
 
-        let reader = ConnReader::Derp(FramedRead::new(reader, DerpCodec));
-        let writer = ConnWriter::Derp(FramedWrite::new(writer, DerpCodec));
+        let cache = self.key_cache.clone();
+
+        let reader = ConnReader::Derp(FramedRead::new(reader, DerpCodec::new(cache.clone())));
+        let writer = ConnWriter::Derp(FramedWrite::new(writer, DerpCodec::new(cache)));
 
         Ok((reader, writer, local_addr))
     }

iroh-relay/src/client.rs

@@ -27,6 +27,7 @@ use tokio_util::{
 };
 use tracing::{debug, info_span, trace, Instrument};
 
+use super::KeyCache;
 use crate::{
     client::streams::{MaybeTlsStreamReader, MaybeTlsStreamWriter},
     defaults::timeouts::CLIENT_RECV_TIMEOUT,
@@ -270,7 +271,7 @@ pub struct ConnBuilder {
 
 pub(crate) enum ConnReader {
     Derp(FramedRead<MaybeTlsStreamReader, DerpCodec>),
-    Ws(SplitStream<WebSocketStream>),
+    Ws(SplitStream<WebSocketStream>, KeyCache),
 }
 
 pub(crate) enum ConnWriter {
@@ -291,9 +292,9 @@ impl Stream for ConnReader {
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         match *self {
             Self::Derp(ref mut ws) => Pin::new(ws).poll_next(cx),
-            Self::Ws(ref mut ws) => match Pin::new(ws).poll_next(cx) {
+            Self::Ws(ref mut ws, ref cache) => match Pin::new(ws).poll_next(cx) {
                 Poll::Ready(Some(Ok(tokio_tungstenite_wasm::Message::Binary(vec)))) => {
-                    Poll::Ready(Some(Frame::decode_from_ws_msg(vec)))
+                    Poll::Ready(Some(Frame::decode_from_ws_msg(vec, cache)))
                 }
                 Poll::Ready(Some(Ok(msg))) => {
                     tracing::warn!(?msg, "Got websocket message of unsupported type, skipping.");

iroh-relay/src/client/conn.rs

@@ -20,6 +20,13 @@ pub const DEFAULT_HTTPS_PORT: u16 = 443;
 /// The default metrics port used by the Relay server.
 pub const DEFAULT_METRICS_PORT: u16 = 9090;
 
+/// The default capacity of the key cache for the relay server.
+///
+/// Sized for 1 million concurrent clients.
+/// memory usage will be (32 + 8 + 8 + 8) * 1_000_000 = 56MB on 64 bit,
+/// which seems reasonable for a server.
+pub const DEFAULT_KEY_CACHE_CAPACITY: usize = 1024 * 1024;
+
 /// Contains all timeouts that we use in `iroh`.
 pub(crate) mod timeouts {
     use std::time::Duration;

iroh-relay/src/defaults.rs

@@ -0,0 +1,63 @@
+use std::{
+    num::NonZeroUsize,
+    sync::{Arc, Mutex},
+};
+
+use iroh_base::PublicKey;
+
+type SignatureError = <PublicKey as TryFrom<&'static [u8]>>::Error;
+type PublicKeyBytes = [u8; PublicKey::LENGTH];
+
+/// A cache for public keys.
+///
+/// This is used solely to make parsing public keys from byte slices more
+/// efficient for the very common case where a large number of identical keys
+/// are being parsed, like in the relay server.
+///
+/// The cache stores only successful parse results.
+#[derive(Debug, Clone)]
+pub enum KeyCache {
+    /// The key cache is disabled.
+    Disabled,
+    /// The key cache is enabled with a fixed capacity. It is shared between
+    /// multiple threads.
+    Shared(Arc<Mutex<lru::LruCache<PublicKey, ()>>>),
+}
+
+impl KeyCache {
+    /// Key cache to be used in tests.
+    #[cfg(test)]
+    pub fn test() -> Self {
+        Self::Disabled
+    }
+
+    /// Create a new key cache with the given capacity.
+    ///
+    /// If the capacity is zero, the cache is disabled and has zero overhead.
+    pub fn new(capacity: usize) -> Self {
+        let Some(capacity) = NonZeroUsize::new(capacity) else {
+            return Self::Disabled;
+        };
+        let cache = lru::LruCache::new(capacity);
+        Self::Shared(Arc::new(Mutex::new(cache)))
+    }
+
+    /// Get a key from a slice of bytes.
+    pub fn key_from_slice(&self, slice: &[u8]) -> Result<PublicKey, SignatureError> {
+        let Self::Shared(cache) = self else {
+            return PublicKey::try_from(slice);
+        };
+        let Ok(bytes) = PublicKeyBytes::try_from(slice) else {
+            // if the size is wrong, use PublicKey::try_from to fail with a
+            // SignatureError.
+            return Err(PublicKey::try_from(slice).expect_err("invalid length"));
+        };
+        let mut cache = cache.lock().expect("not poisoned");
+        if let Some((key, _)) = cache.get_key_value(&bytes) {
+            return Ok(*key);
+        }
+        let key = PublicKey::from_bytes(&bytes)?;
+        cache.put(key, ());
+        Ok(key)
+    }
+}

iroh-relay/src/key_cache.rs

@@ -38,7 +38,9 @@ pub mod quic;
 #[cfg(feature = "server")]
 pub mod server;
 
+mod key_cache;
 mod relay_map;
+pub(crate) use key_cache::KeyCache;
 
 #[cfg(test)]
 mod dns;

iroh-relay/src/lib.rs

@@ -168,6 +168,8 @@ struct Config {
     /// Defaults to `http_bind_addr` with the port set to [`DEFAULT_METRICS_PORT`]
     /// (`[::]:9090` when `http_bind_addr` is set to the default).
     metrics_bind_addr: Option<SocketAddr>,
+    /// The capacity of the key cache.
+    key_cache_capacity: Option<usize>,
 }
 
 impl Config {
@@ -199,6 +201,7 @@ impl Default for Config {
             limits: None,
             enable_metrics: cfg_defaults::enable_metrics(),
             metrics_bind_addr: None,
+            key_cache_capacity: Default::default(),
         }
     }
 }
@@ -570,6 +573,7 @@ async fn build_relay_config(cfg: Config) -> Result<relay::ServerConfig<std::io::
         // if `dangerous_http_only` is set, do not pass in any tls configuration
         tls: relay_tls.and_then(|tls| if dangerous_http_only { None } else { Some(tls) }),
         limits,
+        key_cache_capacity: cfg.key_cache_capacity,
     };
     let stun_config = relay::StunConfig {
         bind_addr: cfg.stun_bind_addr(),