1//! The account's rate limit is shared cluster-wide and enforced before
2//! sending (contract *Execution*, "The account's rate limit is shared
3//! cluster-wide"). Every arrival is timed at the mock, so the limit is
4//! measured where the account would bill it, not inferred from the
5//! queries' durations.
6//!
7//! GCRA admits one period of burst (`postjevsql-core/src/gcra.rs`): at
8//! `jev.max_requests_per_minute = 60` about 61 requests go at once, then
9//! one a second.
10
11use std::sync::{Arc, Mutex};
12use std::time::{Duration, Instant};
13
14use support::mock_jev::{MockJev, Reply};
15use support::{jev_instance_with, noul};
16use tokio_postgres::error::SqlState;
17
18/// Every request's arrival at the mock, in order.
19type Arrivals = Arc<Mutex<Vec<Instant>>>;
20
21/// No `shared_preload_libraries`: the limiter's words are in a named DSM
22/// segment, attached when the first scan begins.
23const SETTINGS: &[(&str, &str)] = &[
24    ("jev.max_requests_per_minute", "60"),
25    ("jev.max_tokens_per_second", "0"),
26];
27
28/// Distinct rows, so neither the cache nor dedupe folds any together.
29async fn rows(client: &tokio_postgres::Client, table: &str, n: i32) {
30    client
31        .batch_execute(&format!(
32            "SET statement_timeout = '60s';
33             CREATE TABLE {table} AS SELECT g AS id, '{table} ' || g AS body FROM generate_series(1, {n}) g;"
34        ))
35        .await
36        .unwrap();
37}
38
39async fn until(what: &str, mut done: impl FnMut() -> bool) {
40    let deadline = Instant::now() + Duration::from_secs(30);
41    while !done() {
42        assert!(Instant::now() < deadline, "timed out waiting for {what}");
43        tokio::time::sleep(Duration::from_millis(20)).await;
44    }
45}
46
47#[tokio::test(flavor = "multi_thread")]
48async fn two_backends_share_one_requests_per_minute_ceiling() {
49    let arrivals: Arrivals = Arc::default();
50    let seen = arrivals.clone();
51    let mock = MockJev::start(move |_| {
52        seen.lock().unwrap().push(Instant::now());
53        Reply::json(200, noul(0.5))
54    })
55    .await;
56    let (pg, a) = jev_instance_with(&mock, SETTINGS).await;
57    let (b, conn) = tokio_postgres::connect(&pg.conn_str(), tokio_postgres::NoTls).await.unwrap();
58    tokio::spawn(conn);
59    rows(&a, "ra", 33).await;
60    rows(&b, "rb", 33).await;
61
62    // 66 requests: each backend alone fits in the burst, together they
63    // do not. A per-backend limit would answer both at once.
64    let (ra, rb) = tokio::join!(
65        a.query("SELECT jev_prob(t, 'Urgent?') FROM ra t", &[]),
66        b.query("SELECT jev_prob(t, 'Urgent?') FROM rb t", &[]),
67    );
68    assert_eq!(ra.unwrap().len(), 33);
69    assert_eq!(rb.unwrap().len(), 33);
70
71    let times = arrivals.lock().unwrap().clone();
72    assert_eq!(times.len(), 66, "one request per row, none retried");
73    let t0 = times[0];
74    // Cumulatively: by t seconds after the first, at most 62 + t sent
75    // (burst, one on the boundary, then one a second; 150 ms of slack
76    // for the two backends' clocks and the mock's own scheduling).
77    for (i, t) in times.iter().enumerate() {
78        let allowed = 62.0 + (t.duration_since(t0).as_secs_f64() + 0.15);
79        assert!((i + 1) as f64 <= allowed, "request {} arrived {:?} after the first", i + 1, t.duration_since(t0));
80    }
81    let span = times[65].duration_since(t0);
82    assert!(span >= Duration::from_millis(3_800), "66 requests in {span:?}: the ceiling was not shared");
83}
84
85#[tokio::test(flavor = "multi_thread")]
86async fn a_429_halves_the_rate_that_follows() {
87    let arrivals: Arrivals = Arc::default();
88    let seen = arrivals.clone();
89    let mock = MockJev::start(move |_| {
90        let mut seen = seen.lock().unwrap();
91        seen.push(Instant::now());
92        // The first request past the burst is throttled, and asks for no
93        // delay, so any spacing after it is the limiter's.
94        if seen.len() == 62 {
95            Reply::json(429, serde_json::json!({"detail": {"error_type": "rate_limit", "message": "slow down"}}))
96                .header("retry-after-ms", "0")
97        } else {
98            Reply::json(200, noul(0.5))
99        }
100    })
101    .await;
102    let (_pg, client) = jev_instance_with(&mock, SETTINGS).await;
103    rows(&client, "r", 65).await;
104
105    let answered = client.query("SELECT jev_prob(t, 'Urgent?') FROM r t", &[]).await.unwrap();
106    assert_eq!(answered.len(), 65);
107
108    let times = arrivals.lock().unwrap().clone();
109    assert_eq!(times.len(), 66, "65 rows and the 429's retry");
110    // At the ceiling the steady spacing is 1 s. After the 429 the rate is
111    // half (2 s), less the small climb back per answer (1/128 each).
112    for pair in times[61..].windows(2) {
113        let gap = pair[1].duration_since(pair[0]);
114        assert!(gap >= Duration::from_millis(1_800), "a gap of {gap:?} after the 429");
115    }
116    let retries: i64 = client.query_one("SELECT retries FROM jev_stats()", &[]).await.unwrap().get(0);
117    assert_eq!(retries, 1);
118}
119
120#[tokio::test(flavor = "multi_thread")]
121async fn pg_cancel_backend_interrupts_the_admission_wait() {
122    let arrivals: Arrivals = Arc::default();
123    let seen = arrivals.clone();
124    let mock = MockJev::start(move |_| {
125        seen.lock().unwrap().push(Instant::now());
126        Reply::json(200, noul(0.5))
127    })
128    .await;
129    // One a minute: two requests fit the burst, the third waits ~60 s.
130    let (pg, client) =
131        jev_instance_with(&mock, &[("jev.max_requests_per_minute", "1"), SETTINGS[1]]).await;
132    rows(&client, "r", 3).await;
133    let pid: i32 = client.query_one("SELECT pg_backend_pid()", &[]).await.unwrap().get(0);
134
135    let query = tokio::spawn(async move {
136        let result = client.query("SELECT jev_prob(t, 'Urgent?') FROM r t", &[]).await.map(drop);
137        (client, result)
138    });
139    until("the burst to arrive", || arrivals.lock().unwrap().len() == 2).await;
140    // Let the third settle into its wait, rather than racing admission.
141    tokio::time::sleep(Duration::from_millis(500)).await;
142
143    let (other, conn) = tokio_postgres::connect(&pg.conn_str(), tokio_postgres::NoTls).await.unwrap();
144    tokio::spawn(conn);
145    let ok: bool = other.query_one("SELECT pg_cancel_backend($1)", &[&pid]).await.unwrap().get(0);
146    assert!(ok);
147
148    let (client, result) = query.await.unwrap();
149    let err = result.expect_err("cancelled");
150    let db = err.as_db_error().expect("an ERROR");
151    assert_eq!(db.code(), &SqlState::QUERY_CANCELED);
152    // statement_timeout (60 s) shares the SQLSTATE; the message says the
153    // cancel ended the wait, not the timeout.
154    assert_eq!(db.message(), "canceling statement due to user request");
155    assert_eq!(arrivals.lock().unwrap().len(), 2, "the waiting request was never sent");
156    client.query_one("SELECT 1", &[]).await.expect("the session is fine");
157}
158
159/// Seeds `jev_cache` with one answered request of no characters and
160/// `tokens` input tokens above the request overhead, so the learned
161/// characters per token fall far below the 2.6 fallback.
162async fn seed_ratio(client: &tokio_postgres::Client, tokens: i64) {
163    client
164        .execute(
165            "INSERT INTO jev_cache (key, model, answered_model, namespace, layout, state, question,
166                                    input_tokens, output_tokens, answer)
167             VALUES ('\\x00', 'jev-1.13.0', 'jev-1.13.0', 'seed', 'row_as_state', '', '', $1, 0, '{}')",
168            &[&(267 + tokens)],
169        )
170        .await
171        .unwrap();
172}
173
174#[tokio::test(flavor = "multi_thread")]
175async fn admission_prices_a_request_at_the_learned_ratio() {
176    let arrivals: Arrivals = Arc::default();
177    let seen = arrivals.clone();
178    let mock = MockJev::start(move |_| {
179        seen.lock().unwrap().push(Instant::now());
180        Reply::json(200, noul(0.5))
181    })
182    .await;
183    // A second's burst of 100,000 tokens: two requests at 2.6 characters
184    // per token fit it many times over.
185    let (_pg, client) = jev_instance_with(
186        &mock,
187        &[("jev.max_requests_per_minute", "0"), ("jev.max_tokens_per_second", "100000")],
188    )
189    .await;
190    rows(&client, "r", 2).await;
191    let gap = |from: usize| {
192        let times = arrivals.lock().unwrap();
193        times[from + 1].duration_since(times[from])
194    };
195
196    client.query("SELECT jev_prob(t, 'Urgent?') FROM r t", &[]).await.unwrap();
197    assert!(gap(0) < Duration::from_millis(500), "at the fallback the pair waited {:?}", gap(0));
198
199    // Learned: each body is now hundreds of thousands of tokens, so the
200    // second waits out the first's charge before it is admitted.
201    seed_ratio(&client, 1_000_000).await;
202    client.batch_execute("SET jev.cache_namespace = 'again'").await.unwrap();
203    client.query("SELECT jev_prob(t, 'Urgent?') FROM r t", &[]).await.unwrap();
204    assert_eq!(arrivals.lock().unwrap().len(), 4);
205    assert!(gap(2) >= Duration::from_secs(2), "at the learned ratio the pair waited only {:?}", gap(2));
206}