From a28beb7c1ad063487fd465cee42e210cedc13a07 Mon Sep 17 00:00:00 2001 From: Tsonglew Date: Sat, 10 Oct 2026 18:15:44 +0800 Subject: [PATCH] fix: move usage persistence off the inference response path --- README.md | 23 ++++ README.zh-CN.md | 18 +++ backend/server/database.py | 16 ++- backend/server/main.py | 32 +++-- backend/server/usage_writer.py | 130 +++++++++++++++++++ backend/tests/conftest.py | 21 +++ backend/tests/test_flow.py | 38 +++--- backend/tests/test_inference_max_len.py | 48 ++++++- backend/tests/test_usage_writer.py | 166 ++++++++++++++++++++++++ 9 files changed, 449 insertions(+), 43 deletions(-) create mode 100644 backend/server/usage_writer.py create mode 100644 backend/tests/conftest.py create mode 100644 backend/tests/test_usage_writer.py diff --git a/README.md b/README.md index 2a2175c..fba9a7f 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,29 @@ curl -X POST http://YOUR_SERVER_IP:8080/v1/systemone \ Successful responses include `answers`, `model`, and actual token counts in `usage`. The published image includes the multilingual model; `model=auto` (the default) and `model=multilingual` are available. See **Documentation** in the console for more request examples. +## Usage statistics and storage + +Inference responses enqueue usage statistics without waiting for SQLite writes. A single +background thread per process reuses a connection, writes batches of up to 100 events, +and coalesces `last_used_at` updates by API key. Partial batches flush after up to one +second when the writer is idle; slow storage can delay statistics further. The `usage` +field in each inference response is unchanged. `duration_ms` measures model prediction +and result validation, not queue time or database writes. + +Statistics are best-effort, eventually consistent telemetry, not a billing or quota +ledger. Each process buffers up to 1,000 queued events plus one in-flight batch. A full +queue drops new statistics with rate-limited warnings. SQLite failures roll back where +possible, log and drop the affected batch, then reconnect for later batches; failed +batches are not retried because an uncertain commit could double-count usage. Neither +condition turns successful inference into an HTTP 503. API-key authentication and +revocation still use the database synchronously. + +Graceful shutdown attempts to drain the queue for up to five seconds. An in-flight +SQLite call cannot be cancelled safely and may outlive that wait; forced termination, +a shutdown timeout, or storage errors can lose unflushed statistics. Keep `/data` on +low-latency persistent storage: background batching reduces write pressure but does +not fix slow authentication reads or database locking. + ## License and Feedback This project is licensed under [Apache License 2.0](LICENSE). For issues and suggestions, visit [GitHub Issues](https://github.com/1Panel-dev/laya-server/issues). diff --git a/README.zh-CN.md b/README.zh-CN.md index ebadb8f..d9aa071 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -48,3 +48,21 @@ curl -X POST http://YOUR_SERVER_IP:8080/v1/systemone \ ## 许可与反馈 本项目采用 [Apache License 2.0](LICENSE) 许可证。问题和建议请提交至 [GitHub Issues](https://github.com/1Panel-dev/laya-server/issues)。 + +## 用量统计与存储 + +推理成功后将统计放入内存队列即可返回，不等待 SQLite 写入。每个进程使用一个后台 +线程复用连接，每批最多写入 100 条记录，并合并同一个 API key 的 `last_used_at` +更新。写入线程空闲时，不满一批的记录最多等待 1 秒开始写入；存储缓慢时统计会进一步 +延迟。响应中的 `usage` 字段保持不变，`duration_ms` 只计算模型调用及结果校验时间， +不包含排队和数据库写入时间。 + +统计采用尽力记录、最终一致的方式，不适合作为计费或严格配额账本。每个进程最多缓存 +1,000 条待写记录，另有一个正在处理的批次。队列满时丢弃新统计并限频记录警告。 +SQLite 写入失败时尽可能回滚，记录错误并丢弃该批次，后续批次重新连接；不重试失败 +批次，以免提交结果不确定时重复计数。这些情况不会让成功的推理返回 HTTP 503。 +API key 验证及撤销仍同步访问数据库。 + +正常停机时最多等待 5 秒清空队列。已经执行中的 SQLite 调用无法安全取消，可能超过 +该等待时间；强制终止、停机超时或存储错误可能导致尚未保存的统计丢失。`/data` 仍应 +使用低延迟的持久化存储，后台批量写入不能解决鉴权读取缓慢或数据库锁竞争问题。 diff --git a/backend/server/database.py b/backend/server/database.py index c39864c..0279f82 100644 --- a/backend/server/database.py +++ b/backend/server/database.py @@ -13,12 +13,20 @@ class Database: def __init__(self, path: Path): self.path = path + def open_connection(self) -> sqlite3.Connection: + connection = sqlite3.connect(self.path, timeout=10) + try: + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA foreign_keys=ON") + connection.execute("PRAGMA busy_timeout=10000") + except sqlite3.Error: + connection.close() + raise + return connection + @contextmanager def connect(self) -> Iterator[sqlite3.Connection]: - connection = sqlite3.connect(self.path, timeout=10) - connection.row_factory = sqlite3.Row - connection.execute("PRAGMA foreign_keys=ON") - connection.execute("PRAGMA busy_timeout=10000") + connection = self.open_connection() try: yield connection connection.commit() diff --git a/backend/server/main.py b/backend/server/main.py index b6ad6d0..9d44287 100644 --- a/backend/server/main.py +++ b/backend/server/main.py @@ -2,9 +2,9 @@ import hmac import logging import secrets -import sqlite3 import threading import time +from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone from typing import Annotated, Any @@ -13,11 +13,13 @@ from fastapi import Depends, FastAPI, Header, HTTPException, Request, Response from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse +from starlette.concurrency import run_in_threadpool from .config import Settings from .database import Database, utc_now from .laya_adapter import LayaAdapter, Predictor from .schemas import CreateKeyRequest, InferenceRequest, LoginRequest +from .usage_writer import UsageEvent, UsageWriter logger = logging.getLogger(__name__) SESSION_COOKIE = "laya_session" @@ -36,7 +38,18 @@ def create_app(settings: Settings | None = None, predictor: Predictor | None = N settings = settings or Settings.from_env() db = Database(settings.database_path) db.initialize() - app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None) + usage_writer = UsageWriter(db) + + @asynccontextmanager + async def lifespan(app: FastAPI): + usage_writer.start() + try: + yield + finally: + # Joining a slow writer must not block the event loop. + await run_in_threadpool(usage_writer.stop) + + app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None, lifespan=lifespan) hasher = PasswordHasher() router_holder: dict[str, Predictor] = {} router_lock = threading.Lock() @@ -129,17 +142,10 @@ def infer(payload: InferenceRequest, source: str, key_id: int | None) -> dict[st logger.error("Laya inference failed: %s", type(exc).__name__) raise error(503, "MODEL_UNAVAILABLE", "Model is unavailable") duration_ms = round((time.monotonic() - started) * 1000) - try: - with db.connect() as connection: - connection.execute( - "INSERT INTO usage_events (created_at,source,key_id,model,input_tokens,output_tokens,duration_ms) VALUES (?,?,?,?,?,?,?)", - (utc_now(), source, key_id, result["model"], usage["input_tokens"], usage["output_tokens"], duration_ms), - ) - if key_id is not None: - connection.execute("UPDATE api_keys SET last_used_at=? WHERE id=?", (utc_now(), key_id)) - except sqlite3.Error: - logger.exception("Failed to persist usage") - raise error(503, "USAGE_UNAVAILABLE", "Usage could not be recorded") + usage_writer.submit(UsageEvent( + utc_now(), source, key_id, result["model"], + usage["input_tokens"], usage["output_tokens"], duration_ms, + )) return result @app.get("/health/live") diff --git a/backend/server/usage_writer.py b/backend/server/usage_writer.py new file mode 100644 index 0000000..d0ce956 --- /dev/null +++ b/backend/server/usage_writer.py @@ -0,0 +1,130 @@ +"""Bounded, best-effort usage telemetry; never part of inference success.""" +import logging +import sqlite3 +import threading +import time +from collections import deque +from dataclasses import dataclass + +from .database import Database + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class UsageEvent: + created_at: str + source: str + key_id: int | None + model: str + input_tokens: int + output_tokens: int + duration_ms: int + + +class UsageWriter: + def __init__(self, db: Database, *, capacity: int = 1000, batch_size: int = 100, + flush_interval: float = 1.0, shutdown_timeout: float = 5.0): + if min(capacity, batch_size, flush_interval, shutdown_timeout) <= 0: + raise ValueError("Usage writer limits must be positive") + self.db = db + self.capacity = capacity + self.batch_size = batch_size + self.flush_interval = flush_interval + self.shutdown_timeout = shutdown_timeout + self._condition = threading.Condition() + self._pending: deque[UsageEvent] = deque() + self._thread: threading.Thread | None = None + self._accepting = False + self._deadline: float | None = None + self.dropped = 0 + + def start(self) -> None: + with self._condition: + if self._thread is not None and self._thread.is_alive(): + raise RuntimeError("Previous usage writer is still running") + self._deadline = None + self._accepting = True + self._thread = threading.Thread(target=self._run, name="laya-usage-writer", daemon=True) + self._thread.start() + + def submit(self, event: UsageEvent) -> bool: + with self._condition: + if not self._accepting or len(self._pending) >= self.capacity: + self.dropped += 1 + # Avoid flooding logs when storage cannot keep up. + if self.dropped == 1 or self.dropped & (self.dropped - 1) == 0: + logger.warning("Usage queue unavailable/full; dropped events=%d", self.dropped) + return False + self._pending.append(event) + self._condition.notify() + return True + + def stop(self) -> None: + with self._condition: + self._accepting = False + self._deadline = time.monotonic() + self.shutdown_timeout + self._condition.notify_all() + if self._thread is not None: + self._thread.join(self.shutdown_timeout) + if self._thread.is_alive(): + logger.warning("Usage shutdown timed out; unflushed statistics may be lost") + + def _next_batch(self) -> list[UsageEvent]: + with self._condition: + while not self._pending and self._accepting: + self._condition.wait() + flush_at = time.monotonic() + self.flush_interval + while self._accepting and len(self._pending) < self.batch_size: + remaining = flush_at - time.monotonic() + if remaining <= 0: + break + self._condition.wait(remaining) + if self._deadline is not None and time.monotonic() >= self._deadline: + self.dropped += len(self._pending) + self._pending.clear() + return [self._pending.popleft() for _ in range(min(self.batch_size, len(self._pending)))] + + @staticmethod + def _write_batch(connection: sqlite3.Connection, batch: list[UsageEvent]) -> None: + latest: dict[int, str] = {} + for event in batch: + if event.key_id is not None: + latest[event.key_id] = max(latest.get(event.key_id, ""), event.created_at) + with connection: + connection.executemany( + "INSERT INTO usage_events (created_at,source,key_id,model,input_tokens,output_tokens,duration_ms) " + "VALUES (?,?,?,?,?,?,?)", + [(e.created_at, e.source, e.key_id, e.model, e.input_tokens, e.output_tokens, e.duration_ms) + for e in batch], + ) + connection.executemany( + "UPDATE api_keys SET last_used_at=MAX(COALESCE(last_used_at,''),?) WHERE id=?", + [(created_at, key_id) for key_id, created_at in latest.items()], + ) + + def _run(self) -> None: + connection = None + try: + while batch := self._next_batch(): + try: + if connection is None: + connection = self.db.open_connection() + self._write_batch(connection, batch) + except sqlite3.Error: + with self._condition: + self.dropped += len(batch) + logger.exception("Failed to persist %d usage events; dropping batch", len(batch)) + # Do not retry an ambiguously committed batch and double-count it. + if connection is not None: + connection.close() + connection = None + finally: + with self._condition: + self._accepting = False + if self._pending: + self.dropped += len(self._pending) + logger.warning("Usage writer stopped; discarding %d pending events", len(self._pending)) + self._pending.clear() + if connection is not None: + connection.close() diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py new file mode 100644 index 0000000..08b03dd --- /dev/null +++ b/backend/tests/conftest.py @@ -0,0 +1,21 @@ +import time +from contextlib import ExitStack + +import pytest +from fastapi.testclient import TestClient + + +@pytest.fixture +def client_factory(): + with ExitStack() as stack: + yield lambda *args, **kwargs: stack.enter_context(TestClient(*args, **kwargs)) + + +def wait_for_usage(client, requests): + deadline = time.monotonic() + 5 + while True: + usage = client.get('/internal/usage').json() + if usage['totals']['requests'] == requests: + return usage + assert time.monotonic() < deadline, usage + time.sleep(0.01) diff --git a/backend/tests/test_flow.py b/backend/tests/test_flow.py index b30bfcb..775a00c 100644 --- a/backend/tests/test_flow.py +++ b/backend/tests/test_flow.py @@ -5,7 +5,7 @@ import pytest from argon2 import PasswordHasher -from fastapi.testclient import TestClient +from conftest import wait_for_usage os.environ.setdefault("LAYA_ADMIN_USERNAME", "test-admin") os.environ.setdefault("LAYA_ADMIN_PASSWORD_HASH", PasswordHasher().hash("test-password")) @@ -41,7 +41,7 @@ def predict(self, state, questions, model=None): } -def test_multilingual_image_routes_auto_and_rejects_unbundled_models(tmp_path): +def test_multilingual_image_routes_auto_and_rejects_unbundled_models(tmp_path, client_factory): model = tmp_path / "models" / "multilingual" for filename in ("rl_agent_config.json", "model.safetensors", "tokenizer/tokenizer.json", "encoder/config.json"): path = model / filename @@ -50,7 +50,7 @@ def test_multilingual_image_routes_auto_and_rejects_unbundled_models(tmp_path): fake = FakePredictor() settings = Settings("admin", PasswordHasher().hash("correct horse battery staple"), tmp_path / "db.sqlite3", 1, tmp_path / "models", None, 1, tmp_path / "missing-dist", "multilingual") - client = TestClient(create_app(settings, fake)) + client = client_factory(create_app(settings, fake)) assert client.get("/health/ready").status_code == 200 assert client.get("/internal/models").status_code == 401 origin = {"Origin": "http://testserver"} @@ -68,7 +68,7 @@ def test_multilingual_image_routes_auto_and_rejects_unbundled_models(tmp_path): assert fake.models == ["multilingual", "multilingual"] -def test_admin_password_from_environment(tmp_path, monkeypatch): +def test_admin_password_from_environment(tmp_path, monkeypatch, client_factory): password = "0123456789" monkeypatch.setenv("LAYA_ADMIN_USERNAME", "configured-admin") monkeypatch.setenv("LAYA_ADMIN_PASSWORD", password) @@ -84,7 +84,7 @@ def test_admin_password_from_environment(tmp_path, monkeypatch): monkeypatch.delenv("LAYA_MODEL_PROFILE") assert settings.admin_password_hash != password assert settings.admin_password_hash.startswith("$argon2id$") - client = TestClient(create_app(settings, FakePredictor())) + client = client_factory(create_app(settings, FakePredictor())) response = client.post( "/internal/auth/login", json={"username": "configured-admin", "password": password}, @@ -95,7 +95,7 @@ def test_admin_password_from_environment(tmp_path, monkeypatch): assert client.post("/internal/api-keys", json={"name": "without-origin"}, headers={"X-CSRF-Token": csrf}).status_code == 201 - https_client = TestClient(create_app(settings, FakePredictor()), base_url="https://testserver") + https_client = client_factory(create_app(settings, FakePredictor()), base_url="https://testserver") secure_login = https_client.post("/internal/auth/login", json={"username": "configured-admin", "password": password}) assert secure_login.status_code == 200 assert all("; secure" in value.lower() for value in secure_login.headers.get_list("set-cookie")) @@ -117,12 +117,12 @@ def test_admin_password_from_environment(tmp_path, monkeypatch): Settings.from_env() -def test_login_key_inference_usage_revoke_logout(tmp_path): +def test_login_key_inference_usage_revoke_logout(tmp_path, client_factory): fake = FakePredictor() db_path = tmp_path / "db.sqlite3" settings = Settings("admin", PasswordHasher().hash("correct horse battery staple"), db_path, 1, tmp_path / "models", None, 1, tmp_path / "missing-dist") - client = TestClient(create_app(settings, fake)) + client = client_factory(create_app(settings, fake)) origin = {"Origin": "http://testserver"} assert client.get("/internal/api-keys").status_code == 401 @@ -171,10 +171,10 @@ def test_login_key_inference_usage_revoke_logout(tmp_path): assert prediction.json()["usage"] == {"input_tokens": 37, "output_tokens": 0} assert set(prediction.json()["answers"]) == set(PAYLOAD["questions"]) assert fake.calls == 1 - assert client.get("/internal/usage").json()["totals"] == {"requests": 1, "input_tokens": 37, "output_tokens": 0} + assert wait_for_usage(client, 1)["totals"] == {"requests": 1, "input_tokens": 37, "output_tokens": 0} debug = client.post("/internal/playground", json=PAYLOAD, headers=write_headers) assert debug.status_code == 200 - usage = client.get("/internal/usage").json() + usage = wait_for_usage(client, 2) assert usage["totals"] == {"requests": 2, "input_tokens": 74, "output_tokens": 0} assert {item["source"] for item in usage["sources"]} == {"api_key", "playground"} @@ -184,13 +184,13 @@ def test_login_key_inference_usage_revoke_logout(tmp_path): assert client.post("/internal/auth/logout", headers=write_headers).status_code == 200 assert client.get("/internal/api-keys").status_code == 401 - restarted = TestClient(create_app(settings, fake)) + restarted = client_factory(create_app(settings, fake)) assert restarted.post("/v1/systemone", json=PAYLOAD, headers={"Authorization": f"Bearer {key}"}).status_code == 401 with sqlite3.connect(db_path) as connection: assert connection.execute("SELECT COUNT(*) FROM usage_events").fetchone()[0] == 2 -def test_rate_limit_expiry_and_model_failure(tmp_path): +def test_rate_limit_expiry_and_model_failure(tmp_path, client_factory): class UnavailablePredictor: def predict(self, state, questions, model=None): raise FileNotFoundError("secret internal model path") @@ -198,7 +198,7 @@ def predict(self, state, questions, model=None): db_path = tmp_path / "db.sqlite3" settings = Settings("admin", PasswordHasher().hash("password"), db_path, 1, tmp_path / "models", None, 1, tmp_path / "missing-dist") - client = TestClient(create_app(settings, UnavailablePredictor())) + client = client_factory(create_app(settings, UnavailablePredictor())) origin = {"Origin": "http://testserver"} assert client.get("/health/ready").json()["detail"]["code"] == "MODEL_UNAVAILABLE" assert client.get("/docs").status_code == 404 @@ -223,11 +223,11 @@ def predict(self, state, questions, model=None): @pytest.mark.parametrize("keep_cookies", [True, False], ids=["stale-cookies", "expired-cookies"]) -def test_login_after_session_expiry(tmp_path, keep_cookies): +def test_login_after_session_expiry(tmp_path, keep_cookies, client_factory): db_path = tmp_path / "db.sqlite3" settings = Settings("admin", PasswordHasher().hash("test-password"), db_path, 1, tmp_path / "models", None, 1, tmp_path / "missing-dist") - client = TestClient(create_app(settings, FakePredictor())) + client = client_factory(create_app(settings, FakePredictor())) credentials = {"username": "admin", "password": "test-password"} origin = {"Origin": "http://testserver"} assert client.post("/internal/auth/login", json=credentials, headers=origin).status_code == 200 @@ -260,7 +260,7 @@ def test_login_after_session_expiry(tmp_path, keep_cookies): assert [key["name"] for key in client.get("/internal/api-keys").json()] == ["after-login"] -def test_concurrent_inference_requests_are_not_rejected(tmp_path): +def test_concurrent_inference_requests_are_not_rejected(tmp_path, client_factory): class ConcurrentPredictor(FakePredictor): def __init__(self): super().__init__() @@ -273,7 +273,7 @@ def predict(self, state, questions, model=None): settings = Settings("admin", PasswordHasher().hash("test-password"), tmp_path / "db.sqlite3", 1, tmp_path / "models", None, 1, tmp_path / "missing-dist") app = create_app(settings, ConcurrentPredictor()) - client = TestClient(app) + client = client_factory(app) origin = {"Origin": "http://testserver"} assert client.post("/internal/auth/login", json={"username": "admin", "password": "test-password"}, headers=origin).status_code == 200 csrf = client.get("/internal/auth/session").json()["csrf_token"] @@ -284,13 +284,13 @@ def request(index: int): if index == 2: return client.post("/internal/playground", json=PAYLOAD, headers={**origin, "X-CSRF-Token": csrf}) - return TestClient(app).post("/v1/systemone", json=PAYLOAD, headers={"Authorization": f"Bearer {key}"}) + return client.post("/v1/systemone", json=PAYLOAD, headers={"Authorization": f"Bearer {key}"}) with ThreadPoolExecutor(max_workers=3) as executor: responses = list(executor.map(request, range(3))) assert [response.status_code for response in responses] == [200, 200, 200] - assert client.get("/internal/usage").json()["totals"]["requests"] == 3 + assert wait_for_usage(client, 3)["totals"]["requests"] == 3 with sqlite3.connect(settings.database_path) as connection: assert dict(connection.execute("SELECT source, COUNT(*) FROM usage_events GROUP BY source").fetchall()) == { "api_key": 2, "playground": 1, diff --git a/backend/tests/test_inference_max_len.py b/backend/tests/test_inference_max_len.py index bcbf70b..f5dbcae 100644 --- a/backend/tests/test_inference_max_len.py +++ b/backend/tests/test_inference_max_len.py @@ -3,7 +3,7 @@ import pytest from argon2 import PasswordHasher -from fastapi.testclient import TestClient +from conftest import wait_for_usage from server.config import Settings @@ -32,7 +32,7 @@ def password_hash(): @pytest.fixture -def inference_client(tmp_path, monkeypatch, password_hash): +def inference_client(tmp_path, monkeypatch, password_hash, client_factory): monkeypatch.setenv("LAYA_ADMIN_USERNAME", "admin") monkeypatch.setenv("LAYA_ADMIN_PASSWORD_HASH", password_hash) monkeypatch.delenv("LAYA_ADMIN_PASSWORD", raising=False) @@ -42,7 +42,7 @@ def inference_client(tmp_path, monkeypatch, password_hash): predictor = RecordingPredictor() settings = Settings("admin", password_hash, tmp_path / "db.sqlite3", 1, tmp_path / "models", None, 1, tmp_path / "missing-dist", "multilingual") - client = TestClient(create_app(settings, predictor)) + client = client_factory(create_app(settings, predictor)) assert client.post("/internal/auth/login", json={"username": "admin", "password": "test-password"}).status_code == 200 csrf = client.get("/internal/auth/session").json()["csrf_token"] playground_headers = {"X-CSRF-Token": csrf} @@ -50,7 +50,6 @@ def inference_client(tmp_path, monkeypatch, password_hash): headers = {"/v1/systemone": {"Authorization": f"Bearer {key}"}, "/internal/playground": playground_headers} yield client, predictor, headers - client.close() @pytest.mark.parametrize("endpoint", ["/v1/systemone", "/internal/playground"]) @@ -61,7 +60,7 @@ def test_max_len_valid_and_default_requests(inference_client, endpoint, options) assert response.status_code == 200, response.text expected = {} if options.get("max_len") is None else options assert predictor.calls == [(PAYLOAD["state"], "multilingual", expected)] - assert client.get("/internal/usage").json()["totals"] == { + assert wait_for_usage(client, 1)["totals"] == { "requests": 1, "input_tokens": 37, "output_tokens": 0, } @@ -93,7 +92,7 @@ def request(case): responses = list(executor.map(request, cases)) assert [response.status_code for response in responses] == [200, 200, 200] assert {state: overrides for state, model, overrides in predictor.calls} == dict(cases) - assert client.get("/internal/usage").json()["totals"]["requests"] == 3 + assert wait_for_usage(client, 3)["totals"]["requests"] == 3 @pytest.mark.parametrize("endpoint", ["/v1/systemone", "/internal/playground"]) @@ -105,4 +104,39 @@ def test_multi_question_usage_is_not_capped_by_max_len(inference_client, endpoin assert response.status_code == 200, response.text assert set(response.json()["answers"]) == set(questions) assert response.json()["usage"]["input_tokens"] == 600 - assert client.get("/internal/usage").json()["totals"]["input_tokens"] == 600 + assert wait_for_usage(client, 1)["totals"]["input_tokens"] == 600 + + +@pytest.mark.parametrize("endpoint", ["/v1/systemone", "/internal/playground"]) +@pytest.mark.parametrize("storage_error", [False, True], ids=["blocked-storage", "failed-storage"]) +def test_usage_storage_does_not_block_successful_inference(inference_client, monkeypatch, endpoint, storage_error): + import sqlite3 + from threading import Event + from server.usage_writer import UsageWriter + + entered = Event() + release = Event() + original = UsageWriter._write_batch + + def blocked_write(connection, batch): + entered.set() + assert release.wait(5) + if storage_error: + raise sqlite3.OperationalError("simulated slow storage failure") + original(connection, batch) + + monkeypatch.setattr(UsageWriter, '_write_batch', staticmethod(blocked_write)) + client, predictor, headers = inference_client + with ThreadPoolExecutor(max_workers=1) as pool: + try: + response = pool.submit(client.post, endpoint, json=PAYLOAD, headers=headers[endpoint]) + assert entered.wait(3) + # The writer is still blocked here: HTTP completion cannot depend on it. + result = response.result(timeout=1) + assert result.status_code == 200 + assert result.json()['usage']['input_tokens'] == 37 + assert client.post(endpoint, json=PAYLOAD, headers=headers[endpoint]).status_code == 200 + if endpoint == '/v1/systemone': + assert client.post(endpoint, json=PAYLOAD).status_code == 401 + finally: + release.set() diff --git a/backend/tests/test_usage_writer.py b/backend/tests/test_usage_writer.py new file mode 100644 index 0000000..7d15b5d --- /dev/null +++ b/backend/tests/test_usage_writer.py @@ -0,0 +1,166 @@ +import sqlite3 +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from server.database import Database +from server.usage_writer import UsageEvent, UsageWriter + + +def event(key_id=None, created_at='2026-10-10T09:00:00+00:00'): + return UsageEvent(created_at, 'api_key' if key_id else 'playground', key_id, 'test', 37, 0, 25) + + +@pytest.fixture +def db(tmp_path): + database = Database(tmp_path / 'usage.sqlite3') + database.initialize() + return database + + +def rows(db): + with db.connect() as connection: + return connection.execute('SELECT * FROM usage_events').fetchall() + + +def test_batches_coalesce_key_updates_and_reuse_connection(db, monkeypatch): + with db.connect() as connection: + connection.execute("INSERT INTO api_keys(id,name,digest,mask,created_at) VALUES (1,'test','hash','mask','now')") + connections = [] + statements = [] + original = db.open_connection + + def open_connection(): + connection = original() + connections.append(connection) + connection.set_trace_callback(statements.append) + return connection + + monkeypatch.setattr(db, 'open_connection', open_connection) + writer = UsageWriter(db, batch_size=3, flush_interval=60) + writer.start() + # Out-of-order timestamps must not move last_used_at backwards. + for stamp in ['09:00:02', '09:00:01', '09:00:03', '09:00:00']: + assert writer.submit(event(1, f'2026-10-10T{stamp}+00:00')) + writer.stop() + assert not writer._thread.is_alive() + assert len(connections) == 1 + assert sum(s == 'COMMIT' for s in statements) == 2 + assert sum(s.startswith('UPDATE api_keys') for s in statements) == 2 + assert len(rows(db)) == 4 + with db.connect() as connection: + assert connection.execute('SELECT last_used_at FROM api_keys').fetchone()[0] == '2026-10-10T09:00:03+00:00' + + +def test_partial_batch_flushes_without_shutdown(db): + writer = UsageWriter(db, flush_interval=0.02) + writer.start() + try: + assert writer.submit(event()) + deadline = time.monotonic() + 3 + while not rows(db): + assert time.monotonic() < deadline + time.sleep(0.01) + finally: + writer.stop() + assert len(rows(db)) == 1 + + +def test_blocked_storage_bounds_queue_and_shutdown(db, monkeypatch, caplog): + entered = threading.Event() + release = threading.Event() + + def blocked_write(connection, batch): + entered.set() + assert release.wait(5) + + writer = UsageWriter(db, capacity=2, batch_size=1, shutdown_timeout=0.05) + monkeypatch.setattr(writer, '_write_batch', blocked_write) + writer.start() + try: + assert writer.submit(event()) + assert entered.wait(3) + with ThreadPoolExecutor(max_workers=4) as pool: + accepted = list(pool.map(lambda _: writer.submit(event()), range(20))) + assert sum(accepted) == 2 + assert writer.dropped == 18 + writer.stop() + assert writer._thread.is_alive() # stop returned despite blocked I/O + assert not writer.submit(event()) + with pytest.raises(RuntimeError, match='still running'): + writer.start() + assert 'shutdown timed out' in caplog.text + finally: + release.set() + writer._thread.join(3) + assert not writer._thread.is_alive() + assert writer.dropped == 21 # 18 overflow, one after stop, two expired + + +def test_failed_transaction_rolls_back_and_later_batch_recovers(db, monkeypatch, caplog): + failed = threading.Event() + original = UsageWriter._write_batch + attempts = 0 + + def fail_first(connection, batch): + nonlocal attempts + attempts += 1 + if attempts == 1: + try: + with connection: + connection.execute("INSERT INTO usage_events(created_at,source,model,input_tokens,output_tokens,duration_ms) VALUES ('now','playground','test',1,0,1)") + raise sqlite3.OperationalError('simulated storage failure') + finally: + failed.set() + original(connection, batch) + + writer = UsageWriter(db, batch_size=1) + monkeypatch.setattr(writer, '_write_batch', fail_first) + writer.start() + try: + writer.submit(event()) + assert failed.wait(3) + writer.submit(event()) + finally: + writer.stop() + assert len(rows(db)) == 1 + assert writer.dropped == 1 + assert 'dropping batch' in caplog.text + + +def test_writer_can_restart_after_clean_shutdown(db): + writer = UsageWriter(db) + for _ in range(2): + writer.start() + assert writer.submit(event()) + writer.stop() + assert len(rows(db)) == 2 + + +def test_connection_failure_does_not_kill_worker(db, monkeypatch, caplog): + original = db.open_connection + failed = threading.Event() + attempts = 0 + + def fail_first_open(): + nonlocal attempts + attempts += 1 + if attempts == 1: + failed.set() + raise sqlite3.OperationalError('database unavailable') + return original() + + monkeypatch.setattr(db, 'open_connection', fail_first_open) + writer = UsageWriter(db, batch_size=1) + writer.start() + try: + writer.submit(event()) + assert failed.wait(3) + writer.submit(event()) + finally: + writer.stop() + assert len(rows(db)) == 1 + assert writer.dropped == 1 + assert 'dropping batch' in caplog.text