From 720cdd1d1440845957248d152e53f0ed890b2a05 Mon Sep 17 00:00:00 2001 From: kshitij Date: Wed, 29 Jul 2026 00:22:50 +0500 Subject: [PATCH] refactor: use atomic_json_write instead of hand-rolled _write_payload MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace 20 lines of manual os.open/O_EXCL/fdopen/fsync/os.replace with the existing atomic_json_write() from utils.py, which is already used by 6+ modules and handles temp-file creation, fsync, atomic replace, mode control, and owner preservation. The only novel helper (_fsync_directory) is retained — atomic_json_write does not do directory fsync. Update test_flush_write_failure_leaves_no_recovery_file to monkeypatch utils.os.replace (the new call path) instead of gateway.shutdown_flush.os.replace. --- gateway/shutdown_flush.py | 29 ++++++++-------------------- tests/gateway/test_shutdown_flush.py | 2 +- 2 files changed, 9 insertions(+), 22 deletions(-) diff --git a/gateway/shutdown_flush.py b/gateway/shutdown_flush.py index 3ffa74ea25644..e144c6ef43e7a 100644 --- a/gateway/shutdown_flush.py +++ b/gateway/shutdown_flush.py @@ -55,29 +55,16 @@ def _fsync_directory(path: Path) -> None: def _write_payload(flush_dir: Path, payload: Dict[str, Any]) -> None: """Atomically write one private, uniquely named recovery payload.""" + from utils import atomic_json_write + file_id = uuid.uuid4().hex final_path = flush_dir / f"pending-{file_id}.json" - temp_path = flush_dir / f".pending-{file_id}.tmp" - file_descriptor = -1 - - try: - file_descriptor = os.open( - temp_path, - os.O_WRONLY | os.O_CREAT | os.O_EXCL, - 0o600, - ) - with os.fdopen(file_descriptor, "w", encoding="utf-8") as handle: - file_descriptor = -1 # The file object now owns the descriptor. - json.dump(payload, handle, ensure_ascii=False, default=str) - handle.flush() - os.fsync(handle.fileno()) - - os.replace(temp_path, final_path) - except Exception: - if file_descriptor >= 0: - os.close(file_descriptor) - temp_path.unlink(missing_ok=True) - raise + atomic_json_write( + final_path, + payload, + mode=0o600, + default=str, + ) try: _fsync_directory(flush_dir) diff --git a/tests/gateway/test_shutdown_flush.py b/tests/gateway/test_shutdown_flush.py index 92b070964142f..a6f57b8137b0b 100644 --- a/tests/gateway/test_shutdown_flush.py +++ b/tests/gateway/test_shutdown_flush.py @@ -74,7 +74,7 @@ def test_flush_write_failure_leaves_no_recovery_file(tmp_path, monkeypatch): def fail_replace(source, destination): raise OSError("simulated replace failure") - monkeypatch.setattr("gateway.shutdown_flush.os.replace", fail_replace) + monkeypatch.setattr("utils.os.replace", fail_replace) assert flush_pending_to_file({"session": "message"}, reason="test") == 0 assert list(flush_dir.iterdir()) == []