use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::{Duration, Instant},
};
use bytes::Bytes;
use scion_quic::h3::client::{H3ResponseBody, Http3Client, RequestError as H3RequestError};
use crate::{
config::Config,
epoch::Epoch,
error::{Error, TimeoutPhase},
origin::Origin,
request::{IntoUrl, Request},
response::Response,
};
enum EpochSlot {
Init,
Active(Arc<Epoch>, u64),
Closed,
}
pub struct Client {
epoch: tokio::sync::RwLock<EpochSlot>,
reset_generation: AtomicU64,
config: Config,
#[cfg(test)]
epoch_factory: Option<EpochFactory>,
}
#[cfg(test)]
type EpochFactory =
Box<dyn Fn() -> futures::future::BoxFuture<'static, Result<Epoch, Error>> + Send + Sync>;
impl Client {
#[must_use]
pub fn new(config: Config) -> Self {
Client {
epoch: tokio::sync::RwLock::new(EpochSlot::Init),
reset_generation: AtomicU64::new(0),
config,
#[cfg(test)]
epoch_factory: None,
}
}
#[cfg(test)]
pub(crate) fn with_epoch_factory(config: Config, factory: EpochFactory) -> Self {
let mut client = Client::new(config);
client.epoch_factory = Some(factory);
client
}
pub async fn request(&self, request: Request) -> Result<Response, Error> {
let timeout = request
.request_timeout()
.unwrap_or(self.config.request_timeout);
let now = tokio::time::Instant::now();
let deadline = now
.checked_add(timeout)
.unwrap_or_else(|| now + Duration::from_secs(86400 * 365 * 30));
let (response, connection) = tokio::time::timeout_at(deadline, self.request_head(&request))
.await
.map_err(|_| {
Error::Timeout {
phase: TimeoutPhase::Request,
timeout,
}
})??;
Ok(Response::new(response, deadline, timeout, connection))
}
pub async fn get(&self, url: impl IntoUrl) -> Result<Response, Error> {
let request = Request::get(url).build()?;
self.request(request).await
}
pub async fn post(&self, url: impl IntoUrl, body: impl Into<Bytes>) -> Result<Response, Error> {
let request = Request::post(url).body(body).build()?;
self.request(request).await
}
pub async fn warm_up(&self, url: impl IntoUrl) -> Result<(), Error> {
let request = Request::get(url).build()?;
let origin = Origin::from_request(&request)?;
let epoch = self.current_epoch().await?;
let origin_client = epoch.origin_client(origin, Instant::now(), &self.config)?;
origin_client.connection().await?;
Ok(())
}
pub fn reset(&self) {
self.reset_generation.fetch_add(1, Ordering::Release);
}
pub async fn close(&self) {
let previous = {
let mut slot = self.epoch.write().await;
std::mem::replace(&mut *slot, EpochSlot::Closed)
};
if let EpochSlot::Active(epoch, _) = previous {
epoch.shutdown().await;
}
}
async fn request_head(
&self,
request: &Request,
) -> Result<(http::Response<H3ResponseBody>, Arc<Http3Client>), Error> {
let origin = Origin::from_request(request)?;
let (host, port) = (origin.host.clone(), origin.port);
let epoch = self.current_epoch().await?;
let origin_client = epoch.origin_client(origin, Instant::now(), &self.config)?;
let (connection, generation) = origin_client.connection().await?;
match connection.request(request.to_http()?).await {
Ok(response) => Ok((response, connection)),
Err(H3RequestError::Establish(_)) => {
let (connection, _) = origin_client.reconnect(generation).await?;
let response = connection
.request(request.to_http()?)
.await
.map_err(|e| Error::from_h3_request_error(&host, port, e))?;
Ok((response, connection))
}
Err(e) => Err(Error::from_h3_request_error(&host, port, e)),
}
}
async fn current_epoch(&self) -> Result<Arc<Epoch>, Error> {
{
let slot = self.epoch.read().await;
match &*slot {
EpochSlot::Closed => return Err(Error::Closed),
EpochSlot::Active(epoch, built)
if *built == self.reset_generation.load(Ordering::Acquire) =>
{
return Ok(epoch.clone());
}
_ => {}
}
}
let mut slot = self.epoch.write().await;
match &*slot {
EpochSlot::Closed => Err(Error::Closed),
EpochSlot::Active(epoch, built)
if *built == self.reset_generation.load(Ordering::Acquire) =>
{
Ok(epoch.clone())
}
_ => {
let generation = self.reset_generation.load(Ordering::Acquire);
let fresh = Arc::new(self.build_epoch().await?);
let previous =
std::mem::replace(&mut *slot, EpochSlot::Active(fresh.clone(), generation));
if let EpochSlot::Active(previous, _) = previous {
tracing::debug!("Replacing SCION connectivity epoch");
tokio::spawn(previous.shutdown());
}
Ok(fresh)
}
}
}
async fn build_epoch(&self) -> Result<Epoch, Error> {
#[cfg(test)]
if let Some(factory) = &self.epoch_factory {
return factory().await;
}
Epoch::build(&self.config).await
}
}
#[cfg(test)]
mod tests {
use std::{
sync::atomic::{AtomicUsize, Ordering::SeqCst},
time::Duration,
};
use scion_quic::quic::config::QuicConfig;
use test_log::test;
use super::*;
use crate::test_support::{
FailingResolver, SERVER_PORT, StaticResolver, TestServerHarness, server_scion_ip,
test_config, test_router,
};
fn harness_client(
harness: Arc<TestServerHarness>,
resolver: Arc<dyn scion_stack::resolver::ScionDnsResolver>,
config: Config,
) -> (Client, Arc<AtomicUsize>) {
let builds = Arc::new(AtomicUsize::new(0));
let factory_builds = builds.clone();
let factory_config = config.clone();
let client = Client::with_epoch_factory(
config,
Box::new(move || {
factory_builds.fetch_add(1, SeqCst);
let harness = harness.clone();
let resolver = resolver.clone();
let config = factory_config.clone();
Box::pin(async move { Ok(Epoch::new(harness, resolver, &config)) })
}),
);
(client, builds)
}
fn url(host: &str, path: &str) -> String {
format!("https://{host}:{SERVER_PORT}{path}")
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn get_resolves_via_dns_path() {
let (router, hits) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, _) = harness_client(harness, resolver, test_config());
let response = client.get(url("localhost", "/hello")).await.unwrap();
assert!(response.is_success());
let (body, _) = response.text(Some(1024)).await.unwrap();
assert_eq!(body, "world");
assert_eq!(hits.load(SeqCst), 1);
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn duration_max_timeout_does_not_panic() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, _) = harness_client(harness, resolver, test_config());
let request = Request::get(url("localhost", "/hello"))
.request_timeout(Duration::MAX)
.build()
.unwrap();
let response = client.request(request).await.unwrap();
assert!(response.is_success());
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn target_bypasses_the_resolver() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let (client, _) = harness_client(harness, Arc::new(FailingResolver), test_config());
let request = Request::post(url("localhost", "/echo"))
.body("ping")
.target(server_scion_ip())
.build()
.unwrap();
let response = client.request(request).await.unwrap();
let (body, _) = response.text(Some(1024)).await.unwrap();
assert_eq!(body, "ping");
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn resolution_failure_surfaces_as_retryable_resolution_error() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let (client, _) = harness_client(harness, Arc::new(FailingResolver), test_config());
let err = client.get(url("localhost", "/hello")).await.err().unwrap();
assert!(matches!(err, Error::Resolution { .. }), "{err}");
assert!(err.is_retryable());
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn establish_failure_reconnects_once_and_retries() {
let (router, hits) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, _) = harness_client(harness.clone(), resolver, test_config());
let response = client.get(url("localhost", "/hello")).await.unwrap();
assert!(response.is_success());
assert_eq!(harness.binds.load(SeqCst), 1);
harness.break_sockets();
tokio::time::sleep(Duration::from_millis(200)).await;
let response = client.get(url("localhost", "/hello")).await.unwrap();
assert!(response.is_success());
assert_eq!(harness.binds.load(SeqCst), 2);
assert_eq!(hits.load(SeqCst), 2);
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn stream_level_errors_are_not_retried() {
let (router, _) = test_router();
let mut server_quic = QuicConfig::builder().verify_peer(false).build();
server_quic.initial_max_streams_bidi = 0;
let harness = TestServerHarness::with_server_quic(router, server_quic);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, _) = harness_client(harness.clone(), resolver, test_config());
let err = client.get(url("localhost", "/hello")).await.err().unwrap();
assert!(matches!(err, Error::ConnectionLimit), "{err}");
assert_eq!(harness.binds.load(SeqCst), 1);
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn reset_is_lazy_and_rebuild_is_single_flight() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, builds) = harness_client(harness, resolver, test_config());
let client = Arc::new(client);
client.get(url("localhost", "/hello")).await.unwrap();
assert_eq!(builds.load(SeqCst), 1);
client.reset();
assert_eq!(builds.load(SeqCst), 1);
let tasks: Vec<_> = (0..3)
.map(|_| {
let client = client.clone();
tokio::spawn(async move { client.get(url("localhost", "/hello")).await })
})
.collect();
for task in tasks {
task.await.unwrap().unwrap();
}
assert_eq!(builds.load(SeqCst), 2);
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn reset_during_rebuild_triggers_another_rebuild() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let builds = Arc::new(AtomicUsize::new(0));
let gate = Arc::new(tokio::sync::Semaphore::new(1));
let factory_builds = builds.clone();
let factory_gate = gate.clone();
let config = test_config();
let factory_config = config.clone();
let client = Arc::new(Client::with_epoch_factory(
config,
Box::new(move || {
let builds = factory_builds.clone();
let gate = factory_gate.clone();
let harness = harness.clone();
let resolver = resolver.clone();
let config = factory_config.clone();
Box::pin(async move {
gate.acquire().await.expect("gate closed").forget();
builds.fetch_add(1, SeqCst);
Ok(Epoch::new(harness, resolver, &config))
})
}),
));
client.get(url("localhost", "/hello")).await.unwrap();
assert_eq!(builds.load(SeqCst), 1);
client.reset();
let pending = {
let client = client.clone();
tokio::spawn(async move { client.get(url("localhost", "/hello")).await })
};
tokio::time::sleep(Duration::from_millis(100)).await;
assert_eq!(builds.load(SeqCst), 1);
client.reset();
gate.add_permits(2);
pending.await.unwrap().unwrap();
assert_eq!(builds.load(SeqCst), 2);
client.get(url("localhost", "/hello")).await.unwrap();
assert_eq!(builds.load(SeqCst), 3);
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn failed_rebuild_keeps_stale_and_retries_next_request() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let builds = Arc::new(AtomicUsize::new(0));
let factory_builds = builds.clone();
let config = test_config();
let factory_config = config.clone();
let client = Client::with_epoch_factory(
config,
Box::new(move || {
let build = factory_builds.fetch_add(1, SeqCst) + 1;
let harness = harness.clone();
let resolver = resolver.clone();
let config = factory_config.clone();
Box::pin(async move {
if build == 2 {
return Err(Error::StackBuild {
retryable: true,
source: "flaky discovery".into(),
});
}
Ok(Epoch::new(harness, resolver, &config))
})
}),
);
client.get(url("localhost", "/hello")).await.unwrap();
client.reset();
let err = client.get(url("localhost", "/hello")).await.err().unwrap();
assert!(matches!(err, Error::StackBuild { .. }), "{err}");
assert!(err.is_retryable());
client.get(url("localhost", "/hello")).await.unwrap();
assert_eq!(builds.load(SeqCst), 3);
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn epochs_do_not_share_stacks_across_reset() {
let (router_one, _) = test_router();
let (router_two, _) = test_router();
let first = TestServerHarness::new(router_one);
let second = TestServerHarness::new(router_two);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let builds = Arc::new(AtomicUsize::new(0));
let factory_builds = builds.clone();
let config = test_config();
let factory_config = config.clone();
let factory_first = first.clone();
let factory_second = second.clone();
let client = Client::with_epoch_factory(
config,
Box::new(move || {
let build = factory_builds.fetch_add(1, SeqCst) + 1;
let harness = if build == 1 {
factory_first.clone()
} else {
factory_second.clone()
};
let resolver = resolver.clone();
let config = factory_config.clone();
Box::pin(async move { Ok(Epoch::new(harness, resolver, &config)) })
}),
);
client.get(url("localhost", "/hello")).await.unwrap();
assert_eq!(
(first.binds.load(SeqCst), second.binds.load(SeqCst)),
(1, 0)
);
client.reset();
client.get(url("localhost", "/hello")).await.unwrap();
assert_eq!(
(first.binds.load(SeqCst), second.binds.load(SeqCst)),
(1, 1)
);
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn stalled_establishment_does_not_block_other_origins() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
harness.stall_first_bind();
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, _) = harness_client(harness.clone(), resolver, test_config());
let client = Arc::new(client);
let stalled = {
let client = client.clone();
tokio::spawn(async move { client.get(url("a.local", "/hello")).await })
};
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(!stalled.is_finished());
let response = client.get(url("b.local", "/hello")).await.unwrap();
assert!(response.is_success());
assert!(!stalled.is_finished());
harness.release_stalled();
stalled.await.unwrap().unwrap();
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn evicted_origin_completes_in_flight_request() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, _) = harness_client(harness, resolver, test_config().with_max_origins(1));
let client = Arc::new(client);
let slow = {
let client = client.clone();
tokio::spawn(async move {
let response = client.get(url("a.local", "/slow")).await?;
response.text(Some(1024)).await
})
};
tokio::time::sleep(Duration::from_millis(150)).await;
client.get(url("b.local", "/hello")).await.unwrap();
let (body, _) = slow.await.unwrap().unwrap();
assert_eq!(body, "slow");
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn close_rejects_new_requests_and_faults_in_flight() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, _) = harness_client(harness, resolver, test_config());
let client = Arc::new(client);
client.get(url("localhost", "/hello")).await.unwrap();
let in_flight = {
let client = client.clone();
tokio::spawn(async move {
let response = client.get(url("localhost", "/slow")).await?;
response.text(Some(1024)).await
})
};
tokio::time::sleep(Duration::from_millis(100)).await;
client.close().await;
let start = tokio::time::Instant::now();
let result = in_flight.await.unwrap();
assert!(result.is_err());
assert!(start.elapsed() < Duration::from_millis(300));
let err = client.get(url("localhost", "/hello")).await.err().unwrap();
assert!(matches!(err, Error::Closed), "{err}");
client.close().await;
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn per_request_timeout_overrides_the_default() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, _) = harness_client(harness, resolver, test_config());
let request = Request::get(url("localhost", "/slow"))
.request_timeout(Duration::from_millis(100))
.build()
.unwrap();
let err = client.request(request).await.err().unwrap();
assert!(
matches!(
err,
Error::Timeout {
phase: TimeoutPhase::Request,
..
}
),
"{err}"
);
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn body_collection_respects_max_size() {
let (router, _) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, _) = harness_client(harness, resolver, test_config());
let response = client.get(url("localhost", "/hello")).await.unwrap();
let err = response.bytes(Some(2)).await.err().unwrap();
assert!(matches!(err, Error::BodyTooLarge { limit: 2 }), "{err}");
}
#[test(tokio::test)]
#[ntest::timeout(10_000)]
async fn warm_up_establishes_without_a_request() {
let (router, hits) = test_router();
let harness = TestServerHarness::new(router);
let resolver = StaticResolver::new(vec![server_scion_ip()]);
let (client, builds) = harness_client(harness.clone(), resolver, test_config());
client.warm_up(url("localhost", "/")).await.unwrap();
assert_eq!(builds.load(SeqCst), 1);
assert_eq!(harness.binds.load(SeqCst), 1);
assert_eq!(hits.load(SeqCst), 0);
client.get(url("localhost", "/hello")).await.unwrap();
assert_eq!(harness.binds.load(SeqCst), 1);
assert_eq!(hits.load(SeqCst), 1);
}
}