diff --git a/scripts/record_daily_account_snapshot.py b/scripts/record_daily_account_snapshot.py index 1b97ce4..1a841e9 100644 --- a/scripts/record_daily_account_snapshot.py +++ b/scripts/record_daily_account_snapshot.py @@ -81,6 +81,32 @@ }) _PRODUCER_WRITE_TIMEOUT_SECONDS = 20.0 _PRODUCER_DECIMAL = re.compile(r"^-?(?:0|[1-9]\d*)(?:\.\d+)?$") +_SNAPSHOT_DIAGNOSTIC_PHASES = frozenset({ + "config", "scheduler_read", "scheduler_validate", "resume", "resume_readback", + "before_run", "scheduler_run", "after_run", "before_store_open", "store_open", + "wait", "list", "read", "validate", "selected", "after_observation", "publish_check", "publish", +}) +_SNAPSHOT_DIAGNOSTIC_COUNTERS = ( + "polls", "list_calls", "listed", "metadata_rejected", "candidates", "reads", + "read_rejected", "read_failed", "object_invalid", "object_truncated", + "candidate_timeout", "generation_changed", "other_rejected", "validation_rejected", "validated", +) +_SNAPSHOT_REJECTION_COUNTERS = { + "gcs_read_failed": "read_failed", "gcs_object_invalid": "object_invalid", + "gcs_object_truncated": "object_truncated", "observation_timeout": "candidate_timeout", + "generation_changed": "generation_changed", +} + + +def _snapshot_diagnostic(diagnostics, *, phase=None, counter=None, amount=1): + # Fixed enums and capped aggregate observations only; never retain provider input. + if diagnostics is None: + return + if phase is not None: + diagnostics["phase"] = phase if phase in _SNAPSHOT_DIAGNOSTIC_PHASES else "other" + if counter is not None: + key = counter if counter in _SNAPSHOT_DIAGNOSTIC_COUNTERS else "other_rejected" + diagnostics[key] = min(9999, diagnostics.get(key, 0) + amount) @dataclass(frozen=True) @@ -500,6 +526,7 @@ def record_daily_account_snapshot( now_reader: Callable[[], datetime], monotonic: Callable[[], float] = time.monotonic, sleep: Callable[[float], None] = time.sleep, + diagnostics: dict[str, Any] | None = None, ) -> DailyAccountRecordResult: """Validate, trigger once, select a fresh same-source object, and POST once.""" @@ -508,6 +535,9 @@ def record_daily_account_snapshot( mode = env.get("ACCOUNT_HISTORY_OBSERVATION_MODE") or "scheduler_archive" if mode != "scheduler_archive": return DailyAccountRecordResult("error", "independent_entrypoint_required" if mode == "independent_get" else "observation_mode_invalid") + if diagnostics is not None: + diagnostics.clear() + diagnostics.update(dict.fromkeys(_SNAPSHOT_DIAGNOSTIC_COUNTERS, 0), phase="config") try: config = _config(env) sync_config = _account_facts_sync_config(env) @@ -515,6 +545,7 @@ def record_daily_account_snapshot( return DailyAccountRecordResult("error", rejected.category, "rejected", rejected.category) session = None + _snapshot_diagnostic(diagnostics, phase="scheduler_read") try: session = session_factory() job = _scheduler_get(session, config) @@ -528,6 +559,7 @@ def record_daily_account_snapshot( paused_resumable = ( config.target_id in _PAUSED_PROBE_CONTRACTS and job.get("state") == "PAUSED" ) + _snapshot_diagnostic(diagnostics, phase="scheduler_validate") try: if paused_resumable: if env.get("RUNTIME_TARGET_ENABLED") != "false": @@ -540,6 +572,7 @@ def record_daily_account_snapshot( job, now_reader, monotonic=monotonic, + diagnostics=diagnostics, observe=lambda started_at, deadline: _observe_archived_snapshot( open_store, config, @@ -548,11 +581,13 @@ def record_daily_account_snapshot( now_reader=now_reader, monotonic=monotonic, sleep=sleep, + diagnostics=diagnostics, ), ) else: _validate_scheduler_job(job, config) started_at = _utc_now(now_reader) + _snapshot_diagnostic(diagnostics, phase="scheduler_run") _scheduler_run(session, config) except _Rejected as rejected: return DailyAccountRecordResult("error", rejected.category) @@ -572,6 +607,7 @@ def record_daily_account_snapshot( now_reader=now_reader, monotonic=monotonic, sleep=sleep, + diagnostics=diagnostics, ) except _Rejected as rejected: return DailyAccountRecordResult("error", rejected.category) @@ -583,8 +619,10 @@ def record_daily_account_snapshot( body, record = selected if sync_config is None: return DailyAccountRecordResult("recorded", publish_status="disabled") + _snapshot_diagnostic(diagnostics, phase="publish_check") if not _is_fresh(record, _utc_now(now_reader)): return DailyAccountRecordResult("recorded", publish_status="rejected", publish_category="observation_stale") + _snapshot_diagnostic(diagnostics, phase="publish") publish_status, publish_category = _publish_account_facts( sync_config, body, record, http_post=http_post ) @@ -927,10 +965,12 @@ def _run_paused_probe( *, monotonic: Callable[[], float], observe: Callable[[datetime, float], tuple[bytes, dict[str, Any]] | None], + diagnostics: dict[str, Any] | None = None, ) -> tuple[bytes, dict[str, Any]] | None: original = deepcopy(dict(original_job)) original_last_attempt = original.get("lastAttemptTime") deadline = monotonic() + WAIT_SECONDS + _snapshot_diagnostic(diagnostics, phase="resume") try: _scheduler_action(session, config, "resume") except _Rejected as rejected: @@ -946,6 +986,7 @@ def _run_paused_probe( raise _Rejected("scheduler_resume_state_unknown") raise rejected + _snapshot_diagnostic(diagnostics, phase="resume_readback") try: resumed = _read_scheduler_job(session, config) except Exception: @@ -965,15 +1006,18 @@ def _run_paused_probe( raise _Rejected("scheduler_resume_readback_mismatch") try: + _snapshot_diagnostic(diagnostics, phase="before_run") if monotonic() >= deadline: raise _Rejected("observation_timeout") started_at = _utc_now(now_reader) try: + _snapshot_diagnostic(diagnostics, phase="scheduler_run") _scheduler_run(session, config) except _Rejected: raise except Exception: raise _Rejected("scheduler_run_unknown") from None + _snapshot_diagnostic(diagnostics, phase="after_run") if monotonic() >= deadline: raise _Rejected("observation_timeout") return observe(started_at, deadline) @@ -1019,11 +1063,14 @@ def _observe_archived_snapshot( now_reader: Callable[[], datetime], monotonic: Callable[[], float], sleep: Callable[[float], None], + diagnostics: dict[str, Any] | None = None, ) -> tuple[bytes, dict[str, Any]] | None: + _snapshot_diagnostic(diagnostics, phase="before_store_open") if deadline is not None and monotonic() >= deadline: return None try: store = open_store(config.project_id) + _snapshot_diagnostic(diagnostics, phase="store_open") if deadline is None: deadline = monotonic() + WAIT_SECONDS if monotonic() >= deadline: @@ -1037,8 +1084,10 @@ def _observe_archived_snapshot( now_reader=now_reader, monotonic=monotonic, sleep=sleep, + diagnostics=diagnostics, ) if monotonic() >= deadline: + _snapshot_diagnostic(diagnostics, phase="after_observation" if selected is not None else None) return None return selected except _Rejected: @@ -1063,38 +1112,54 @@ def _wait_for_observation( now_reader: Callable[[], datetime], monotonic: Callable[[], float], sleep: Callable[[float], None], + diagnostics: dict[str, Any] | None = None, ) -> tuple[bytes, dict[str, Any]] | None: while True: + _snapshot_diagnostic(diagnostics, phase="wait") if monotonic() >= deadline: return None now = _utc_now(now_reader) + _snapshot_diagnostic(diagnostics, counter="polls") dates = {started_at.date(), now.date()} candidates = _list_candidates( - client, config, dates, deadline=deadline, monotonic=monotonic + client, config, dates, deadline=deadline, monotonic=monotonic, diagnostics=diagnostics ) if monotonic() >= deadline: return None valid: list[tuple[datetime, bytes, dict[str, Any]]] = [] for candidate in candidates: + candidate_phase = "read" + _snapshot_diagnostic(diagnostics, phase="read", counter="reads") try: payload, raw = _read_candidate( client, config, candidate, deadline=deadline, monotonic=monotonic ) + candidate_phase = "validate" + _snapshot_diagnostic(diagnostics, phase="validate") record = _validate_history_object(payload, config, candidate, started_at, now) except _Rejected as rejected: + _snapshot_diagnostic( + diagnostics, counter="validation_rejected" if candidate_phase == "validate" else "read_rejected", + ) + _snapshot_diagnostic( + diagnostics, counter=_SNAPSHOT_REJECTION_COUNTERS.get(rejected.category, "other_rejected"), + ) if rejected.category == "generation_changed": raise continue + _snapshot_diagnostic(diagnostics, counter="validated") if monotonic() >= deadline: return None valid.append((datetime.fromisoformat(record["observed_finished_at"]), raw, record)) if valid: _finished, raw, record = max(valid, key=lambda entry: entry[0]) if monotonic() < deadline and _is_fresh(record, _utc_now(now_reader)): + _snapshot_diagnostic(diagnostics, phase="selected") return raw, record remaining = deadline - monotonic() if remaining <= 0: return None + _snapshot_diagnostic(diagnostics, phase="wait") sleep(min(POLL_SECONDS, remaining)) @@ -1105,13 +1170,16 @@ def _list_candidates( *, deadline: float, monotonic: Callable[[], float], + diagnostics: dict[str, Any] | None = None, ) -> list[dict[str, Any]]: + _snapshot_diagnostic(diagnostics, phase="list") candidates: list[dict[str, Any]] = [] for day in sorted(dates): remaining = deadline - monotonic() if remaining <= 0: raise _Rejected("observation_timeout") prefix = f"{config.prefix_path}/{config.target_id}/{config.source_binding_id}/{day.isoformat()}/" + _snapshot_diagnostic(diagnostics, counter="list_calls") try: blobs = client.list_blobs( config.bucket, @@ -1134,6 +1202,7 @@ def _list_candidates( raise _Rejected("gcs_list_failed") from None if monotonic() >= deadline: raise _Rejected("observation_timeout") + _snapshot_diagnostic(diagnostics, counter="listed", amount=len(page)) if has_next_page or len(page) > MAX_OBJECTS_PER_DAY: raise _Rejected("gcs_listing_truncated") for blob in page: @@ -1150,10 +1219,12 @@ def _list_candidates( or not isinstance(size, int) or size < 0 ): + _snapshot_diagnostic(diagnostics, counter="metadata_rejected") continue if size > MAX_OBJECT_BYTES: raise _Rejected("gcs_object_too_large") candidates.append({"name": name, "generation": int(generation), "size": size}) + _snapshot_diagnostic(diagnostics, counter="candidates") if len(candidates) > MAX_OBJECTS_PER_DAY * 2: raise _Rejected("gcs_listing_truncated") return candidates @@ -1821,16 +1892,20 @@ def main( if args: print("error: config_invalid") return 1 + diagnostics: dict[str, Any] = {} result = record_daily_account_snapshot( os.environ if environ is None else environ, open_store=open_store or _open_store, session_factory=session_factory or _authorized_session, http_post=http_post or _default_http_post, now_reader=now_reader or (lambda: datetime.now(timezone.utc)), + diagnostics=diagnostics, ) record_part = f"record=error:{result.category}" if result.status == "error" else f"record={result.status}" publish_part = f"account_facts_publish={result.publish_status}" print(f"{record_part} {publish_part}") + if diagnostics: + print("snapshot_diagnostics=" + json.dumps(diagnostics, separators=(",", ":"), sort_keys=True)) if result.status == "error" or result.publish_status in {"rejected", "unknown"}: return 1 return 0 diff --git a/tests/test_daily_account_snapshot.py b/tests/test_daily_account_snapshot.py index 4bbc8f0..2a030a1 100644 --- a/tests/test_daily_account_snapshot.py +++ b/tests/test_daily_account_snapshot.py @@ -369,7 +369,7 @@ def close(self): self.closed = True -def _record(env=None, spies=None, *, times=None, monotonic=None, sleep=None): +def _record(env=None, spies=None, *, times=None, monotonic=None, sleep=None, diagnostics=None): spies = spies or _Spies() now_reader = _Clock(times or [T0, T0 + timedelta(seconds=5)]) clock = {"now": 0.0} @@ -391,6 +391,7 @@ def sleep_fn(seconds): now_reader=now_reader, monotonic=monotonic_fn, sleep=sleep_fn, + diagnostics=diagnostics, ) return result, spies @@ -1084,11 +1085,180 @@ def test_qrs_unknown_is_reported_once_and_secrets_are_redacted(capsys): ) output = capsys.readouterr().out assert code == 1 - assert output.strip() == "record=recorded account_facts_publish=unknown" + lines = output.splitlines() + assert lines[0] == "record=recorded account_facts_publish=unknown" + assert json.loads(lines[1].removeprefix("snapshot_diagnostics="))["phase"] == "publish" assert SECRET not in output and QRS_TOKEN not in output assert len(spies.posts) == 1 +@pytest.mark.parametrize( + ("case", "counter", "expected_status", "expected_category"), + [ + ("empty", "listed", "error", "observation_timeout"), + ("valid", "validated", "recorded", ""), + ("metadata", "metadata_rejected", "error", "observation_timeout"), + ("invalid_json", "object_invalid", "error", "observation_timeout"), + ("invalid_contract", "validation_rejected", "error", "observation_timeout"), + ("read_failure", "read_failed", "error", "observation_timeout"), + ("generation", "generation_changed", "error", "generation_changed"), + ("unknown_category", "other_rejected", "error", "observation_timeout"), + ], +) +def test_snapshot_diagnostics_distinguish_synthetic_observation_outcomes( + monkeypatch, case, counter, expected_status, expected_category, +): + obj = _object(_history()) + if case == "metadata": + obj.generation = True + elif case == "invalid_json": + obj = _object(_history(), raw=("{" + SECRET).encode()) + elif case == "invalid_contract": + payload = _history() + payload["source_binding"]["id"] = OTHER_BINDING + obj = _object(_history(), raw=json.dumps(payload).encode()) + elif case == "read_failure": + obj.error = RuntimeError(SECRET) + elif case == "generation": + obj.error = RuntimeError("generation changed " + SECRET) + elif case == "unknown_category": + def reject(*args, **kwargs): + raise snapshots._Rejected(SECRET) + monkeypatch.setattr(snapshots, "_read_candidate", reject) + spies = _Spies(objects=[] if case == "empty" else [obj]) + diagnostics = {SECRET: SECRET} + result, spies = _record( + _env(ACCOUNT_FACTS_SYNC_ENABLED="true", ACCOUNT_FACTS_SYNC_URL=QRS_URL, ACCOUNT_FACTS_SYNC_TOKEN=QRS_TOKEN), + spies, diagnostics=diagnostics, + ) + assert (result.status, result.category) == (expected_status, expected_category) + assert len(spies.posts) == (1 if case == "valid" else 0) + assert diagnostics["polls"] > 0 + if case == "empty": + assert diagnostics["listed"] == diagnostics["reads"] == 0 + else: + assert diagnostics[counter] > 0 + if case == "read_failure": + assert diagnostics["read_rejected"] > 0 and diagnostics["validation_rejected"] == 0 + if case == "invalid_contract": + assert diagnostics["object_invalid"] > 0 and diagnostics["read_rejected"] == 0 + serialized = json.dumps(diagnostics) + for private in (SECRET, QRS_TOKEN, BINDING, OTHER_BINDING, PREFIX, SERVICE_URL, T0.isoformat(), "10.25"): + assert private not in serialized + assert set(diagnostics) == {"phase", *snapshots._SNAPSHOT_DIAGNOSTIC_COUNTERS} + assert all(type(value) is int and 0 <= value <= 9999 for key, value in diagnostics.items() if key != "phase") + + +@pytest.mark.parametrize("expired_at", ["resume", "run"]) +def test_snapshot_diagnostics_paused_deadline_phase_preserves_restore_and_run_count(expired_at): + ticks = {"now": 0.0} + spies = _Spies(job=_job("sg", state="PAUSED")) + original_post = spies.session.post + + def post(url, **kwargs): + response = original_post(url, **kwargs) + if url.endswith(":" + expired_at): + ticks["now"] = snapshots.WAIT_SECONDS + return response + + spies.session.post = post + diagnostics = {} + result, spies = _record( + _env("sg", RUNTIME_TARGET_ENABLED="false"), spies, + monotonic=lambda: ticks["now"], diagnostics=diagnostics, + ) + assert result.category == "observation_timeout" + assert diagnostics["phase"] == ("before_run" if expired_at == "resume" else "after_run") + actions = [call[1].rsplit(":", 1)[-1] for call in spies.session.calls if call[0] == "post"] + assert actions == (["resume", "pause"] if expired_at == "resume" else ["resume", "run", "pause"]) + assert spies.session.job["state"] == "PAUSED" + assert spies.events[-1] == "scheduler_get" + assert spies.open_calls == spies.posts == [] + + +@pytest.mark.parametrize("expired_at", ["open", "list", "read"]) +def test_snapshot_diagnostics_storage_deadline_phase_keeps_original_budget(expired_at): + ticks = {"now": 0.0} + obj = _object(_history(target_id="sg")) + spies = _Spies(objects=[obj], job=_job("sg", state="PAUSED")) + if expired_at == "open": + original_open = spies.open_store + + def slow_open(project): + storage = original_open(project) + ticks["now"] = snapshots.WAIT_SECONDS + return storage + + spies.open_store = slow_open + elif expired_at == "list": + original_list = spies.storage.list_blobs + + def slow_list(*args, **kwargs): + page = original_list(*args, **kwargs) + page.before_page = lambda: ticks.update(now=snapshots.WAIT_SECONDS) + return page + + spies.storage.list_blobs = slow_list + else: + original_read = obj.download_as_bytes + + def slow_read(**kwargs): + ticks["now"] = snapshots.WAIT_SECONDS + return original_read(**kwargs) + + obj.download_as_bytes = slow_read + diagnostics = {} + result, spies = _record( + _env("sg", RUNTIME_TARGET_ENABLED="false"), spies, + monotonic=lambda: ticks["now"], diagnostics=diagnostics, + ) + assert snapshots.WAIT_SECONDS == 180.0 + assert result.category == "observation_timeout" and spies.posts == [] + assert diagnostics["phase"] == {"open": "store_open", "list": "list", "read": "read"}[expired_at] + assert spies.session.job["state"] == "PAUSED" + if expired_at == "read": + assert diagnostics["candidate_timeout"] == 1 + else: + assert obj.reads == [] + + +@pytest.mark.parametrize("case", ["disabled", "empty", "valid", "generation"]) +def test_snapshot_diagnostics_cli_keeps_status_line_exit_and_private_output_boundary(monkeypatch, capsys, case): + spies = _Spies(objects=[] if case in {"disabled", "empty"} else [_object(_history())]) + if case == "generation": + spies.storage.objects[0].error = RuntimeError("generation " + SECRET) + ticks = {"now": 0.0} + original_record = snapshots.record_daily_account_snapshot + + def deterministic_record(*args, **kwargs): + return original_record( + *args, **kwargs, monotonic=lambda: ticks["now"], + sleep=lambda seconds: ticks.update(now=ticks["now"] + seconds), + ) + + monkeypatch.setattr(snapshots, "record_daily_account_snapshot", deterministic_record) + code = snapshots.main( + [], environ=_env(ACCOUNT_HISTORY_RECORDING_ENABLED="false" if case == "disabled" else "true"), + open_store=spies.open_store, session_factory=spies.session_factory, http_post=spies.http_post, + now_reader=_Clock([T0, T0 + timedelta(seconds=5)]), + ) + output = capsys.readouterr().out + lines = output.splitlines() + expected_status = {"disabled": "disabled", "empty": "error:observation_timeout", "valid": "recorded", "generation": "error:generation_changed"} + assert lines[0] == f"record={expected_status[case]} account_facts_publish=disabled" + assert code == int(case in {"empty", "generation"}) + if case == "disabled": + assert len(lines) == 1 and spies.session.calls == spies.open_calls == [] + else: + assert len(lines) == 2 and lines[1].startswith("snapshot_diagnostics=") + diagnostics = json.loads(lines[1].split("=", 1)[1]) + assert diagnostics["phase"] in snapshots._SNAPSHOT_DIAGNOSTIC_PHASES + assert len(lines[1]) < 600 + for private in (SECRET, QRS_TOKEN, BINDING, PREFIX, SERVICE_URL, T0.isoformat(), "10.25"): + assert private not in output + assert spies.posts == [] + + def test_authorized_session_disables_transport_retries(monkeypatch): import google.auth from google.auth.transport.requests import AuthorizedSession