Skip to content

Commit

Permalink
feat(iroh-relay)!: use explicit key cache (#3053)
Browse files Browse the repository at this point in the history
## 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 <[email protected]>
rklaehn and dignifiedquire authored Dec 16, 2024
1 parent ac74c53 commit d4f72fa
Showing 18 changed files with 211 additions and 47 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 7 additions & 0 deletions iroh-base/src/key.rs
Original file line number Diff line number Diff line change
@@ -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))
1 change: 1 addition & 0 deletions iroh-relay/Cargo.toml
Original file line number Diff line number Diff line change
@@ -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"] }
22 changes: 19 additions & 3 deletions iroh-relay/src/client.rs
Original file line number Diff line number Diff line change
@@ -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))
}
7 changes: 4 additions & 3 deletions iroh-relay/src/client/conn.rs
Original file line number Diff line number Diff line change
@@ -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.");
7 changes: 7 additions & 0 deletions iroh-relay/src/defaults.rs
Original file line number Diff line number Diff line change
@@ -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;
63 changes: 63 additions & 0 deletions iroh-relay/src/key_cache.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}
2 changes: 2 additions & 0 deletions iroh-relay/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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;
4 changes: 4 additions & 0 deletions iroh-relay/src/main.rs
Original file line number Diff line number Diff line change
@@ -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(),
Loading

0 comments on commit d4f72fa

Please sign in to comment.