From 1f1e2d547cceda68da71855186abc38a3a1ba713 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Thu, 31 Oct 2024 11:49:00 +0000
Subject: [PATCH] optimize override ips; utilize all ips from cache

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/service/resolver/actual.rs |  6 +++---
 src/service/resolver/cache.rs  | 10 +++++++---
 src/service/resolver/dns.rs    | 12 +++++-------
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/src/service/resolver/actual.rs b/src/service/resolver/actual.rs
index 61eedca5..5dc03d14 100644
--- a/src/service/resolver/actual.rs
+++ b/src/service/resolver/actual.rs
@@ -10,7 +10,7 @@ use ipaddress::IPAddress;
 use ruma::ServerName;
 
 use super::{
-	cache::{CachedDest, CachedOverride},
+	cache::{CachedDest, CachedOverride, MAX_IPS},
 	fed::{add_port_to_hostname, get_ip_with_port, FedDest, PortString},
 };
 
@@ -266,9 +266,9 @@ impl super::Service {
 				}
 
 				self.set_cached_override(
-					overname.to_owned(),
+					overname,
 					CachedOverride {
-						ips: override_ip.iter().collect(),
+						ips: override_ip.into_iter().take(MAX_IPS).collect(),
 						port,
 						expire: CachedOverride::default_expire(),
 					},
diff --git a/src/service/resolver/cache.rs b/src/service/resolver/cache.rs
index 465b5985..a13399dc 100644
--- a/src/service/resolver/cache.rs
+++ b/src/service/resolver/cache.rs
@@ -5,6 +5,7 @@ use std::{
 	time::SystemTime,
 };
 
+use arrayvec::ArrayVec;
 use conduit::{trace, utils::rand};
 use ruma::{OwnedServerName, ServerName};
 
@@ -24,7 +25,7 @@ pub struct CachedDest {
 
 #[derive(Clone, Debug)]
 pub struct CachedOverride {
-	pub ips: Vec<IpAddr>,
+	pub ips: IpAddrs,
 	pub port: u16,
 	pub expire: SystemTime,
 }
@@ -32,6 +33,9 @@ pub struct CachedOverride {
 pub type WellKnownMap = HashMap<OwnedServerName, CachedDest>;
 pub type TlsNameMap = HashMap<String, CachedOverride>;
 
+pub type IpAddrs = ArrayVec<IpAddr, MAX_IPS>;
+pub(crate) const MAX_IPS: usize = 3;
+
 impl Cache {
 	pub(super) fn new() -> Arc<Self> {
 		Arc::new(Self {
@@ -61,13 +65,13 @@ impl super::Service {
 			.cloned()
 	}
 
-	pub fn set_cached_override(&self, name: String, over: CachedOverride) -> Option<CachedOverride> {
+	pub fn set_cached_override(&self, name: &str, over: CachedOverride) -> Option<CachedOverride> {
 		trace!(?name, ?over, "set cached override");
 		self.cache
 			.overrides
 			.write()
 			.expect("locked for writing")
-			.insert(name, over)
+			.insert(name.into(), over)
 	}
 
 	#[must_use]
diff --git a/src/service/resolver/dns.rs b/src/service/resolver/dns.rs
index 89129e03..d3e9f5c9 100644
--- a/src/service/resolver/dns.rs
+++ b/src/service/resolver/dns.rs
@@ -1,4 +1,4 @@
-use std::{iter, net::SocketAddr, sync::Arc, time::Duration};
+use std::{net::SocketAddr, sync::Arc, time::Duration};
 
 use conduit::{err, Result, Server};
 use futures::FutureExt;
@@ -101,14 +101,12 @@ impl Resolve for Hooked {
 }
 
 async fn cached_to_reqwest(cached: CachedOverride) -> ResolvingResult {
-	let first_ip = cached
+	let addrs = cached
 		.ips
-		.first()
-		.expect("must provide at least one override");
+		.into_iter()
+		.map(move |ip| SocketAddr::new(ip, cached.port));
 
-	let saddr = SocketAddr::new(*first_ip, cached.port);
-
-	Ok(Box::new(iter::once(saddr)))
+	Ok(Box::new(addrs))
 }
 
 async fn resolve_to_reqwest(resolver: Arc<TokioAsyncResolver>, name: Name) -> ResolvingResult {