fix(poll): fix the api poll queue on large query

Author: kshitijrajsharma

api/duck.py

@@ -12,10 +12,11 @@
 import queue
 import threading
 import time
+from concurrent.futures import ThreadPoolExecutor
 from urllib.parse import urlparse
 
 import duckdb
-from litestar.exceptions import TooManyRequestsException
+from litestar.exceptions import HTTPException, TooManyRequestsException
 
 from osmsg import query
 from osmsg.db.schema import _apply_runtime_pragmas
@@ -115,15 +116,50 @@ def _acquire(pool: queue.Queue) -> duckdb.DuckDBPyConnection:
         raise TooManyRequestsException(detail="Server at capacity; retry shortly.") from None
 
 
+# A timed-out all-time query caches nothing, so its retry is just as slow; on timeout we re-run it in the
+# background to fill the cache the retry reads. One at a time on its own connection so it can't starve the pool.
+_warm_pool = ThreadPoolExecutor(max_workers=1)
+_warm_inflight: set[tuple[str, str]] = set()
+_warm_lock = threading.Lock()
+
+
+def _warm(fn, hashtag, kwargs, key) -> None:
+    con = _connect()
+    try:
+        fn(con, hashtag, _sources(), **kwargs)
+    except duckdb.Error:
+        pass
+    finally:
+        con.close()
+        with _warm_lock:
+            _warm_inflight.discard(key)
+
+
+def _enqueue_warm(fn, hashtag, kwargs) -> None:
+    """Re-run an all-time query off the request path to fill its cache, deduped per key, best-effort.
+    No-op when caching is off."""
+    if _QUERY_CACHE_DIR is None:
+        return
+    key = (fn.__name__, repr(hashtag))
+    with _warm_lock:
+        if key in _warm_inflight:
+            return
+        _warm_inflight.add(key)
+    _warm_pool.submit(_warm, fn, hashtag, kwargs, key)
+
+
 def _run(fn, hashtag, **kwargs):
     """Borrow a warm pooled connection (429 past _POOL_WAIT caps concurrency), run under a watchdog that
     interrupts a query past _QUERY_TIMEOUT, and return it (replaced if it errored or was interrupted)."""
     pool = _pool_ready()
     con = _acquire(pool)
     done = threading.Event()
+    interrupted = False
 
     def _watchdog() -> None:
+        nonlocal interrupted
         if not done.wait(_QUERY_TIMEOUT):
+            interrupted = True
             with contextlib.suppress(duckdb.Error):
                 con.interrupt()
 
@@ -134,6 +170,11 @@ def _watchdog() -> None:
         return fn(con, hashtag, _sources(), **kwargs)
     except duckdb.Error:
         healthy = False  # interrupted or DB error -> the connection may be dirty, recycle it
+        if interrupted:
+            # Only all-time queries hit the shared cache, so only they are worth warming.
+            if kwargs.get("start") is None and kwargs.get("end") is None:
+                _enqueue_warm(fn, hashtag, kwargs)
+            raise HTTPException(status_code=503, detail="Server is busy, please try again in a moment.") from None
         raise
     finally:
         done.set()

frontend/app.js

@@ -472,6 +472,7 @@ function sleep(ms, signal) {
     signal?.addEventListener("abort", () => { clearTimeout(t); reject(new DOMException("Aborted", "AbortError")); }, { once: true });
   });
 }
+// A 429 means the API shed load ("retry shortly"); retry once after a short pause before showing busy.
 async function apiGet(name, params, signal) {
   for (let attempt = 0; ; attempt++) {
     const res = await fetch(endpoint(name, params), { headers: { accept: "application/json" }, mode: "cors", signal });
@@ -484,6 +485,12 @@ async function apiGet(name, params, signal) {
       err.busy = true;
       throw err;
     }
+    if (res.status === 503) {
+      // Watchdog interrupted a large query; retrying now would just time out again, so surface busy.
+      const err = new Error("Server is busy");
+      err.busy = true;
+      throw err;
+    }
     if (!res.ok) throw new Error(`HTTP ${res.status} ${res.statusText || ""}`.trim());
     return res.json();
   }
@@ -627,6 +634,8 @@ async function runQuery() {
   } catch (err) {
     if (err?.name !== "AbortError") {
       console.warn("OSMSG query failed:", err);
+      // Surface a busy shed even when the first (summary) call is the one rejected, so the user is not
+      // left staring at loading skeletons.
       if (err?.busy && alive()) {
         showError(err);
         $("#ov-strip-totals").innerHTML = "";

tests/test_api_ratelimit.py

@@ -1,15 +1,34 @@
 """Load-shed and rate-limit keying for the API."""
 
 import queue
+import time
 from types import SimpleNamespace
 
+import duckdb
 import pytest
-from litestar.exceptions import TooManyRequestsException
+from litestar.exceptions import HTTPException, TooManyRequestsException
 
 from api import duck
 from api.app import _client_identifier, app, rate_limit_config
 
 
+class _FakeCon:
+    def interrupt(self):
+        pass
+
+    def close(self):
+        pass
+
+
+def _stub_pool(monkeypatch):
+    """Point _run at fake connections so it never touches DuckDB/Postgres."""
+    monkeypatch.setattr(duck, "_sources", lambda: None)
+    monkeypatch.setattr(duck, "_connect", lambda: _FakeCon())
+    pool: queue.Queue = queue.Queue()
+    pool.put(_FakeCon())
+    monkeypatch.setattr(duck, "_pool_ready", lambda: pool)
+
+
 def test_acquire_returns_free_connection():
     pool: queue.Queue = queue.Queue()
     sentinel = object()
@@ -48,3 +67,62 @@ def test_rate_limit_middleware_is_wired():
     assert [m.middleware for m in app.middleware] == [RateLimitMiddleware]
     assert rate_limit_config.rate_limit == ("minute", 120)
     assert rate_limit_config.identifier_for_request is _client_identifier
+
+
+def _fake_query(name):
+    def fn(*args, **kwargs):
+        return None
+
+    fn.__name__ = name
+    return fn
+
+
+def test_enqueue_warm_single_flight_dedups(monkeypatch):
+    monkeypatch.setattr(duck, "_QUERY_CACHE_DIR", "/tmp/qc")
+    submitted: list = []
+    monkeypatch.setattr(duck._warm_pool, "submit", lambda *a: submitted.append(a))
+    duck._warm_inflight.clear()
+    summary = _fake_query("summary")
+    allt = {"start": None, "end": None}
+    duck._enqueue_warm(summary, "hotosm", allt)
+    duck._enqueue_warm(summary, "hotosm", allt)  # same key -> dropped
+    duck._enqueue_warm(summary, "osmnepal", allt)  # different key -> submitted
+    assert len(submitted) == 2
+    assert len(duck._warm_inflight) == 2
+
+
+def test_enqueue_warm_noop_when_cache_disabled(monkeypatch):
+    monkeypatch.setattr(duck, "_QUERY_CACHE_DIR", None)
+    submitted: list = []
+    monkeypatch.setattr(duck._warm_pool, "submit", lambda *a: submitted.append(a))
+    duck._warm_inflight.clear()
+    duck._enqueue_warm(_fake_query("summary"), "hotosm", {"start": None, "end": None})
+    assert submitted == []
+    assert len(duck._warm_inflight) == 0
+
+
+def test_run_maps_interrupt_to_503(monkeypatch):
+    monkeypatch.setattr(duck, "_QUERY_TIMEOUT", 0.0)  # watchdog fires immediately -> interrupted
+    monkeypatch.setattr(duck, "_QUERY_CACHE_DIR", None)  # skip the warm side effect
+    _stub_pool(monkeypatch)
+
+    def boom(con, hashtag, sources, **kwargs):
+        time.sleep(0.05)  # let the zero-timeout watchdog set `interrupted` before we raise
+        raise duckdb.Error("interrupted")
+
+    boom.__name__ = "summary"
+    with pytest.raises(HTTPException) as ei:
+        duck._run(boom, "hotosm", start=None, end=None)
+    assert ei.value.status_code == 503
+
+
+def test_run_reraises_real_db_error_not_as_busy(monkeypatch):
+    monkeypatch.setattr(duck, "_QUERY_TIMEOUT", 100.0)  # watchdog never fires
+    _stub_pool(monkeypatch)
+
+    def boom(con, hashtag, sources, **kwargs):
+        raise duckdb.Error("real failure")
+
+    boom.__name__ = "summary"
+    with pytest.raises(duckdb.Error):  # a genuine error is a 500, not masked as 503 busy
+        duck._run(boom, "hotosm", start=None, end=None)