Spaces:
Sleeping
Sleeping
File size: 5,419 Bytes
2415446 a1bab2d 2415446 a1bab2d 2415446 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | """Atomic JSON persistence for messaging session state."""
from __future__ import annotations
import contextlib
import json
import os
import tempfile
import threading
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
from loguru import logger
@dataclass(frozen=True)
class _PendingWrite:
generation: int
snapshot: dict[str, Any]
class DebouncedJsonPersistence:
"""Thread-safe debounced JSON writer with atomic replace semantics."""
def __init__(
self,
storage_path: str,
*,
snapshot: Callable[[], dict[str, Any]],
on_dirty: Callable[[bool], None],
debounce_secs: float = 0.5,
) -> None:
self.storage_path = storage_path
self._snapshot = snapshot
self._on_dirty = on_dirty
self._debounce_secs = debounce_secs
self._save_timer: threading.Timer | None = None
self._timer_lock = threading.Lock()
self._writer_lock = threading.Lock()
self._save_generation = 0
def load_json(self) -> dict[str, Any]:
if not os.path.exists(self.storage_path):
return {}
with open(self.storage_path, encoding="utf-8") as file:
data = json.load(file)
return data if isinstance(data, dict) else {}
def schedule_save(self) -> None:
self._on_dirty(True)
with self._timer_lock:
if self._save_timer is not None:
self._save_timer.cancel()
self._save_generation += 1
generation = self._save_generation
timer = threading.Timer(
self._debounce_secs,
self._save_from_timer,
args=(generation,),
)
timer.daemon = True
self._save_timer = timer
timer.start()
def flush(self) -> None:
self._on_dirty(True)
pending = self._snapshot_for_write()
if pending is None:
return
self._write_pending(pending)
def _save_from_timer(self, generation: int) -> None:
try:
pending = self._snapshot_for_write(expected_generation=generation)
if pending is None:
return
self._write_pending(pending)
except Exception as e:
self._on_dirty(True)
logger.error(
"Failed to save sessions: exc_type={}",
type(e).__name__,
)
def _write_pending(self, pending: _PendingWrite) -> None:
try:
written = self._write_if_current(pending)
except Exception:
self._on_dirty(True)
raise
if written:
self._mark_clean_if_current(pending.generation)
def _write_if_current(self, pending: _PendingWrite) -> bool:
"""Serialize writers and reject a snapshot superseded before replace."""
with self._writer_lock:
with self._timer_lock:
if pending.generation != self._save_generation:
return False
self._write_file(pending.snapshot)
return True
def _snapshot_for_write(
self, *, expected_generation: int | None = None
) -> _PendingWrite | None:
generation = self._claim_timer(expected_generation)
if generation is None:
return None
snapshot = self._snapshot()
return _PendingWrite(generation=generation, snapshot=snapshot)
def _claim_timer(self, expected_generation: int | None) -> int | None:
with self._timer_lock:
if expected_generation is not None and (
expected_generation != self._save_generation or self._save_timer is None
):
return None
if self._save_timer is not None:
self._save_timer.cancel()
self._save_timer = None
return self._save_generation
def _mark_clean_if_current(self, generation: int) -> None:
with self._timer_lock:
is_current = (
self._save_timer is None and generation == self._save_generation
)
if is_current:
self._on_dirty(False)
def write_data(self, data: dict[str, Any]) -> None:
"""Write authoritative state after invalidating older pending snapshots."""
self._on_dirty(True)
with self._timer_lock:
if self._save_timer is not None:
self._save_timer.cancel()
self._save_timer = None
self._save_generation += 1
pending = _PendingWrite(
generation=self._save_generation,
snapshot=data,
)
self._write_pending(pending)
def _write_file(self, data: dict[str, Any]) -> None:
abs_target = os.path.abspath(self.storage_path)
dir_name = os.path.dirname(abs_target) or "."
fd, tmp_path = tempfile.mkstemp(
dir=dir_name,
prefix=".sessions.",
suffix=".tmp.json",
)
try:
with os.fdopen(fd, "w", encoding="utf-8") as file:
json.dump(data, file, indent=2)
file.flush()
os.fsync(file.fileno())
os.replace(tmp_path, abs_target)
except BaseException:
with contextlib.suppress(OSError):
os.unlink(tmp_path)
raise
|