Single-binary self-hosted market watcher for stocks, ETFs, indexes, and futures: live charts, key stats, fundamentals, SEC filings, and SSE streaming.
axumdockerfinancerustself-hostedsqlitestocksvite
1//! Persistent, per-endpoint request guard.
2//!
3//! Every outbound call to a third-party data endpoint (today only Stooq) passes
4//! through an `EndpointGuard`. The guard is this project's hard guarantee that a
5//! third-party rate limit can never be hit: not by a burst, not by a buggy
6//! loop, and not across process restarts.
7//!
8//! It combines three mechanisms, all backed by the `endpoint_guard` table so
9//! they survive restarts and are shared by every job and every process (the
10//! server and the `finance seed` subcommand both write the same row):
11//!
12//! - **A reactive circuit breaker.** A request that returns an explicit
13//! rate-limit signal (HTTP 429/503) opens the breaker immediately; so does a
14//! streak of ordinary failures. While open, every request is refused. Each
15//! consecutive trip backs off longer (30m, 1h, 2h, 4h ... capped at 24h).
16//! Once the backoff elapses the breaker goes *half-open* and lets exactly
17//! one probe through: it closes on success, re-opens (longer) on failure.
18//!
19//! - **A hard per-hour request budget.** At most `HOURLY_BUDGET` requests are
20//! let through per rolling clock hour. When the budget is spent, jobs are
21//! refused until the hour rolls. This caps a runaway loop even when the
22//! upstream never returns an error.
23//!
24//! - **Pacing.** Consecutive requests are spaced at least `MIN_GAP` apart;
25//! `acquire` sleeps as needed before granting, so callers no longer pace
26//! themselves.
27//!
28//! For a Python reader: picture one row of `endpoint_guard` as a small state
29//! machine persisted to disk. `acquire()` asks "may I send now?" (and blocks
30//! for pacing); `record_success` / `record_failure` feed the outcome back in.
31//!
32//! Usage in a bulk loop:
33//! ```ignore
34//! let guard = EndpointGuard::with_budget(pool.clone(), "yahoo", 1000);
35//! for ticker in tickers {
36//! match guard.acquire().await? {
37//! Permit::Granted => {}
38//! Permit::Denied(why) => break, // breaker open or budget spent: stop the run
39//! }
40//! match provider.daily(ticker, None).await {
41//! Ok(bars) => guard.record_success().await?,
42//! Err(err) => guard.record_failure(&err).await?,
43//! }
44//! }
45//! ```
46
47use std::collections::HashMap;
48use std::sync::{Arc, Mutex as StdMutex, OnceLock};
49use std::time::Duration;
50
51use sqlx::SqlitePool;
52
53use crate::db::now_ms;
54use crate::providers::RateLimited;
55
56/// Minimum spacing between two requests to the same endpoint. Matches the
57/// anti-spam policy (>= 1.5s per request). `acquire` enforces it.
58const MIN_GAP: Duration = Duration::from_millis(1500);
59
60/// Consecutive ordinary failures that trip the breaker while it is closed. An
61/// explicit rate-limit signal (429/503) trips it immediately, regardless.
62const FAIL_THRESHOLD: i64 = 4;
63
64/// A half-open probe that records no result within this long is treated as
65/// abandoned (its process likely crashed mid-probe), so a fresh probe is let
66/// through rather than wedging the breaker half-open forever.
67const STALE_PROBE_SECS: i64 = 10 * 60;
68
69const HOUR_MS: i64 = 3600 * 1000;
70
71/// One async mutex per endpoint name, shared by every `EndpointGuard`
72/// instance in the process (instances are constructed cheaply all over, so
73/// the lock cannot live on the struct).
74fn endpoint_lock(endpoint: &str) -> Arc<tokio::sync::Mutex<()>> {
75 static LOCKS: OnceLock<StdMutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>> =
76 OnceLock::new();
77 let map = LOCKS.get_or_init(|| StdMutex::new(HashMap::new()));
78 let mut m = map.lock().expect("endpoint lock map poisoned");
79 m.entry(endpoint.to_string()).or_default().clone()
80}
81
82/// The guard's verdict for one request attempt.
83pub enum Permit {
84 /// Cleared to send. Pacing has already been applied (this call slept as
85 /// needed) and the request has been counted against the hourly budget.
86 Granted,
87 /// The request must not be sent: the circuit breaker is open or the hourly
88 /// budget is spent. The caller should stop its run. The string explains
89 /// why and is suitable for a log line or a `fetch_log` detail.
90 Denied(String),
91}
92
93/// A persistent guard over one outbound data endpoint. Cheap to construct
94/// (`SqlitePool` is an `Arc` internally); all real state lives in the
95/// `endpoint_guard` row, so separate instances for the same endpoint stay
96/// consistent.
97pub struct EndpointGuard {
98 pool: SqlitePool,
99 endpoint: String,
100 /// This endpoint's hard per-hour request ceiling.
101 hourly_budget: i64,
102}
103
104/// The subset of an `endpoint_guard` row the guard logic reads back. The table
105/// carries more columns (`opened_at`, `last_ok_at`, `last_error`,
106/// `hourly_budget`, ...) for the data-health page; they are written here but
107/// not read here.
108#[derive(sqlx::FromRow)]
109struct GuardRow {
110 state: String,
111 fail_streak: i64,
112 trip_count: i64,
113 retry_at: Option<i64>,
114 hour_start: Option<i64>,
115 hour_count: i64,
116 last_request_at: Option<i64>,
117 updated_at: i64,
118}
119
120impl EndpointGuard {
121 /// A guard with an explicit per-hour request budget. Each endpoint sets its
122 /// own ceiling (e.g. 1000 for Yahoo, 600 for SEC) — see the constants in
123 /// `scheduler.rs`.
124 pub fn with_budget(pool: SqlitePool, endpoint: &str, hourly_budget: i64) -> Self {
125 Self {
126 pool,
127 endpoint: endpoint.to_string(),
128 hourly_budget,
129 }
130 }
131
132 /// Ensure this endpoint's `endpoint_guard` row exists and its persisted
133 /// `hourly_budget` matches this guard. Used at startup to register the
134 /// known endpoints, so the data-health page shows each one with its right
135 /// budget from boot rather than only after that endpoint's first request.
136 pub async fn ensure_registered(&self) -> anyhow::Result<()> {
137 self.load(now_ms()).await?;
138 Ok(())
139 }
140
141 /// Ask permission to send one request.
142 ///
143 /// On `Permit::Granted` the call has already slept for pacing and counted
144 /// the request against the hourly budget; the caller should send the
145 /// request straight away. On `Permit::Denied` the caller must not send and
146 /// should stop its run (the breaker is open or the budget is spent).
147 pub async fn acquire(&self) -> anyhow::Result<Permit> {
148 // Serialize the read-decide-sleep-commit sequence per endpoint. It is
149 // not atomic against the shared row: two concurrent acquires could
150 // both read the same `last_request_at`, sleep identical pacing, and
151 // send simultaneously (or both claim the single half-open probe
152 // slot). The DB still coordinates across processes (server vs seed);
153 // this closes the race between tasks inside one process.
154 let lock = endpoint_lock(&self.endpoint);
155 let _serialized = lock.lock().await;
156 let now = now_ms();
157 let g = self.load(now).await?;
158
159 // 1. Hard per-hour budget. Checked first: it holds even when nothing
160 // has failed, so it is the backstop against a runaway loop.
161 if g.hour_count >= self.hourly_budget {
162 let resets_in = human_secs((g.hour_start.unwrap_or(now) + HOUR_MS - now) / 1000);
163 return Ok(Permit::Denied(format!(
164 "{} hourly budget spent ({}/{}), resets in {resets_in}",
165 self.endpoint, g.hour_count, self.hourly_budget
166 )));
167 }
168
169 // 2. Circuit breaker. `probing` means this acquire is taking the single
170 // half-open probe slot.
171 let probing = match g.state.as_str() {
172 "open" => {
173 let retry_at = g.retry_at.unwrap_or(now);
174 if now < retry_at {
175 return Ok(Permit::Denied(format!(
176 "{} circuit breaker open, retry in {}",
177 self.endpoint,
178 human_secs((retry_at - now) / 1000)
179 )));
180 }
181 true // backoff elapsed: this caller becomes the half-open probe
182 }
183 "half_open" => {
184 // A probe is in flight. Allow a fresh one only if the previous
185 // probe looks abandoned (no result recorded for a long time).
186 if now - g.updated_at <= STALE_PROBE_SECS * 1000 {
187 return Ok(Permit::Denied(format!(
188 "{} circuit breaker half-open, probe in flight",
189 self.endpoint
190 )));
191 }
192 true
193 }
194 _ => false, // closed
195 };
196
197 // 3. Granted. Pace, then commit the bookkeeping.
198 if let Some(last) = g.last_request_at {
199 let wait = MIN_GAP.as_millis() as i64 - (now - last);
200 if wait > 0 {
201 tokio::time::sleep(Duration::from_millis(wait as u64)).await;
202 }
203 }
204 let sent = now_ms();
205 let new_state = if probing { "half_open" } else { g.state.as_str() };
206 sqlx::query(
207 "UPDATE endpoint_guard SET \
208 state = ?, hour_count = hour_count + 1, \
209 last_request_at = ?, updated_at = ? WHERE endpoint = ?",
210 )
211 .bind(new_state)
212 .bind(sent)
213 .bind(sent)
214 .bind(&self.endpoint)
215 .execute(&self.pool)
216 .await?;
217
218 Ok(Permit::Granted)
219 }
220
221 /// Record that the last `acquire`d request succeeded. Closes the breaker
222 /// and clears every failure counter.
223 pub async fn record_success(&self) -> anyhow::Result<()> {
224 let now = now_ms();
225 sqlx::query(
226 "UPDATE endpoint_guard SET \
227 state = 'closed', fail_streak = 0, trip_count = 0, \
228 opened_at = NULL, retry_at = NULL, \
229 last_ok_at = ?, updated_at = ? WHERE endpoint = ?",
230 )
231 .bind(now)
232 .bind(now)
233 .bind(&self.endpoint)
234 .execute(&self.pool)
235 .await?;
236 Ok(())
237 }
238
239 /// Record that the last `acquire`d request failed.
240 ///
241 /// The breaker trips (opens) when any of these holds: the error is an
242 /// explicit rate-limit signal ([`RateLimited`]); a half-open probe failed;
243 /// or the ordinary-failure streak reached `FAIL_THRESHOLD`. Otherwise the
244 /// streak is just incremented. A trip backs off exponentially, honouring a
245 /// `Retry-After` when it is longer than the computed backoff.
246 pub async fn record_failure(&self, err: &anyhow::Error) -> anyhow::Result<()> {
247 let now = now_ms();
248 let rate_limited = err.downcast_ref::<RateLimited>();
249 let g = self.load(now).await?;
250
251 let streak = g.fail_streak + 1;
252 let trip = rate_limited.is_some() // explicit upstream rate-limit signal
253 || g.state != "closed" // a half-open probe failed
254 || streak >= FAIL_THRESHOLD; // too many ordinary failures in a row
255 let msg = format!("{err:#}");
256
257 if trip {
258 let trip_count = g.trip_count + 1;
259 let mut backoff = backoff_secs(trip_count);
260 if let Some(ra) = rate_limited.and_then(|r| r.retry_after_secs) {
261 backoff = backoff.max(ra);
262 }
263 let retry_at = now + backoff * 1000;
264 sqlx::query(
265 "UPDATE endpoint_guard SET \
266 state = 'open', fail_streak = 0, trip_count = ?, \
267 opened_at = ?, retry_at = ?, \
268 last_error = ?, last_error_at = ?, updated_at = ? WHERE endpoint = ?",
269 )
270 .bind(trip_count)
271 .bind(now)
272 .bind(retry_at)
273 .bind(&msg)
274 .bind(now)
275 .bind(now)
276 .bind(&self.endpoint)
277 .execute(&self.pool)
278 .await?;
279 tracing::warn!(
280 "[guard] {} breaker OPEN (trip #{trip_count}), backoff {}: {msg}",
281 self.endpoint,
282 human_secs(backoff)
283 );
284 } else {
285 sqlx::query(
286 "UPDATE endpoint_guard SET \
287 fail_streak = ?, last_error = ?, last_error_at = ?, updated_at = ? \
288 WHERE endpoint = ?",
289 )
290 .bind(streak)
291 .bind(&msg)
292 .bind(now)
293 .bind(now)
294 .bind(&self.endpoint)
295 .execute(&self.pool)
296 .await?;
297 }
298 Ok(())
299 }
300
301 /// Load the guard row, creating a default one on first use and rolling the
302 /// per-hour budget window if the clock hour has elapsed. The hour roll is
303 /// persisted here (not just held in memory) so a later `hour_count + 1` is
304 /// always counting within the right hour.
305 async fn load(&self, now: i64) -> anyhow::Result<GuardRow> {
306 // Create the row on first use, and keep `hourly_budget` in step with
307 // how this guard was constructed — it differs per endpoint (see
308 // `with_budget`), and the data-health page reads it straight from here.
309 // `updated_at` is left untouched on the correcting update: it tracks
310 // state-machine changes, not routine bookkeeping.
311 sqlx::query(
312 "INSERT INTO endpoint_guard (endpoint, hourly_budget, updated_at) VALUES (?, ?, ?) \
313 ON CONFLICT(endpoint) DO UPDATE SET hourly_budget = excluded.hourly_budget \
314 WHERE endpoint_guard.hourly_budget <> excluded.hourly_budget",
315 )
316 .bind(&self.endpoint)
317 .bind(self.hourly_budget)
318 .bind(now)
319 .execute(&self.pool)
320 .await?;
321
322 // Roll the budget window. `updated_at` is deliberately left untouched:
323 // it tracks state-machine changes, and a half-open staleness check
324 // depends on it not being bumped by routine budget bookkeeping.
325 sqlx::query(
326 "UPDATE endpoint_guard SET hour_start = ?, hour_count = 0 \
327 WHERE endpoint = ? AND (hour_start IS NULL OR ? - hour_start >= ?)",
328 )
329 .bind(now)
330 .bind(&self.endpoint)
331 .bind(now)
332 .bind(HOUR_MS)
333 .execute(&self.pool)
334 .await?;
335
336 let row = sqlx::query_as::<_, GuardRow>(
337 "SELECT state, fail_streak, trip_count, retry_at, hour_start, \
338 hour_count, last_request_at, updated_at \
339 FROM endpoint_guard WHERE endpoint = ?",
340 )
341 .bind(&self.endpoint)
342 .fetch_one(&self.pool)
343 .await?;
344 Ok(row)
345 }
346}
347
348/// Backoff for the n-th consecutive trip: 30m, 1h, 2h, 4h, ... capped at 24h.
349fn backoff_secs(trip_count: i64) -> i64 {
350 const BASE: i64 = 30 * 60;
351 const CAP: i64 = 24 * 3600;
352 // trip_count is >= 1; clamp the shift so `1 << shift` cannot overflow.
353 let shift = (trip_count - 1).clamp(0, 16) as u32;
354 BASE.saturating_mul(1_i64 << shift).min(CAP)
355}
356
357/// A coarse, human-readable duration for log lines and status messages.
358fn human_secs(secs: i64) -> String {
359 let s = secs.max(0);
360 if s < 60 {
361 format!("{s}s")
362 } else if s < 3600 {
363 format!("{}m", s / 60)
364 } else {
365 format!("{}h{}m", s / 3600, (s % 3600) / 60)
366 }
367}