use std::collections::HashSet;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use bytes::Bytes;
use http_body_util::{Empty, Full};
use hyper::service::service_fn;
use hyper::{Request, Response, StatusCode};
use hyper_util::client::legacy::Client;
use hyper_util::client::legacy::connect::HttpConnector;
use hyper_util::rt::{TokioExecutor, TokioIo};
use tokio::net::TcpListener;
use plecto_control::{Control, Host, Manifest, MemoryStore};
use plecto_host::test_support::TestSigner;
use plecto_server::serve;
async fn spawn_labeled_upstream(label: &'static str) -> SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
loop {
let (stream, _) = listener.accept().await.unwrap();
tokio::spawn(async move {
let _ = hyper::server::conn::http1::Builder::new()
.serve_connection(
TokioIo::new(stream),
service_fn(move |_req| async move {
Ok::<Response<Full<Bytes>>, Infallible>(
Response::builder()
.status(200)
.header("x-instance", label)
.body(Full::new(Bytes::from_static(b"ok")))
.unwrap(),
)
}),
)
.await;
});
}
});
addr
}
async fn dead_addr() -> SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
drop(listener);
addr
}
async fn spawn_proxy(control: Arc<Control>) -> SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let _ = serve(control, listener).await;
});
addr
}
fn control_for(addresses: &[SocketAddr]) -> Arc<Control> {
let signer = TestSigner::new().unwrap();
let addrs = addresses
.iter()
.map(|a| format!("\"{a}\""))
.collect::<Vec<_>>()
.join(", ");
let toml = format!(
r#"
[[upstream]]
name = "pool"
addresses = [{addrs}]
[upstream.health]
path = "/healthz"
interval_ms = 50
timeout_ms = 100
[[route]]
upstream = "pool"
[route.match]
path_prefix = "/api"
"#
);
let manifest = Manifest::from_toml(&toml).unwrap();
let host = Host::new(signer.trust_policy().unwrap()).unwrap();
Arc::new(Control::load(host, &manifest, Box::new(MemoryStore::new())).unwrap())
}
fn client() -> Client<HttpConnector, Empty<Bytes>> {
Client::builder(TokioExecutor::new()).build_http()
}
async fn get(
client: &Client<HttpConnector, Empty<Bytes>>,
proxy: SocketAddr,
) -> (StatusCode, Option<String>) {
let resp = client
.request(
Request::builder()
.method("GET")
.uri(format!("http://{proxy}/api/x"))
.body(Empty::<Bytes>::new())
.unwrap(),
)
.await
.expect("proxy request");
let instance = resp
.headers()
.get("x-instance")
.and_then(|v| v.to_str().ok())
.map(str::to_string);
(resp.status(), instance)
}
#[tokio::test]
async fn round_robin_distributes_over_healthy_instances() {
let a = spawn_labeled_upstream("a").await;
let b = spawn_labeled_upstream("b").await;
let proxy = spawn_proxy(control_for(&[a, b])).await;
let client = client();
let mut seen = HashSet::new();
for _ in 0..200 {
let (status, instance) = get(&client, proxy).await;
if status == StatusCode::OK
&& let Some(i) = instance
{
seen.insert(i);
}
if seen.len() >= 2 {
break;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
assert_eq!(
seen,
HashSet::from(["a".to_string(), "b".to_string()]),
"round-robin must distribute across both healthy instances"
);
}
#[tokio::test]
async fn dead_instance_is_ejected_and_traffic_goes_to_the_healthy_one() {
let live = spawn_labeled_upstream("live").await;
let dead = dead_addr().await;
let proxy = spawn_proxy(control_for(&[live, dead])).await;
let client = client();
for _ in 0..200 {
if get(&client, proxy).await.0 != StatusCode::SERVICE_UNAVAILABLE {
break;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
for _ in 0..20 {
let (status, instance) = get(&client, proxy).await;
assert_eq!(
status,
StatusCode::OK,
"all traffic forwards to the healthy instance"
);
assert_eq!(
instance.as_deref(),
Some("live"),
"the dead instance is ejected"
);
}
}
#[tokio::test]
async fn all_unhealthy_fails_closed_503() {
let proxy = spawn_proxy(control_for(&[dead_addr().await, dead_addr().await])).await;
let client = client();
tokio::time::sleep(Duration::from_millis(150)).await;
let resp = client
.request(
Request::builder()
.method("GET")
.uri(format!("http://{proxy}/api/x"))
.body(Empty::<Bytes>::new())
.unwrap(),
)
.await
.expect("proxy request");
assert_eq!(
resp.status(),
StatusCode::SERVICE_UNAVAILABLE,
"no healthy instance → fail closed with 503"
);
assert_eq!(
resp.headers()
.get("x-plecto-fault")
.and_then(|v| v.to_str().ok()),
Some("no-healthy-upstream"),
"the 503 carries the no-healthy-upstream fault marker"
);
}