From 51802d07cf43c516dd0f17dbd88a77e51f64879c Mon Sep 17 00:00:00 2001 From: yunyaozhou Date: Tue, 4 Aug 2026 04:40:58 +0800 Subject: [PATCH] feat: revalidate workspace snapshots with ETags --- README.md | 5 +- pyproject.toml | 2 +- sub2api_quota_tui.py | 118 ++++++++++++++++-- tests/test_payload.py | 277 ++++++++++++++++++++++++++++++++++++++++-- uv.lock | 2 +- 5 files changed, 383 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 63b13af..51686a0 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,10 @@ `shusub2` 是 Sub2API 聚合运维 TUI。默认模式只读取 server2 Pricing Monitor 后台生成的内存 snapshot，不直接访问 Accounts helper、`sub2api-status`、Sub2API admin API、PostgreSQL、SSH 或 pricing upstream。每次数据刷新只请求一次 workspace -projection；Dashboard、Accounts、Sources、Requests 和 Errors 共用同一份进程内缓存。后台 snapshot +projection；Dashboard、Accounts、Sources、Requests 和 Errors 共用同一份进程内缓存。缓存保留 +服务端 workspace weak ETag；TTL 到期或按 `r` 后只做条件 GET，未变时 bodyless `304` 会续期既有 +已验证 snapshot，且不会附加 `refresh=1` 或触发服务端采集。workspace URL 必须直接指向最终 +端点；客户端拒绝重定向，以免把内存 validator 带往其它 origin。后台 snapshot 读取之外，默认启动仍会进行一次短超时的 Gitea 版本检查；使用 `--no-version-check` 或 `SHUSUB2_NO_VERSION_CHECK=1` 可关闭它。 diff --git a/pyproject.toml b/pyproject.toml index 9a10d4f..797da32 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "shusub2" -version = "0.3.3" +version = "0.3.4" description = "Aggregated operations TUI for Sub2API" readme = "README.md" requires-python = ">=3.11" diff --git a/sub2api_quota_tui.py b/sub2api_quota_tui.py index 3fde0b8..85e15a8 100644 --- a/sub2api_quota_tui.py +++ b/sub2api_quota_tui.py @@ -16,6 +16,7 @@ import subprocess import sys import threading import time +import urllib.error import urllib.parse import urllib.request import zlib @@ -26,7 +27,7 @@ from zoneinfo import ZoneInfo, ZoneInfoNotFoundError APP_NAME = "shusub2" -FALLBACK_VERSION = "0.3.3" +FALLBACK_VERSION = "0.3.4" DEFAULT_WORKSPACE_URL = "https://price.tailbeb9ad.ts.net/api/ui-data?view=workspace" DEFAULT_WORKSPACE_URL_CONFIG_FILE = "~/.config/shusub2/workspace-url" DEFAULT_API_URL = "http://127.0.0.1:18318/api/tui/accounts" @@ -61,6 +62,39 @@ MAX_WORKSPACE_REQUESTS = 100 MAX_WORKSPACE_ERROR_EVENTS = 256 MAX_WORKSPACE_KEYS = 256 USAGE_TIMEZONE = "Asia/Shanghai" + + +class WorkspaceNotModified(RuntimeError): + """A conditional workspace GET confirmed the validated snapshot is unchanged.""" + + def __init__(self, etag: str = "") -> None: + super().__init__("workspace payload not modified") + self.etag = str(etag or "").strip() + + +class WorkspaceNoRedirectHandler(urllib.request.HTTPRedirectHandler): + """Keep workspace validators confined to the configured endpoint origin.""" + + def redirect_request( + self, + request: urllib.request.Request, + response: Any, + code: int, + message: str, + headers: Any, + new_url: str, + ) -> None: + return None + + +class WorkspacePayload(dict[str, Any]): + """Validated workspace JSON with an opaque HTTP validator kept out of the payload.""" + + def __init__(self, payload: dict[str, Any], etag: str = "") -> None: + super().__init__(payload) + self.etag = str(etag or "").strip() + + MONITOR_OK_STATUSES = {"operational", "ok", "success"} MONITOR_FAILED_STATUSES = {"error", "failed", "failure"} MONITOR_STOPWORDS = {"response", "responses", "monitor"} @@ -586,12 +620,48 @@ def workspace_period_label(traffic: Any) -> str: return f"today {period['date']} ({period['timezone']})" -def fetch_workspace_payload(workspace_url: str, timeout: int) -> dict[str, Any]: - payload = fetch_payload( +def fetch_workspace_response( + workspace_url: str, + timeout: int, + *, + if_none_match: str = "", +) -> tuple[dict[str, Any], str]: + """Fetch one workspace representation or signal a bodyless 304 revalidation.""" + validator = str(if_none_match or "").strip() + headers = {"Accept": "application/json", "Accept-Encoding": "gzip"} + if validator: + headers["If-None-Match"] = validator + request = urllib.request.Request(workspace_url, headers=headers) + opener = urllib.request.build_opener( + urllib.request.ProxyHandler({}), + WorkspaceNoRedirectHandler(), + ) + try: + with opener.open(request, timeout=timeout) as response: + payload = decode_json_response( + response, + "Pricing Monitor workspace did not return a JSON object", + maximum_bytes=MAX_WORKSPACE_JSON_BYTES, + ) + etag = str(response.headers.get("ETag") or "").strip() + except urllib.error.HTTPError as exc: + if exc.code == 304: + headers = exc.headers or {} + raise WorkspaceNotModified(str(headers.get("ETag") or validator)) from exc + raise + return payload, etag + + +def fetch_workspace_payload( + workspace_url: str, + timeout: int, + *, + if_none_match: str = "", +) -> WorkspacePayload: + payload, etag = fetch_workspace_response( workspace_url, timeout, - maximum_bytes=MAX_WORKSPACE_JSON_BYTES, - use_proxy=False, + if_none_match=if_none_match, ) traffic = payload.get("traffic") if ( @@ -635,7 +705,7 @@ def fetch_workspace_payload(workspace_url: str, timeout: int) -> dict[str, Any]: isinstance(source, dict) for source in payload["sources"] ): raise RuntimeError("Pricing Monitor returned invalid workspace sources") - return payload + return WorkspacePayload(payload, etag) def workspace_accounts_payload(payload: dict[str, Any]) -> dict[str, Any]: @@ -781,12 +851,16 @@ class WorkspaceCache: self.retry_seconds = min(30, self.ttl_seconds) self.error = "" self.network_fetches = 0 + self.etag = "" self.lock = threading.RLock() + def _has_current_payload(self) -> bool: + return bool(self.payload) and workspace_traffic_is_current(self.payload) + def get(self, *, force: bool = False) -> dict[str, Any]: with self.lock: now = time.monotonic() - current_payload = bool(self.payload) and workspace_traffic_is_current(self.payload) + current_payload = self._has_current_payload() if not force: if self.error and self.last_attempt_at and now - self.last_attempt_at < self.retry_seconds: if current_payload: @@ -800,19 +874,41 @@ class WorkspaceCache: return copy.deepcopy(self.payload) self.last_attempt_at = now self.network_fetches += 1 + if not current_payload: + self.etag = "" + request_etag = self.etag try: - payload = fetch_workspace_payload(self.url, self.timeout) + if request_etag: + payload = fetch_workspace_payload( + self.url, + self.timeout, + if_none_match=request_etag, + ) + else: + payload = fetch_workspace_payload(self.url, self.timeout) + except WorkspaceNotModified as unchanged: + if not self._has_current_payload(): + self.etag = "" + self.error = "workspace current-day traffic unavailable" + raise RuntimeError(self.error) from unchanged + self.etag = unchanged.etag or request_etag + self.fetched_at = time.monotonic() + self.error = "" + return copy.deepcopy(self.payload) except Exception as exc: self.error = "workspace unavailable" - if current_payload: + if self._has_current_payload(): return copy.deepcopy(self.payload) + self.etag = "" raise RuntimeError(self.error) from exc if not workspace_traffic_is_current(payload): self.error = "workspace current-day traffic unavailable" - if current_payload: + if self._has_current_payload(): return copy.deepcopy(self.payload) + self.etag = "" raise RuntimeError(self.error) - self.payload = copy.deepcopy(payload) + self.payload = copy.deepcopy(dict(payload)) + self.etag = str(getattr(payload, "etag", "") or "").strip() self.fetched_at = time.monotonic() self.error = "" return copy.deepcopy(self.payload) diff --git a/tests/test_payload.py b/tests/test_payload.py index 2c665a9..ad78ac0 100644 --- a/tests/test_payload.py +++ b/tests/test_payload.py @@ -13,6 +13,7 @@ import tempfile import threading import time import unittest +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from unittest import mock @@ -844,7 +845,7 @@ class WorkspaceTests(unittest.TestCase): mod = load_module() nonfinite = workspace_payload_fixture() nonfinite["traffic"]["requests"][0]["actual_cost"] = float("nan") - with mock.patch.object(mod, "fetch_payload", return_value=nonfinite): + with mock.patch.object(mod, "fetch_workspace_response", return_value=(nonfinite, "")): with self.assertRaisesRegex(RuntimeError, "invalid workspace projection"): mod.fetch_workspace_payload("https://workspace.example.test/data", 3) @@ -852,13 +853,13 @@ class WorkspaceTests(unittest.TestCase): oversized["traffic"]["requests"] = [ {"id": str(index)} for index in range(mod.MAX_WORKSPACE_REQUESTS + 1) ] - with mock.patch.object(mod, "fetch_payload", return_value=oversized): + with mock.patch.object(mod, "fetch_workspace_response", return_value=(oversized, "")): with self.assertRaisesRegex(RuntimeError, "invalid workspace requests"): mod.fetch_workspace_payload("https://workspace.example.test/data", 3) invalid_period = workspace_payload_fixture() invalid_period["traffic"].pop("ends_at") - with mock.patch.object(mod, "fetch_payload", return_value=invalid_period): + with mock.patch.object(mod, "fetch_workspace_response", return_value=(invalid_period, "")): with self.assertRaisesRegex(RuntimeError, "invalid workspace traffic period"): mod.fetch_workspace_payload("https://workspace.example.test/data", 3) @@ -874,28 +875,290 @@ class WorkspaceTests(unittest.TestCase): "ends_at": (utc_start + dt.timedelta(days=1)).isoformat(), } ) - with mock.patch.object(mod, "fetch_payload", return_value=utc_period): + with mock.patch.object(mod, "fetch_workspace_response", return_value=(utc_period, "")): with self.assertRaisesRegex(RuntimeError, "invalid workspace traffic period"): mod.fetch_workspace_payload("https://workspace.example.test/data", 3) incomplete_usage = workspace_payload_fixture() incomplete_usage["traffic"]["usage_rollups_full_day"] = False - with mock.patch.object(mod, "fetch_payload", return_value=incomplete_usage): + with mock.patch.object(mod, "fetch_workspace_response", return_value=(incomplete_usage, "")): with self.assertRaisesRegex(RuntimeError, "incomplete workspace traffic usage"): mod.fetch_workspace_payload("https://workspace.example.test/data", 3) noncanonical_request = workspace_payload_fixture() noncanonical_request["traffic"]["requests"][0]["instance"] = "server4" - with mock.patch.object(mod, "fetch_payload", return_value=noncanonical_request): + with mock.patch.object(mod, "fetch_workspace_response", return_value=(noncanonical_request, "")): with self.assertRaisesRegex(RuntimeError, "noncanonical workspace requests"): mod.fetch_workspace_payload("https://workspace.example.test/data", 3) missing_key_provenance = workspace_payload_fixture() missing_key_provenance["traffic"]["keys"][0].pop("instance") - with mock.patch.object(mod, "fetch_payload", return_value=missing_key_provenance): + with mock.patch.object(mod, "fetch_workspace_response", return_value=(missing_key_provenance, "")): with self.assertRaisesRegex(RuntimeError, "noncanonical workspace keys"): mod.fetch_workspace_payload("https://workspace.example.test/data", 3) + def test_workspace_conditional_response_sends_validator_and_accepts_bodyless_304(self) -> None: + mod = load_module() + validator = 'W/"workspace-v1"' + requests = [] + + class NotModifiedOpener: + def open(self, request, timeout): + requests.append((request, timeout)) + raise mod.urllib.error.HTTPError( + request.full_url, + 304, + "Not Modified", + {"ETag": validator}, + None, + ) + + with mock.patch.object(mod.urllib.request, "build_opener", return_value=NotModifiedOpener()): + with self.assertRaises(mod.WorkspaceNotModified) as raised: + mod.fetch_workspace_response( + "https://workspace.example.test/api/ui-data?view=workspace", + 3, + if_none_match=validator, + ) + + self.assertEqual(raised.exception.etag, validator) + self.assertEqual(len(requests), 1) + request, timeout = requests[0] + self.assertEqual(timeout, 3) + self.assertEqual( + request.full_url, + "https://workspace.example.test/api/ui-data?view=workspace", + ) + self.assertEqual(request.get_header("Accept-encoding"), "gzip") + self.assertEqual(request.get_header("If-none-match"), validator) + + with mock.patch.object( + mod, + "fetch_workspace_response", + return_value=(workspace_payload_fixture(), validator), + ): + verified = mod.fetch_workspace_payload( + "https://workspace.example.test/api/ui-data?view=workspace", 3 + ) + self.assertEqual(verified.etag, validator) + + def test_workspace_response_rejects_redirects_without_forwarding_validators(self) -> None: + mod = load_module() + validator = 'W/"workspace-v1"' + origin_validators: list[str] = [] + redirected_validators: list[str] = [] + + class RedirectTargetHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: # noqa: N802 + redirected_validators.append(self.headers.get("If-None-Match", "")) + self.send_response(200) + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: # noqa: A003 + return + + target = ThreadingHTTPServer(("127.0.0.1", 0), RedirectTargetHandler) + target_thread = threading.Thread(target=target.serve_forever, daemon=True) + target_thread.start() + self.addCleanup(target.server_close) + self.addCleanup(lambda: target_thread.join(timeout=2)) + self.addCleanup(target.shutdown) + target_host, target_port = target.server_address[:2] + + class RedirectOriginHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: # noqa: N802 + origin_validators.append(self.headers.get("If-None-Match", "")) + self.send_response(302) + self.send_header("Location", f"http://{target_host}:{target_port}/redirected") + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: # noqa: A003 + return + + origin = ThreadingHTTPServer(("127.0.0.1", 0), RedirectOriginHandler) + origin_thread = threading.Thread(target=origin.serve_forever, daemon=True) + origin_thread.start() + self.addCleanup(origin.server_close) + self.addCleanup(lambda: origin_thread.join(timeout=2)) + self.addCleanup(origin.shutdown) + origin_host, origin_port = origin.server_address[:2] + + with self.assertRaises(mod.urllib.error.HTTPError) as redirected: + mod.fetch_workspace_response( + f"http://{origin_host}:{origin_port}/workspace", + 3, + if_none_match=validator, + ) + + self.assertEqual(redirected.exception.code, 302) + self.assertEqual(origin_validators, [validator]) + self.assertEqual(redirected_validators, []) + + def test_workspace_cache_revalidates_current_payload_with_etag(self) -> None: + mod = load_module() + initial = workspace_payload_fixture() + changed = workspace_payload_fixture() + changed["accounts"]["totals"]["today_requests"] = 3 + initial_validator = 'W/"workspace-v1"' + changed_validator = 'W/"workspace-v2"' + cache = mod.WorkspaceCache("https://workspace.example.test/api/ui-data?view=workspace", 3, 300) + + with mock.patch.object( + mod, + "fetch_workspace_payload", + side_effect=[ + mod.WorkspacePayload(initial, initial_validator), + mod.WorkspaceNotModified(initial_validator), + mod.WorkspacePayload(changed, changed_validator), + ], + ) as fetch: + first = cache.get() + revalidated = cache.get(force=True) + replaced = cache.get(force=True) + + self.assertEqual(first, revalidated) + self.assertIsNot(first, revalidated) + self.assertEqual(replaced["accounts"]["totals"]["today_requests"], 3) + self.assertEqual(cache.etag, changed_validator) + self.assertEqual(cache.error, "") + self.assertEqual(cache.network_fetches, 3) + self.assertEqual( + fetch.call_args_list, + [ + mock.call(cache.url, 3), + mock.call(cache.url, 3, if_none_match=initial_validator), + mock.call(cache.url, 3, if_none_match=initial_validator), + ], + ) + + def test_workspace_cache_drops_validator_outside_current_day_and_fails_closed_on_304(self) -> None: + mod = load_module() + validator = 'W/"workspace-v1"' + cache = mod.WorkspaceCache("https://workspace.example.test/api/ui-data?view=workspace", 3, 300) + expired = workspace_payload_fixture() + zone = dt.timezone(dt.timedelta(hours=8), name="Asia/Shanghai") + local_start = dt.datetime.combine( + dt.datetime.now(zone).date() - dt.timedelta(days=1), + dt.time.min, + tzinfo=zone, + ) + expired["traffic"].update( + { + "date": local_start.date().isoformat(), + "started_at": local_start.isoformat(), + "ends_at": (local_start + dt.timedelta(days=1)).isoformat(), + } + ) + cache.payload = expired + cache.etag = validator + + with mock.patch.object( + mod, + "fetch_workspace_payload", + return_value=mod.WorkspacePayload(workspace_payload_fixture(), ""), + ) as fetch: + cache.get(force=True) + cache.get(force=True) + + self.assertEqual(fetch.call_args_list, [mock.call(cache.url, 3), mock.call(cache.url, 3)]) + self.assertEqual(cache.etag, "") + + empty = mod.WorkspaceCache(cache.url, 3, 300) + empty.etag = validator + with mock.patch.object( + mod, + "fetch_workspace_payload", + side_effect=mod.WorkspaceNotModified(validator), + ) as fetch: + with self.assertRaisesRegex(RuntimeError, "current-day traffic unavailable"): + empty.get(force=True) + fetch.assert_called_once_with(empty.url, 3) + self.assertEqual(empty.etag, "") + + def test_workspace_cache_rechecks_current_day_after_network_completion(self) -> None: + mod = load_module() + validator = 'W/"workspace-v1"' + + def cached_workspace() -> object: + cache = mod.WorkspaceCache( + "https://workspace.example.test/api/ui-data?view=workspace", 3, 300 + ) + cache.payload = workspace_payload_fixture() + cache.etag = validator + return cache + + not_modified_cache = cached_workspace() + with ( + mock.patch.object( + mod, + "fetch_workspace_payload", + side_effect=mod.WorkspaceNotModified(validator), + ) as fetch, + mock.patch.object(mod, "workspace_traffic_is_current", side_effect=[True, False]), + ): + with self.assertRaisesRegex(RuntimeError, "current-day traffic unavailable"): + not_modified_cache.get(force=True) + fetch.assert_called_once_with( + not_modified_cache.url, 3, if_none_match=validator + ) + self.assertEqual(not_modified_cache.etag, "") + + failed_cache = cached_workspace() + with ( + mock.patch.object( + mod, + "fetch_workspace_payload", + side_effect=RuntimeError("private failure"), + ) as fetch, + mock.patch.object(mod, "workspace_traffic_is_current", side_effect=[True, False]), + ): + with self.assertRaisesRegex(RuntimeError, "workspace unavailable"): + failed_cache.get(force=True) + fetch.assert_called_once_with(failed_cache.url, 3, if_none_match=validator) + self.assertEqual(failed_cache.etag, "") + + invalid = workspace_payload_fixture() + invalid["traffic"]["usage_rollups_full_day"] = False + invalid_cache = cached_workspace() + with ( + mock.patch.object( + mod, + "fetch_workspace_response", + return_value=(invalid, 'W/"workspace-v2"'), + ) as fetch, + mock.patch.object(mod, "workspace_traffic_is_current", side_effect=[True, False]), + ): + with self.assertRaisesRegex(RuntimeError, "workspace unavailable"): + invalid_cache.get(force=True) + fetch.assert_called_once_with(invalid_cache.url, 3, if_none_match=validator) + self.assertEqual(invalid_cache.etag, "") + + def test_workspace_cache_retains_last_good_validator_when_a_200_is_invalid(self) -> None: + mod = load_module() + validator = 'W/"workspace-v1"' + invalid = workspace_payload_fixture() + invalid["traffic"]["usage_rollups_full_day"] = False + cache = mod.WorkspaceCache("https://workspace.example.test/api/ui-data?view=workspace", 3, 300) + + with mock.patch.object( + mod, + "fetch_workspace_response", + side_effect=[ + (workspace_payload_fixture(), validator), + (invalid, 'W/"workspace-v2"'), + ], + ) as fetch: + first = cache.get() + stale = cache.get(force=True) + + self.assertEqual(stale, first) + self.assertEqual(cache.etag, validator) + self.assertEqual(cache.error, "workspace unavailable") + self.assertEqual( + fetch.call_args_list[1], + mock.call(cache.url, 3, if_none_match=validator), + ) + def test_workspace_cache_coalesces_reads_and_returns_copies(self) -> None: mod = load_module() calls: list[str] = [] diff --git a/uv.lock b/uv.lock index 37389b2..39a04f3 100644 --- a/uv.lock +++ b/uv.lock @@ -85,7 +85,7 @@ wheels = [ [[package]] name = "shusub2" -version = "0.3.3" +version = "0.3.4" source = { editable = "." } dependencies = [ { name = "textual" },