mirror of
https://github.com/block/buzz.git
synced 2026-11-07 18:01:06 +00:00
fix(relay): exclude ephemeral activity from message quota (#7736)
## Summary Stop charging ephemeral WebSocket activity against the per-minute `Messages` quota. Presence, typing, and observer traffic should not consume the allowance intended for persistable event attempts. - Use the existing `buzz_core::kind::is_ephemeral` predicate (kinds `20000..=29999`) at the second admission gate. All authenticated EVENT/REQ/COUNT traffic still consumes the shared `WsEvents` budget first. Non-ephemeral EVENT attempts remain charged, including attempts later rejected by authorization or validation. - Add bounded `bucket` labels (`messages`, `ws_operations`, `api_calls`) to `buzz_admission_rejections_total`. HTTP changes are metric labeling only. Wire rejection text remains byte-for-byte unchanged, preserving existing client parsers. - Add real-Redis admission regressions and select them in the existing relay integration CI job. No quota increases, identity/keying changes, client changes, new limiter, or deployment. Downstream authorization and fail-closed Redis behavior are unchanged. ### Related issue Duplicate search found overlapping work: open #4912 and closed #4902. They propose the same ephemeral-accounting correction alongside broader wire/client/config changes. The relay fix is a deliberately narrower four-file implementation on the current admission module, with metrics-only diagnostics and real-Redis regression coverage. A separate test-only CI repair is documented below. It does not close or claim to supersede those PRs. Related symptom: #7411. This PR does **not** resolve subscription bursts exhausting the shared WS budget; that budget is intentionally preserved. It also does not establish or fix the cause of HTTP quota exhaustion. ### Testing Validation ran against base `779af8886caae1317b4de962082429867ab61503` plus the exact patch now committed as `70ea0528594e25e6dcb40cd74620bd5ef0177592` (binary diff SHA-256 `5064e1cbb87ecb9eff10336bae14b912cac442e1edc247aa4f093de30e4f9717`). Normal commit/push hooks subsequently passed on the committed head, without bypass. Passed: - Independent review and rerun: 11 rejection tests (including real Redis) plus 4 admission tests. Covers ephemeral range boundaries, human/agent budgets, persistable events, ephemeral admission after message exhaustion, mixed EVENT/REQ/COUNT flood exhaustion, Redis outage, exact wire payloads, and bucket metrics. Three mutations were caught: reintroducing ephemeral billing, exempting stored events, and bypassing WS rejection. - `just test-unit`: all 17 runner groups; `just test-integration`: all 3 runner groups. Auth integration reported no tests; ignored database integration tests are not claimed as executed. - Relay all-target Clippy, formatting, workflow actionlint, pre-commit/signoff hooks, and all applicable pre-push lanes (`push-head-scope`, `branch-skew`, `file-size-check`, `rust-tests`, `desktop-tauri-checks`). - Release relay/CLI/admin build and isolated release-binary workflow: channel creation, signed root/reply sends, exact-ID readback, and thread retrieval. Debug and release WS smoke admitted three presence events with a Messages limit of two, admitted two persistable events, rejected the third, still admitted presence afterward, and rejected author mismatch. Debug Prometheus confirmed `bucket="messages"`. **Outstanding validation gap:** the additional `cargo test -p buzz-relay --no-fail-fast` run was not green: library 1,055 passed / 2 failed / 100 ignored; binary 13 passed / 1 ignored; boot integration 9 passed. Failures were `api::mesh_demo::tests::demo_join_forwarded_arm_round_trips_echo` (504 instead of 200) and `telemetry::tests::trace_context_lookup_does_not_enable_callsites` (tracing assertion). Neither failure was diagnosed or repaired in this scope, and this PR does not claim a fully green package suite. `just ci` was not run locally; hosted CI remains to validate the PR. All runtime tests used isolated local services. Owned services have been stopped; nothing was deployed and no production data was touched. ### CI cancellation-test repair At the original relay head, hosted Rust unit tests failed on the pre-existing cancellation harness race tracked in #4945. With owner authorization, this PR now adapts the one-file repair from #5003 (original commit `cc8dc6124957acfd67fe2336396b58e0ff43ec60`), preserving author credit. #5003's old hosted workflows failed without starting jobs; their status is not used as validation evidence. Commit `55fe5f16618a51d03ad48d6f2f97b2f8b463607a` holds the fake provider response until a successful cancellation acknowledgement, then replays all pre-acknowledgement usage/prompt frames in order. It retains the cancellation and usage-before-response assertions and explicitly requires prompt completion. No production code changes in this follow-up. Validation on parent `70ea052859` plus this test patch, before the pre-commit-only assertion wrapping: full `cargo test -p buzz-agent --no-fail-fast` passed (715 tests, one ignored); 100 targeted repetitions passed with four concurrent workers; all-target buzz-agent Clippy passed. The unmodified parent also passed 50 local repetitions, so this is not claimed as a local reproduction of the intermittent CI failure. The synchronization defect was confirmed from the source, original CI log, and #4945. Hosted CI must validate the updated head. Independent source review cleared exact commit `55fe5f16618a51d03ad48d6f2f97b2f8b463607a`. Normal pre-commit/signoff and all applicable pre-push hooks passed without bypass. The original relay-package mesh-demo and telemetry failures remain disclosed above and were not modified by this test-only repair. --------- Signed-off-by: Wes <wesbillman@users.noreply.github.com> Co-authored-by: Carl <32a2e2c9d428ee08902cab75d956da2c1d235a22d4766b0dd4138bf6e2e5db1d@buzz.block.builderlab.xyz> Co-authored-by: kiranmagic7 <262980978+kiranmagic7@users.noreply.github.com>
This commit is contained in:
8
.github/workflows/_ci-relay.yml
vendored
8
.github/workflows/_ci-relay.yml
vendored
@@ -497,6 +497,14 @@ jobs:
|
||||
--archive-file target/ci/backend-integration-tests.tar.zst \
|
||||
-E 'package(buzz-relay) and test(/api::git::settings::tests::external_infra::/)' \
|
||||
--run-ignored ignored-only
|
||||
- name: Relay quota accounting and rejection regressions
|
||||
env:
|
||||
BUZZ_TEST_REDIS_URL: redis://localhost:6379
|
||||
run: |
|
||||
cargo nextest run \
|
||||
--archive-file target/ci/backend-integration-tests.tar.zst \
|
||||
-E 'package(buzz-relay) and (test(/^rejection::/) or test(/^admission::/))' \
|
||||
--run-ignored all
|
||||
- name: Workflow message provenance unit tests
|
||||
# The relay's workflow_sink suite is not selected by the infra-free
|
||||
# unit job. Its ignored database cases run in the isolated PostgreSQL
|
||||
|
||||
@@ -1320,7 +1320,7 @@ async fn mid_turn_usage_includes_earlier_turns() {
|
||||
/// Setup: round 1 is a tool call WITH usage (tokens are captured). After the
|
||||
/// tool_call_update notification (proving round 1 is fully processed), we gate
|
||||
/// the round-2 LLM response behind a `oneshot` barrier that only releases after
|
||||
/// cancel is sent. This guarantees the turn exits with `stopReason: "cancelled"`
|
||||
/// cancel is acknowledged. This guarantees the turn exits with `stopReason: "cancelled"`
|
||||
/// deterministically, even on a slow CI worker.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn cancelled_turn_with_usage_emits_notification_before_response() {
|
||||
@@ -1335,7 +1335,7 @@ async fn cancelled_turn_with_usage_emits_notification_before_response() {
|
||||
// in-flight TCP request can resolve. The queue is empty for round 2, so the
|
||||
// agent receives the fallback "no canned response" body which it treats as
|
||||
// an LLM error; the cancel check at the round boundary fires first because
|
||||
// the gate is only released after cancel is enqueued.
|
||||
// the gate is only released after cancel is acknowledged.
|
||||
let responses = vec![openai_tool_call_with_usage(
|
||||
"call_cancel_test",
|
||||
"fake__noop",
|
||||
@@ -1371,7 +1371,7 @@ async fn cancelled_turn_with_usage_emits_notification_before_response() {
|
||||
}
|
||||
}
|
||||
// For request 2+ (round 2), wait for the gate to open before
|
||||
// responding. This ensures cancel is sent before round 2 resolves,
|
||||
// responding. This ensures cancel is processed before round 2 resolves,
|
||||
// making stopReason: cancelled deterministic.
|
||||
if req_num >= 2 {
|
||||
let rx = gate.lock().await.take();
|
||||
@@ -1415,20 +1415,27 @@ async fn cancelled_turn_with_usage_emits_notification_before_response() {
|
||||
})
|
||||
.await;
|
||||
|
||||
// Now send cancel and release the round-2 gate. Cancel is enqueued before
|
||||
// round 2 can respond, so the turn exits with stopReason: cancelled.
|
||||
// Writing to stdin does not prove the agent processed cancel. Keep the
|
||||
// provider blocked until the acknowledgement, retaining all earlier frames
|
||||
// so usage/prompt ordering is checked even if the turn finishes before ACK.
|
||||
let c_id = h.send("session/cancel", json!({"sessionId": sid})).await;
|
||||
let (frames_before_cancel_ack, cancel_ack) =
|
||||
recv_until_with_drain(&mut h, |v| v["id"] == json!(c_id)).await;
|
||||
assert_eq!(cancel_ack.get("result"), Some(&Value::Null), "{cancel_ack}");
|
||||
assert!(cancel_ack.get("error").is_none(), "{cancel_ack}");
|
||||
let _ = gate_tx.send(()); // unblock round 2
|
||||
|
||||
let mut saw_usage_before_prompt_response = false;
|
||||
let mut saw_usage = false;
|
||||
let mut saw_cancel_ok = false;
|
||||
let mut saw_prompt_response = false;
|
||||
for _ in 0..40 {
|
||||
let v = h.recv().await;
|
||||
if v["id"] == json!(c_id) {
|
||||
saw_cancel_ok = true;
|
||||
} else if is_usage_update(&v) {
|
||||
let mut pending_frames = VecDeque::from(frames_before_cancel_ack);
|
||||
let frame_budget = 40 + pending_frames.len();
|
||||
for _ in 0..frame_budget {
|
||||
let v = match pending_frames.pop_front() {
|
||||
Some(v) => v,
|
||||
None => h.recv().await,
|
||||
};
|
||||
if is_usage_update(&v) {
|
||||
saw_usage = true;
|
||||
if !saw_prompt_response {
|
||||
saw_usage_before_prompt_response = true;
|
||||
@@ -1441,11 +1448,14 @@ async fn cancelled_turn_with_usage_emits_notification_before_response() {
|
||||
"turn must end with stopReason: cancelled"
|
||||
);
|
||||
}
|
||||
if saw_usage && saw_prompt_response && saw_cancel_ok {
|
||||
if saw_usage && saw_prompt_response {
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(saw_cancel_ok, "session/cancel was not acknowledged");
|
||||
assert!(
|
||||
saw_prompt_response,
|
||||
"session/prompt did not finish after cancel"
|
||||
);
|
||||
assert!(
|
||||
saw_usage,
|
||||
"expected usage_update notification for cancelled turn with observed tokens"
|
||||
|
||||
@@ -39,14 +39,14 @@ pub(crate) async fn enforce_http_admission(
|
||||
{
|
||||
Ok(()) => Ok(()),
|
||||
Err(crate::admission::AdmissionError::Exceeded { reset_in_secs }) => {
|
||||
metrics::counter!("buzz_admission_rejections_total", "transport" => "http", "reason" => "quota").increment(1);
|
||||
metrics::counter!("buzz_admission_rejections_total", "transport" => "http", "reason" => "quota", "bucket" => "api_calls").increment(1);
|
||||
Err(api_error(
|
||||
StatusCode::TOO_MANY_REQUESTS,
|
||||
&format!("rate-limited: quota exceeded; retry in {reset_in_secs}s"),
|
||||
))
|
||||
}
|
||||
Err(crate::admission::AdmissionError::Unavailable) => {
|
||||
metrics::counter!("buzz_admission_rejections_total", "transport" => "http", "reason" => "unavailable").increment(1);
|
||||
metrics::counter!("buzz_admission_rejections_total", "transport" => "http", "reason" => "unavailable", "bucket" => "api_calls").increment(1);
|
||||
Err(api_error(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"rate-limited: shared admission unavailable",
|
||||
|
||||
@@ -75,11 +75,12 @@ pub(crate) async fn enforce_ws_admission(
|
||||
ws_limit,
|
||||
)
|
||||
.await;
|
||||
if !send_admission_result(conn, ws_result, msg) {
|
||||
if !send_admission_result(conn, ws_result, msg, "ws_operations") {
|
||||
return false;
|
||||
}
|
||||
|
||||
if is_event {
|
||||
if matches!(msg, ClientMessage::Event(event) if !buzz_core::kind::is_ephemeral(event.kind.as_u16() as u32))
|
||||
{
|
||||
let message_limit = if is_agent {
|
||||
limits.agent_standard_messages_per_min
|
||||
} else {
|
||||
@@ -94,9 +95,9 @@ pub(crate) async fn enforce_ws_admission(
|
||||
message_limit,
|
||||
)
|
||||
.await;
|
||||
// The per-minute message quota only applies to EVENTs, and its
|
||||
// rejection must be as correlatable as the burst quota's.
|
||||
if !send_admission_result(conn, message_result, msg) {
|
||||
// Only persistable EVENT attempts consume the minute quota. Ephemeral activity
|
||||
// still passes the shared WS flood limit above and normal event auth.
|
||||
if !send_admission_result(conn, message_result, msg, "messages") {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
@@ -115,12 +116,13 @@ fn send_admission_result(
|
||||
conn: &ConnectionState,
|
||||
result: Result<(), AdmissionError>,
|
||||
msg: &ClientMessage,
|
||||
bucket: &'static str,
|
||||
) -> bool {
|
||||
let target = rejection_target_for(msg);
|
||||
match result {
|
||||
Ok(()) => true,
|
||||
Err(AdmissionError::Exceeded { reset_in_secs }) => {
|
||||
metrics::counter!("buzz_admission_rejections_total", "transport" => "websocket", "reason" => "quota").increment(1);
|
||||
metrics::counter!("buzz_admission_rejections_total", "transport" => "websocket", "reason" => "quota", "bucket" => bucket).increment(1);
|
||||
conn.send(request_rejection_message(
|
||||
target,
|
||||
&format!("rate-limited: quota exceeded; retry in {reset_in_secs}s"),
|
||||
@@ -128,7 +130,7 @@ fn send_admission_result(
|
||||
false
|
||||
}
|
||||
Err(AdmissionError::Unavailable) => {
|
||||
metrics::counter!("buzz_admission_rejections_total", "transport" => "websocket", "reason" => "unavailable").increment(1);
|
||||
metrics::counter!("buzz_admission_rejections_total", "transport" => "websocket", "reason" => "unavailable", "bucket" => bucket).increment(1);
|
||||
conn.send(request_rejection_message(
|
||||
target,
|
||||
"rate-limited: shared admission unavailable",
|
||||
@@ -198,6 +200,7 @@ mod tests {
|
||||
&conn,
|
||||
Err(AdmissionError::Exceeded { reset_in_secs: 7 }),
|
||||
&msg,
|
||||
"ws_operations",
|
||||
);
|
||||
|
||||
assert!(!admitted, "an over-quota frame is not admitted");
|
||||
@@ -227,7 +230,12 @@ mod tests {
|
||||
let (conn, mut rx) = test_conn();
|
||||
let (msg, event_id) = parsed_event_message();
|
||||
|
||||
send_admission_result(&conn, Err(AdmissionError::Unavailable), &msg);
|
||||
send_admission_result(
|
||||
&conn,
|
||||
Err(AdmissionError::Unavailable),
|
||||
&msg,
|
||||
"ws_operations",
|
||||
);
|
||||
|
||||
let frame = sent_frame(&mut rx);
|
||||
assert_eq!(frame[0], "OK");
|
||||
@@ -247,6 +255,7 @@ mod tests {
|
||||
&conn,
|
||||
Err(AdmissionError::Exceeded { reset_in_secs: 7 }),
|
||||
&msg,
|
||||
"ws_operations",
|
||||
);
|
||||
|
||||
let frame = sent_frame(&mut rx);
|
||||
@@ -269,6 +278,7 @@ mod tests {
|
||||
&conn,
|
||||
Err(AdmissionError::Exceeded { reset_in_secs: 7 }),
|
||||
&msg,
|
||||
"ws_operations",
|
||||
);
|
||||
|
||||
let frame = sent_frame(&mut rx);
|
||||
@@ -333,3 +343,6 @@ mod tests {
|
||||
assert_eq!(frame[1], "history-abc");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod quota_tests;
|
||||
|
||||
250
crates/buzz-relay/src/rejection/quota_tests.rs
Normal file
250
crates/buzz-relay/src/rejection/quota_tests.rs
Normal file
@@ -0,0 +1,250 @@
|
||||
//! Real Redis counters through the production admission entrypoints. Tests use
|
||||
//! disposable identities and expiring keys, never FLUSHDB or a shared identity.
|
||||
|
||||
use super::*;
|
||||
use axum::extract::ws::Message as WsMessage;
|
||||
use buzz_auth::{AuthService, RateLimitConfig};
|
||||
use buzz_pubsub::rate_limiter::RedisRateLimiter;
|
||||
use metrics_util::debugging::{DebugValue, DebuggingRecorder};
|
||||
use nostr::{EventBuilder, Keys, Kind};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::connection::tests::{authenticated_state, read_frame, test_conn_with_auth};
|
||||
|
||||
async fn state_with_limits(messages: u64, agent_messages: u64, ws: u64) -> AppState {
|
||||
let url = std::env::var("BUZZ_TEST_REDIS_URL").expect("explicit isolated BUZZ_TEST_REDIS_URL");
|
||||
let mut state = (*crate::state::tests::test_state().await).clone();
|
||||
let mut auth = state.auth.config().clone();
|
||||
auth.rate_limits = RateLimitConfig {
|
||||
human_messages_per_min: messages,
|
||||
agent_standard_messages_per_min: agent_messages,
|
||||
human_ws_events_per_sec: ws,
|
||||
human_api_calls_per_min: 2,
|
||||
..RateLimitConfig::default()
|
||||
};
|
||||
state.auth = Arc::new(AuthService::new(auth));
|
||||
state.redis_pool = deadpool_redis::Config::from_url(url)
|
||||
.create_pool(Some(deadpool_redis::Runtime::Tokio1))
|
||||
.unwrap();
|
||||
state.admission_rate_limiter = Arc::new(RedisRateLimiter::new(state.redis_pool.clone()));
|
||||
state
|
||||
}
|
||||
|
||||
fn connection(agent: bool) -> (Arc<ConnectionState>, mpsc::Receiver<WsMessage>) {
|
||||
let AuthState::Authenticated(mut auth) = authenticated_state() else {
|
||||
unreachable!()
|
||||
};
|
||||
if agent {
|
||||
auth.agent_owner_pubkey = Some(Keys::generate().public_key());
|
||||
}
|
||||
test_conn_with_auth(AuthState::Authenticated(auth))
|
||||
}
|
||||
|
||||
fn event(kind: u16) -> ClientMessage {
|
||||
let event = EventBuilder::new(Kind::from(kind), "")
|
||||
.sign_with_keys(&Keys::generate())
|
||||
.unwrap();
|
||||
ClientMessage::parse(&serde_json::json!(["EVENT", event]).to_string()).unwrap()
|
||||
}
|
||||
|
||||
fn pubkey(conn: &ConnectionState) -> nostr::PublicKey {
|
||||
let AuthState::Authenticated(auth) = conn.auth_state_snapshot() else {
|
||||
unreachable!()
|
||||
};
|
||||
auth.pubkey
|
||||
}
|
||||
|
||||
async fn counter(state: &AppState, conn: &ConnectionState, bucket: LimitType) -> Option<u64> {
|
||||
let key = buzz_auth::rate_limit::rate_limit_key(&conn.tenant, &pubkey(conn), &bucket);
|
||||
let mut redis = state.redis_pool.get().await.unwrap();
|
||||
redis::cmd("GET")
|
||||
.arg(key)
|
||||
.query_async(&mut *redis)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn assert_metric(
|
||||
recorder: &DebuggingRecorder,
|
||||
transport: &str,
|
||||
reason: &str,
|
||||
bucket: &str,
|
||||
count: u64,
|
||||
) {
|
||||
let snapshot = recorder.snapshotter().snapshot().into_vec();
|
||||
let labels = [
|
||||
("transport", transport),
|
||||
("reason", reason),
|
||||
("bucket", bucket),
|
||||
];
|
||||
let actual = snapshot.iter().find_map(|(key, _, _, value)| {
|
||||
(key.key().name() == "buzz_admission_rejections_total"
|
||||
&& labels
|
||||
.iter()
|
||||
.all(|(k, v)| key.key().labels().any(|l| l.key() == *k && l.value() == *v)))
|
||||
.then(|| match value {
|
||||
DebugValue::Counter(value) => *value,
|
||||
_ => panic!("expected counter"),
|
||||
})
|
||||
});
|
||||
assert_eq!(actual, Some(count));
|
||||
}
|
||||
|
||||
fn assert_quota_frame(frame: &serde_json::Value, msg: &ClientMessage, max_retry: u64) {
|
||||
let reason = match msg {
|
||||
ClientMessage::Event(event) => {
|
||||
assert_eq!(frame[0], "OK");
|
||||
assert_eq!(frame[1], event.id.to_hex());
|
||||
assert_eq!(frame[2], false);
|
||||
frame[3].as_str().unwrap()
|
||||
}
|
||||
ClientMessage::Req { sub_id, .. } | ClientMessage::Count { sub_id, .. } => {
|
||||
assert_eq!(frame[0], "CLOSED");
|
||||
assert_eq!(frame[1], *sub_id);
|
||||
frame[2].as_str().unwrap()
|
||||
}
|
||||
_ => panic!("unexpected frame"),
|
||||
};
|
||||
// Exact grammar: clients including buzz-app anchor both ends of this hint.
|
||||
let seconds: u64 = reason
|
||||
.strip_prefix("rate-limited: quota exceeded; retry in ")
|
||||
.unwrap()
|
||||
.strip_suffix('s')
|
||||
.unwrap()
|
||||
.parse()
|
||||
.unwrap();
|
||||
assert!(seconds <= max_retry);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires explicit isolated BUZZ_TEST_REDIS_URL"]
|
||||
async fn ephemeral_events_leave_human_and_agent_message_budgets_untouched() {
|
||||
let state = state_with_limits(2, 3, 100).await;
|
||||
let recorder = DebuggingRecorder::new();
|
||||
let _guard = metrics::set_default_local_recorder(&recorder);
|
||||
// Both range boundaries, typing, presence, and observer frames. Adjacent
|
||||
// non-ephemeral kinds below must still consume the message allowance.
|
||||
let ephemeral_kinds = [20000, 20001, 20002, 24200, 29999];
|
||||
for (agent, limit) in [(false, 2), (true, 3)] {
|
||||
for stored_kind in [0, 1, 9, 19999, 30000, 40002] {
|
||||
let (conn, mut rx) = connection(agent);
|
||||
for kind in ephemeral_kinds {
|
||||
assert!(enforce_ws_admission(&event(kind), &conn, &state).await);
|
||||
}
|
||||
assert_eq!(counter(&state, &conn, LimitType::Messages).await, None);
|
||||
assert_eq!(counter(&state, &conn, LimitType::WsEvents).await, Some(5));
|
||||
for expected in 1..=limit {
|
||||
assert!(enforce_ws_admission(&event(stored_kind), &conn, &state).await);
|
||||
assert_eq!(
|
||||
counter(&state, &conn, LimitType::Messages).await,
|
||||
Some(expected)
|
||||
);
|
||||
}
|
||||
let rejected = event(stored_kind);
|
||||
assert!(!enforce_ws_admission(&rejected, &conn, &state).await);
|
||||
assert_quota_frame(&read_frame(&mut rx), &rejected, 60);
|
||||
for kind in ephemeral_kinds {
|
||||
assert!(enforce_ws_admission(&event(kind), &conn, &state).await);
|
||||
}
|
||||
assert_eq!(
|
||||
counter(&state, &conn, LimitType::Messages).await,
|
||||
Some(limit + 1)
|
||||
);
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
}
|
||||
assert_metric(&recorder, "websocket", "quota", "messages", 12);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires explicit isolated BUZZ_TEST_REDIS_URL"]
|
||||
async fn ws_flood_budget_is_shared_by_ephemeral_stored_req_and_count() {
|
||||
let state = state_with_limits(10, 20, 1).await; // Five operations per window.
|
||||
let recorder = DebuggingRecorder::new();
|
||||
let _guard = metrics::set_default_local_recorder(&recorder);
|
||||
let (conn, mut rx) = connection(false);
|
||||
let req = ClientMessage::parse(r#"["REQ","history",{"kinds":[9]}]"#).unwrap();
|
||||
let count = ClientMessage::parse(r#"["COUNT","count",{"kinds":[9]}]"#).unwrap();
|
||||
for msg in [
|
||||
event(20001),
|
||||
event(20002),
|
||||
event(9),
|
||||
req.clone(),
|
||||
count.clone(),
|
||||
] {
|
||||
assert!(enforce_ws_admission(&msg, &conn, &state).await);
|
||||
}
|
||||
assert_eq!(counter(&state, &conn, LimitType::Messages).await, Some(1));
|
||||
for msg in [event(24200), event(9), req, count] {
|
||||
assert!(!enforce_ws_admission(&msg, &conn, &state).await);
|
||||
assert_quota_frame(&read_frame(&mut rx), &msg, 5);
|
||||
}
|
||||
assert_eq!(counter(&state, &conn, LimitType::Messages).await, Some(1));
|
||||
assert_metric(&recorder, "websocket", "quota", "ws_operations", 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires explicit isolated BUZZ_TEST_REDIS_URL"]
|
||||
async fn http_quota_preserves_wire_error_and_reports_api_bucket() {
|
||||
let state = state_with_limits(0, 0, 0).await;
|
||||
let recorder = DebuggingRecorder::new();
|
||||
let _guard = metrics::set_default_local_recorder(&recorder);
|
||||
let (conn, _) = connection(false);
|
||||
for _ in 0..2 {
|
||||
crate::api::bridge::enforce_http_admission(&state, &conn.tenant, &pubkey(&conn))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let (status, body) =
|
||||
crate::api::bridge::enforce_http_admission(&state, &conn.tenant, &pubkey(&conn))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(status, axum::http::StatusCode::TOO_MANY_REQUESTS);
|
||||
let seconds: u64 = body.0["error"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.strip_prefix("rate-limited: quota exceeded; retry in ")
|
||||
.unwrap()
|
||||
.strip_suffix('s')
|
||||
.unwrap()
|
||||
.parse()
|
||||
.unwrap();
|
||||
assert!(seconds <= 60);
|
||||
assert_eq!(counter(&state, &conn, LimitType::Messages).await, None);
|
||||
assert_eq!(counter(&state, &conn, LimitType::WsEvents).await, None);
|
||||
assert_metric(&recorder, "http", "quota", "api_calls", 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ephemeral_and_http_admission_still_fail_closed_when_redis_is_unavailable() {
|
||||
let state = crate::state::tests::test_state().await;
|
||||
let recorder = DebuggingRecorder::new();
|
||||
let _guard = metrics::set_default_local_recorder(&recorder);
|
||||
let (conn, mut rx) = connection(false);
|
||||
let msg = event(20002);
|
||||
assert!(!enforce_ws_admission(&msg, &conn, &state).await);
|
||||
let ClientMessage::Event(event) = msg else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(
|
||||
read_frame(&mut rx),
|
||||
serde_json::json!([
|
||||
"OK",
|
||||
event.id.to_hex(),
|
||||
false,
|
||||
"rate-limited: shared admission unavailable"
|
||||
])
|
||||
);
|
||||
assert_metric(&recorder, "websocket", "unavailable", "ws_operations", 1);
|
||||
let (status, body) =
|
||||
crate::api::bridge::enforce_http_admission(&state, &conn.tenant, &pubkey(&conn))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(status, axum::http::StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
body.0["error"],
|
||||
"rate-limited: shared admission unavailable"
|
||||
);
|
||||
assert_metric(&recorder, "http", "unavailable", "api_calls", 1);
|
||||
}
|
||||
Reference in New Issue
Block a user