use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use bytes::Bytes;
use futures_core::future::BoxFuture;
use http::StatusCode;
use time::OffsetDateTime;
use cratefield_core::{
Budget, BudgetRouter, BudgetedHttpClient, Database, Defer, EVENT_QUOTA_EXHAUSTED,
EVENT_QUOTA_WARNING, EventBus, HttpClient, HttpError, Idempotent, RetryPolicy, Statement,
Usage,
};
use crate::fakes::ManualClock;
#[allow(clippy::disallowed_types)]
use std::sync::Mutex as StdMutex;
#[allow(clippy::disallowed_types)]
type Log<T> = StdMutex<T>;
fn clock_start() -> OffsetDateTime {
OffsetDateTime::UNIX_EPOCH + time::Duration::seconds(1)
}
struct ScriptedClient {
sent: Log<Vec<http::Request<Bytes>>>,
script: Log<Vec<Result<http::Response<Bytes>, HttpError>>>,
}
impl ScriptedClient {
fn new() -> Self {
Self {
sent: Log::new(Vec::new()),
script: Log::new(Vec::new()),
}
}
fn push(&self, outcome: Result<http::Response<Bytes>, HttpError>) {
self.script.lock().expect("script lock").push(outcome);
}
fn sends(&self) -> usize {
self.sent.lock().expect("sent lock").len()
}
fn captured(&self) -> Vec<http::Request<Bytes>> {
self.sent.lock().expect("sent lock").clone()
}
}
#[async_trait]
impl HttpClient for ScriptedClient {
async fn send(
&self,
request: http::Request<Bytes>,
) -> Result<http::Response<Bytes>, HttpError> {
self.sent.lock().expect("sent lock").push(request);
let mut script = self.script.lock().expect("script lock");
if script.is_empty() {
Ok(http::Response::new(Bytes::new()))
} else {
script.remove(0)
}
}
}
#[derive(Clone, Default)]
struct RecordingDefer(Arc<Log<Vec<BoxFuture<'static, ()>>>>);
impl Defer for RecordingDefer {
fn wait_until(&self, fut: BoxFuture<'static, ()>) {
self.0.lock().expect("defer lock").push(fut);
}
}
impl RecordingDefer {
async fn drain(&self) {
loop {
let next = {
let mut pending = self.0.lock().expect("defer lock");
(!pending.is_empty()).then(|| pending.remove(0))
};
let Some(fut) = next else { break };
fut.await;
}
}
}
struct Fixture {
db: Arc<dyn Database>,
clock: ManualClock,
inner: Arc<ScriptedClient>,
sleeps: Arc<Log<Vec<Duration>>>,
events: Arc<Log<Vec<(String, serde_json::Value)>>>,
defer: RecordingDefer,
client: BudgetedHttpClient,
}
impl Fixture {
fn new(
db: &Arc<dyn Database>,
buckets: &str,
usage: &str,
policy: RetryPolicy,
router: BudgetRouter,
) -> Self {
let clock = ManualClock::new(clock_start());
let inner = Arc::new(ScriptedClient::new());
let sleeps: Arc<Log<Vec<Duration>>> = Arc::new(Log::new(Vec::new()));
let events: Arc<Log<Vec<(String, serde_json::Value)>>> = Arc::new(Log::new(Vec::new()));
let warnings = Arc::clone(&events);
let exhaustions = Arc::clone(&events);
let bus = EventBus::new()
.on(
EVENT_QUOTA_WARNING,
Arc::new(move |_scope, payload| {
warnings
.lock()
.expect("event record")
.push((EVENT_QUOTA_WARNING.to_owned(), payload));
Box::pin(std::future::ready(Ok(())))
}),
)
.on(
EVENT_QUOTA_EXHAUSTED,
Arc::new(move |_scope, payload| {
exhaustions
.lock()
.expect("event record")
.push((EVENT_QUOTA_EXHAUSTED.to_owned(), payload));
Box::pin(std::future::ready(Ok(())))
}),
);
let defer = RecordingDefer::default();
let sleeper_sleeps = Arc::clone(&sleeps);
let sleeper_clock = clock.clone();
let client = BudgetedHttpClient::new(
Arc::clone(&inner) as Arc<dyn HttpClient>,
Arc::new(clock.clone()),
Arc::clone(db),
router,
)
.with_retry(policy)
.with_events(bus)
.with_defer(Arc::new(defer.clone()) as Arc<dyn Defer>)
.with_bucket_table(buckets)
.with_quota_table(usage)
.with_sleeper(Arc::new(move |after| {
sleeper_sleeps.lock().expect("sleeps lock").push(after);
sleeper_clock.advance(after);
Box::pin(std::future::ready(()))
}));
Self {
db: Arc::clone(db),
clock,
inner,
sleeps,
events,
defer,
client,
}
}
fn sends(&self) -> usize {
self.inner.sends()
}
fn sleep_log(&self) -> Vec<Duration> {
self.sleeps.lock().expect("sleeps lock").clone()
}
async fn drain_events(&self) -> Vec<(String, serde_json::Value)> {
self.defer.drain().await;
std::mem::take(&mut *self.events.lock().expect("event record"))
}
async fn bucket_keys(&self, buckets: &str) -> Vec<String> {
let rows = self
.db
.query(&Statement::new(format!(
"SELECT bucket FROM {buckets} ORDER BY bucket"
)))
.await
.expect("bucket rows read");
rows.rows
.iter()
.filter_map(|row| row.get::<String>("bucket"))
.collect()
}
}
fn fixed_router(budget: &Budget) -> BudgetRouter {
let budget = budget.clone();
Arc::new(move |_request| Some(budget.clone()))
}
fn get_request(path: &str) -> http::Request<Bytes> {
http::Request::builder()
.method(http::Method::GET)
.uri(format!("https://upstream.example.test{path}"))
.body(Bytes::new())
.expect("test request")
}
fn post_request(path: &str) -> http::Request<Bytes> {
http::Request::builder()
.method(http::Method::POST)
.uri(format!("https://upstream.example.test{path}"))
.body(Bytes::new())
.expect("test request")
}
fn throttled(status: StatusCode, retry_after: Option<&str>) -> http::Response<Bytes> {
let mut builder = http::Response::builder().status(status);
if let Some(value) = retry_after {
builder = builder.header(http::header::RETRY_AFTER, value);
}
builder.body(Bytes::new()).expect("test response")
}
fn ok_response(body: &[u8]) -> http::Response<Bytes> {
http::Response::builder()
.status(StatusCode::OK)
.body(Bytes::copy_from_slice(body))
.expect("test response")
}
async fn create_probe_tables(db: &Arc<dyn Database>, buckets: &str, usage: &str) {
db.execute(&Statement::new(
BudgetedHttpClient::create_bucket_table_sql(buckets),
))
.await
.expect("bucket probe table");
db.execute(&Statement::new(Usage::new(usage).create_table_sql()))
.await
.expect("quota probe table");
clear_probe_tables(db, buckets, usage).await;
}
async fn clear_probe_tables(db: &Arc<dyn Database>, buckets: &str, usage: &str) {
for table in [buckets, usage] {
db.execute(&Statement::new(format!("DELETE FROM {table}")))
.await
.expect("clear probe");
}
}
pub async fn assert_budget_pacing(db: &Arc<dyn Database>) {
let buckets = "budget_probe_pacing_buckets";
let usage = "budget_probe_pacing_usage";
create_probe_tables(db, buckets, usage).await;
let budget = |key: &str| Budget {
key: key.to_owned(),
capacity: 2.0,
refill_per_sec: 1.0,
cost: 1.0,
daily_limit: None,
};
let router: BudgetRouter = {
let (a, b) = (budget("budget-probe:a"), budget("budget-probe:b"));
Arc::new(move |request| match request.uri().path() {
"/a" => Some(a.clone()),
"/b" => Some(b.clone()),
_ => None,
})
};
let policy = RetryPolicy::default();
let fixture = Fixture::new(db, buckets, usage, policy, router);
fixture
.client
.send(get_request("/a"))
.await
.expect("the first send is within the fresh bucket");
fixture
.client
.send(get_request("/a"))
.await
.expect("the second send spends the last token");
match fixture.client.send(get_request("/a")).await {
Ok(_) => {
let waits = fixture.sleep_log();
assert!(!waits.is_empty(), "a paced send records the waits it made");
assert!(
waits.iter().all(|wait| *wait <= policy.max_delay()),
"pacing never waits past max_delay: {waits:?}"
);
}
Err(HttpError::BudgetExhausted { what, retry_after }) => {
assert_eq!(what, "budget-probe:a");
assert!(
retry_after > Duration::ZERO,
"a refusal carries the refill wait, got {retry_after:?}"
);
}
Err(other) => panic!("a dry bucket paces or refuses, nothing else: {other}"),
}
fixture.clock.advance(Duration::from_secs(1));
fixture
.client
.send(get_request("/a"))
.await
.expect("the refilled bucket carries the request");
fixture
.client
.send(get_request("/b"))
.await
.expect("a second key has its own full bucket");
assert_eq!(
fixture.bucket_keys(buckets).await,
vec!["budget-probe:a", "budget-probe:b"],
"one bucket row per key"
);
clear_probe_tables(db, buckets, usage).await;
}
pub async fn assert_budget_daily_quota(db: &Arc<dyn Database>) {
let buckets = "budget_probe_quota_buckets";
let usage = "budget_probe_quota_usage";
create_probe_tables(db, buckets, usage).await;
let budget = Budget {
key: "budget-probe:quota".to_owned(),
capacity: 100.0,
refill_per_sec: 0.0,
cost: 1.0,
daily_limit: Some(10),
};
let fixture = Fixture::new(
db,
buckets,
usage,
RetryPolicy::default(),
fixed_router(&budget),
);
check_quota_alerts(&fixture).await;
check_untracked_quota(db, buckets, usage).await;
clear_probe_tables(db, buckets, usage).await;
}
async fn check_quota_alerts(fixture: &Fixture) {
for spent in 1..=8 {
fixture
.client
.send(get_request("/quota"))
.await
.expect("within quota");
assert_eq!(
fixture.sends(),
spent,
"every send within the quota goes out"
);
}
let events = fixture.drain_events().await;
let [(name, payload)] = events.as_slice() else {
panic!("exactly one alert by the eighth send, got {events:?}");
};
assert_eq!(name, EVENT_QUOTA_WARNING, "the eighth spend is the warning");
assert_eq!(payload["subject"], "budget-probe:quota");
assert_eq!(payload["used"], 8);
assert_eq!(payload["limit"], 10);
fixture
.client
.send(get_request("/quota"))
.await
.expect("the ninth send is under the limit");
fixture
.client
.send(get_request("/quota"))
.await
.expect("the tenth send spends the day");
let events = fixture.drain_events().await;
let [(name, payload)] = events.as_slice() else {
panic!("exactly one alert for the tenth send, got {events:?}");
};
assert_eq!(
name, EVENT_QUOTA_EXHAUSTED,
"the tenth spend is the exhaustion"
);
assert_eq!(payload["used"], 10);
let err = fixture
.client
.send(get_request("/quota"))
.await
.expect_err("the day is spent");
assert!(
matches!(
&err,
HttpError::BudgetExhausted { what, retry_after }
if what == "budget-probe:quota"
&& *retry_after == Duration::from_secs(86_400 - 1)
),
"the retry_after reaches the UTC day boundary: got {err}"
);
assert_eq!(fixture.sends(), 10, "the refused request never went out");
fixture.clock.advance(Duration::from_secs(86_400 - 1));
fixture
.client
.send(get_request("/quota"))
.await
.expect("the next day's quota is fresh");
}
async fn check_untracked_quota(db: &Arc<dyn Database>, buckets: &str, usage: &str) {
let budget = Budget {
key: "budget-probe:untracked".to_owned(),
daily_limit: None,
..budget_shape()
};
let fixture = Fixture::new(
db,
buckets,
usage,
RetryPolicy::default(),
fixed_router(&budget),
);
for _ in 0..12 {
fixture
.client
.send(get_request("/quota"))
.await
.expect("an untracked budget never refuses on quota");
}
let events = fixture.drain_events().await;
assert!(events.is_empty(), "no limit, no alerts: {events:?}");
}
fn budget_shape() -> Budget {
Budget {
key: String::new(),
capacity: 100.0,
refill_per_sec: 0.0,
cost: 1.0,
daily_limit: None,
}
}
pub async fn assert_budget_retries(db: &Arc<dyn Database>) {
let buckets = "budget_probe_retry_buckets";
let usage = "budget_probe_retry_usage";
create_probe_tables(db, buckets, usage).await;
check_stated_retry_after_is_honoured(db, buckets, usage).await;
check_a_post_is_not_retried(db, buckets, usage).await;
check_the_idempotent_extension_vouches(db, buckets, usage).await;
check_transport_retries_on_the_schedule(db, buckets, usage).await;
clear_probe_tables(db, buckets, usage).await;
}
fn retry_fixture(db: &Arc<dyn Database>, buckets: &str, usage: &str, key: &str) -> Fixture {
let budget = Budget {
key: key.to_owned(),
..budget_shape()
};
let router = fixed_router(&budget);
Fixture::new(db, buckets, usage, RetryPolicy::default(), router)
}
async fn check_stated_retry_after_is_honoured(db: &Arc<dyn Database>, buckets: &str, usage: &str) {
let fixture = retry_fixture(db, buckets, usage, "budget-probe:retry-stated");
fixture
.inner
.push(Ok(throttled(StatusCode::TOO_MANY_REQUESTS, Some("7"))));
fixture.inner.push(Ok(ok_response(b"ok")));
let response = fixture
.client
.send(get_request("/retry"))
.await
.expect("the retry lands after the stated pause");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(response.body().as_ref(), b"ok");
assert_eq!(fixture.sends(), 2, "one throttled attempt, one retry");
assert_eq!(
fixture.sleep_log(),
vec![Duration::from_secs(7)],
"the header's pause, verbatim"
);
}
async fn check_a_post_is_not_retried(db: &Arc<dyn Database>, buckets: &str, usage: &str) {
let fixture = retry_fixture(db, buckets, usage, "budget-probe:retry-post");
fixture
.inner
.push(Ok(throttled(StatusCode::TOO_MANY_REQUESTS, Some("7"))));
let response = fixture
.client
.send(post_request("/retry"))
.await
.expect("a throttled POST is a response, not an error");
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert_eq!(fixture.sends(), 1, "a POST owns its weather");
assert!(fixture.sleep_log().is_empty(), "no wait was made");
}
async fn check_the_idempotent_extension_vouches(
db: &Arc<dyn Database>,
buckets: &str,
usage: &str,
) {
let fixture = retry_fixture(db, buckets, usage, "budget-probe:retry-vouched");
fixture
.inner
.push(Ok(throttled(StatusCode::TOO_MANY_REQUESTS, None)));
fixture.inner.push(Ok(ok_response(b"ok")));
let mut vouched = post_request("/retry");
vouched.extensions_mut().insert(Idempotent(true));
let response = fixture
.client
.send(vouched)
.await
.expect("the vouched retry lands");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(fixture.sends(), 2, "the vouch makes the POST retryable");
assert_eq!(
fixture.sleep_log(),
vec![RetryPolicy::default().base()],
"no header stated: the schedule's first backoff"
);
}
async fn check_transport_retries_on_the_schedule(
db: &Arc<dyn Database>,
buckets: &str,
usage: &str,
) {
let fixture = retry_fixture(db, buckets, usage, "budget-probe:retry-weather");
fixture
.inner
.push(Err(HttpError::Transport("reset".to_owned())));
fixture
.inner
.push(Err(HttpError::Transport("reset".to_owned())));
fixture.inner.push(Ok(ok_response(b"ok")));
let response = fixture
.client
.send(get_request("/retry"))
.await
.expect("the weather clears on the third attempt");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
fixture.sends(),
3,
"the first send plus exactly the two retries"
);
let policy = RetryPolicy::default();
assert_eq!(
fixture.sleep_log(),
vec![policy.base(), policy.base() * 2],
"exponential: base, then double"
);
}
pub async fn assert_budget_passthrough(db: &Arc<dyn Database>) {
let buckets = "budget_probe_passthrough_buckets";
let usage = "budget_probe_passthrough_usage";
create_probe_tables(db, buckets, usage).await;
let router: BudgetRouter = Arc::new(|_request| None);
let fixture = Fixture::new(db, buckets, usage, RetryPolicy::default(), router);
let mut request = http::Request::builder()
.method(http::Method::POST)
.uri("https://upstream.example.test/v1/things?x=1")
.header("x-probe", "yes")
.body(Bytes::from_static(b"payload"))
.expect("test request");
request.extensions_mut().insert(Idempotent(true));
fixture.inner.push(Ok(http::Response::builder()
.status(StatusCode::ACCEPTED)
.header("x-probe-echo", "same")
.body(Bytes::from_static(b"untouched"))
.expect("test response")));
let response = fixture
.client
.send(request)
.await
.expect("an unrouted request passes through");
assert_eq!(response.status(), StatusCode::ACCEPTED);
assert_eq!(response.body().as_ref(), b"untouched");
assert_eq!(
response
.headers()
.get("x-probe-echo")
.map(http::HeaderValue::as_bytes),
Some(b"same".as_slice()),
"the response's own headers survive"
);
let captured = fixture.inner.captured();
let [sent] = captured.as_slice() else {
panic!(
"the inner client saw exactly one request, got {}",
captured.len()
);
};
assert_eq!(sent.method(), http::Method::POST);
assert_eq!(sent.uri(), "https://upstream.example.test/v1/things?x=1");
assert_eq!(
sent.headers()
.get("x-probe")
.map(http::HeaderValue::as_bytes),
Some(b"yes".as_slice())
);
assert_eq!(sent.body().as_ref(), b"payload");
assert_eq!(
sent.extensions().get::<Idempotent>(),
Some(&Idempotent(true)),
"the extensions ride along"
);
assert!(fixture.sleep_log().is_empty(), "no budget, no waits");
assert_eq!(row_count(db, buckets).await, 0, "no bucket rows written");
assert_eq!(row_count(db, usage).await, 0, "no quota rows written");
clear_probe_tables(db, buckets, usage).await;
}
async fn row_count(db: &Arc<dyn Database>, table: &str) -> i64 {
let rows = db
.query(&Statement::new(format!(
"SELECT COUNT(*) AS n FROM {table}"
)))
.await
.expect("count probe rows");
rows.first()
.and_then(|row| row.get::<i64>("n"))
.expect("one count row")
}