ratelimit.rsThe account's rate limit is shared cluster-wide and enforced before sending (contract Execution, "The account's rate limit is shared cluster-wide"). Every arrival is timed at the mock, so the limit is measured where the account would bill it, not inferred from the queries' durations.
GCRA admits one period of burst (postjevsql-core/src/gcra.rs): at
jev.max_requests_per_minute = 60 about 61 requests go at once, then
one a second.
Every request's arrival at the mock, in order.
19type Arrivals = Arc<Mutex<Vec<Instant>>>;
No shared_preload_libraries: the limiter's words are in a named DSM
segment, attached when the first scan begins.
Distinct rows, so neither the cache nor dedupe folds any together.
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}
Seeds jev_cache with one answered request of no characters and
tokens input tokens above the request overhead, so the learned
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}
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}