diff --git a/modules/hashes.py b/modules/hashes.py index 0c6c1021e..d4658c7ae 100644 --- a/modules/hashes.py +++ b/modules/hashes.py @@ -1,5 +1,6 @@ import hashlib import os.path +import threading from collections import defaultdict from typing import Literal, TypeAlias, TypedDict from rich import progress, errors @@ -29,6 +30,7 @@ cache_filename = os.path.join(data_path, "data", "cache.json") progress_ok = True # defaultdict allows for easily using new stores without needing to define them ahead of time _data: defaultdict[str, HashStore] = defaultdict(HashStore) +cache_lock = threading.Lock() def load_cache(): @@ -38,9 +40,9 @@ def load_cache(): def save_cache(): - # Don't include empty hash stores - filtered = filter(lambda item: len(item[1]) > 0, _data.items()) - writefile(dict(filtered), cache_filename) + with cache_lock: # snapshot and write together, so a later save never lands under an earlier snapshot + snapshot = {store: dict(data) for store, data in list(_data.items()) if len(data) > 0} # dict() of a store runs under the GIL, so add_hash from another thread cannot interrupt it + writefile(snapshot, cache_filename) def cache(store: KnownHashStores | str | None = None) -> HashStore: diff --git a/modules/json_helpers.py b/modules/json_helpers.py index 5d9a6d707..b780ae45f 100644 --- a/modules/json_helpers.py +++ b/modules/json_helpers.py @@ -1,14 +1,40 @@ import os import sys import contextlib +import threading import time import json from typing import overload, Literal -import fasteners import orjson from modules.logger import log -locking_available = True # used by file read/write locking +path_locks: dict[str, threading.RLock] = {} +path_locks_guard = threading.Lock() + + +def path_lock(filename: str | os.PathLike[str]) -> threading.RLock: + """One lock per file path; threads of this process serialize on it, other processes are covered by atomic replace.""" + key = os.path.normcase(os.path.realpath(filename)) + with path_locks_guard: + lock = path_locks.get(key) + if lock is None: + lock = path_locks[key] = threading.RLock() + with contextlib.suppress(OSError): + os.remove(f"{key}.lock") # left behind by the file lock this replaces + return lock + + +def read_bytes(filename: str | os.PathLike[str], attempts: int = 5, delay: float = 0.01) -> bytes: + """Read a whole file, waiting out an open that Windows refuses while a replace of the same name is in flight.""" + for attempt in range(attempts): + try: + with open(filename, "rb") as file: + return file.read() + except PermissionError: + if attempt == attempts - 1: + raise + time.sleep(delay) + return b"" def replace_file(source: str, target: str, attempts: int = 10, delay: float = 0.05): @@ -30,51 +56,28 @@ def readfile(filename: str | os.PathLike[str], silent: bool = False, lock: bool @overload def readfile(filename: str | os.PathLike[str], silent: bool = False, lock: bool = False) -> dict | list: ... def readfile(filename: str | os.PathLike[str], silent: bool = False, lock: bool = False, *, as_type="") -> dict | list: - global locking_available # pylint: disable=global-statement + """Read a JSON file; lock=True serializes with writers of the same path in this process.""" data = {} if as_type == "dict" else [] - lock_file = None - locked = False - if lock and locking_available: + with path_lock(filename) if lock else contextlib.nullcontext(): try: - lock_file = fasteners.InterProcessReaderWriterLock(f"{filename}.lock") - lock_file.logger.disabled = True # type: ignore - False positive. Bad typing in Fasteners. - locked = lock_file.acquire_read_lock(blocking=True, timeout=3) - if not locked: - log.warning(f'File read lock: file="{filename}" timeout') - except Exception as err: - lock_file = None - locking_available = False - log.error(f'File read lock: file="{filename}" {err}') - locked = False - - try: - # if not os.path.exists(filename): - # return {} - t0 = time.time() - with open(filename, "rb") as file: - b = file.read() + t0 = time.time() + b = read_bytes(filename) if len(b) == 0: + if not silent: + log.warning(f'Read: file="{filename}" empty') return {} if as_type == "dict" else [] data = orjson.loads(b) # pylint: disable=no-member - # if type(data) is str: - # data = json.loads(data) - t1 = time.time() - if not silent: - fn = f"{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}" # pylint: disable=protected-access - log.debug(f'Read: file="{filename}" json={len(data)} bytes={os.path.getsize(filename)} time={t1 - t0:.3f} fn={fn}') - except FileNotFoundError as err: - if not silent: - log.debug(f'Read failed: file="{filename}" {err}') - except Exception as err: - if not silent: - log.error(f'Read failed: file="{filename}" {err}') - - try: - if locking_available and lock_file is not None: - lock_file.release_read_lock() # the lock file stays: removing it after release races other holders - except Exception as err: - log.error(f'File read lock release: file="{filename}" {err}') + t1 = time.time() + if not silent: + fn = f"{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}" # pylint: disable=protected-access + log.debug(f'Read: file="{filename}" json={len(data)} bytes={os.path.getsize(filename)} time={t1 - t0:.3f} fn={fn}') + except FileNotFoundError as err: + if not silent: + log.debug(f'Read failed: file="{filename}" {err}') + except Exception as err: + if not silent: + log.error(f'Read failed: file="{filename}" {err}') if isinstance(data, list) and as_type == "dict": if not data: @@ -92,72 +95,55 @@ def readfile(filename: str | os.PathLike[str], silent: bool = False, lock: bool return data -def writefile(obj: dict | list, filename: str | os.PathLike[str], mode="w", silent=False, atomic=False): +def writefile(obj: dict | list, filename: str | os.PathLike[str], mode="w", silent=False, atomic=True): + """Write obj as JSON through a temp file and replace; writes to the same path from this process run one at a time, in call order.""" import copy import tempfile - global locking_available # pylint: disable=global-statement - lock_file = None - locked = False - def default(obj): log.error(f'Save: file="{filename}" not a valid object: {obj}') return str(obj) - try: - t0 = time.time() - snapshot = obj.copy() if isinstance(obj, (dict, list)) else obj # dict.copy and list.copy run under the GIL, so a concurrent insert cannot interrupt them - data = copy.deepcopy(snapshot) - for k, v in list(data.items()) if isinstance(data, dict) else []: # validate each key-by-key to avoid global exceptions - try: - _tmp = json.dumps(v, indent=2, default=default, allow_nan=False, ensure_ascii=False) - except Exception as err: - if not silent: - log.error(f'Save: file="{filename}" key="{k}" value="{v}" {err}') - del data[k] - output = json.dumps(data, indent=2, default=default) - except Exception as err: - log.error(f'Save failed: file="{filename}" {err}') - return + if mode != "w": + atomic = False # append cannot go through a temp file - try: - if locking_available: - lock_file = fasteners.InterProcessReaderWriterLock(f"{filename}.lock") if locking_available else None - lock_file.logger.disabled = True # type: ignore - False positive. Bad typing in Fasteners. - locked = lock_file.acquire_write_lock(blocking=True, timeout=3) if lock_file is not None else False - if not locked: - log.warning(f'File write lock: file="{filename}" timeout') - except Exception as err: - locking_available = False - lock_file = None - log.error(f'File write lock: file="{filename}" {err}') - locked = False + with path_lock(filename): + try: + t0 = time.time() + snapshot = obj.copy() if isinstance(obj, (dict, list)) else obj # dict.copy and list.copy run under the GIL, so a concurrent insert cannot interrupt them + data = copy.deepcopy(snapshot) + for k, v in list(data.items()) if isinstance(data, dict) else []: # validate each key-by-key to avoid global exceptions + try: + _tmp = json.dumps(v, indent=2, default=default, allow_nan=False, ensure_ascii=False) + except Exception as err: + if not silent: + log.error(f'Save: file="{filename}" key="{k}" value="{v}" {err}') + del data[k] + output = json.dumps(data, indent=2, default=default) + except Exception as err: + log.error(f'Save failed: file="{filename}" {err}') + return - try: - if atomic: - fd, temp_name = tempfile.mkstemp(dir=os.path.dirname(os.path.abspath(filename)), prefix=f"{os.path.basename(filename)}.", suffix=".tmp") - try: - with os.fdopen(fd, mode, encoding="utf8") as f: - f.write(output) - f.flush() - os.fsync(f.fileno()) - replace_file(temp_name, filename) - except BaseException: - with contextlib.suppress(OSError): - os.remove(temp_name) - raise - else: - with open(filename, mode=mode, encoding="utf8") as file: - file.write(output) - t1 = time.time() - if not silent: - datalength = len(data) - log.debug(f'Save: file="{filename}" json={datalength} bytes={len(output)} time={t1 - t0:.3f}') - except Exception as err: - log.error(f'Save failed: file="{filename}" {err}') - - try: - if locking_available and lock_file is not None: - lock_file.release_write_lock() # the lock file stays: removing it after release races other holders - except Exception as err: - log.error(f'File write lock release: file="{filename}" {err}') + try: + if atomic: + target = os.path.realpath(filename) # replace the file a symlink points at, not the symlink + fd, temp_name = tempfile.mkstemp(dir=os.path.dirname(target), prefix=f"{os.path.basename(target)}.", suffix=".tmp") + try: + with os.fdopen(fd, mode, encoding="utf8") as f: + f.write(output) + f.flush() + os.fsync(f.fileno()) + replace_file(temp_name, target) + except BaseException: + with contextlib.suppress(OSError): + os.remove(temp_name) + raise + else: + with open(filename, mode=mode, encoding="utf8") as file: + file.write(output) + t1 = time.time() + if not silent: + datalength = len(data) + log.debug(f'Save: file="{filename}" json={datalength} bytes={len(output)} time={t1 - t0:.3f}') + except Exception as err: + log.error(f'Save failed: file="{filename}" {err}') diff --git a/requirements.txt b/requirements.txt index 7d5167b82..080d9c2ce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,7 +14,6 @@ psutil pyyaml toml voluptuous -fasteners limits orjson ftfy diff --git a/test/test-json-helpers.py b/test/test-json-helpers.py index 95e91ce31..22a677970 100644 --- a/test/test-json-helpers.py +++ b/test/test-json-helpers.py @@ -6,9 +6,16 @@ Covers: - atomic saves replace a fresh or existing target and leave no temp file, including when the save fails - an atomic save waits out a target another thread briefly holds open, which Windows otherwise refuses -- locking stays available under overlapping locked reads and writes and the lock file is left in place +- overlapping locked reads and atomic writes never fail or read a torn file, and no lock file is created +- locked readers never see a torn file while unlocked writers rewrite it in place +- a lock file left behind by the former file lock is removed - concurrent inserts into a shared dict never drop a save +- default writes are atomic, so unlocked readers never see a torn file either +- an atomic write through a symlink replaces the file it points at and keeps the link +- a read waits out an open that is refused while a replace is in flight +- an empty file reads as empty and is reported unless the read is silent - readfile returns what writefile wrote, as dict and as list +- the hash cache saves cleanly while other threads keep adding hashes No running server required. @@ -23,6 +30,7 @@ import shutil import sys import tempfile import threading +import time script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, script_dir) @@ -34,17 +42,22 @@ from modules import json_helpers # pylint: disable=wrong-import-position class ErrorLog: - """Counts error lines from json_helpers while passing everything else to the real logger.""" + """Counts error and warning lines from json_helpers while passing everything else to the real logger.""" def __init__(self, inner): self.inner = inner self.errors = [] + self.warnings = [] self.lock = threading.Lock() def error(self, message, *args, **kwargs): with self.lock: self.errors.append(str(message)) + def warning(self, message, *args, **kwargs): + with self.lock: + self.warnings.append(str(message)) + def count(self, needle): return sum(1 for m in self.errors if needle in m) @@ -116,25 +129,65 @@ def test_atomic_save_waits_for_open_target(folder): assert leftovers(folder, target) == [], f'temp files left: {leftovers(folder, target)}' -def test_locking_survives_overlap(folder, threads=8, iterations=25): +def test_overlapping_locked_access(folder, threads=8, iterations=25): jh = fresh_helper() target = os.path.join(folder, 'locked.json') jh.writefile({'seed': 0}, target, silent=True) + torn = [] def worker(n): for i in range(iterations): if (n + i) % 2: jh.writefile({'n': n, 'i': i}, target, silent=True, atomic=True) - else: - jh.readfile(target, silent=True, lock=True, as_type='dict') + elif not jh.readfile(target, silent=True, lock=True, as_type='dict'): + torn.append((n, i)) run_threads(worker, threads) - assert jh.locking_available, 'locking switched itself off' assert jh.log.errors == [], jh.log.errors - assert os.path.exists(target + '.lock'), 'lock file was removed' + assert torn == [], f'{len(torn)} reads returned nothing' + assert not os.path.exists(target + '.lock'), 'a lock file was created' assert leftovers(folder, target) == [], f'temp files left: {leftovers(folder, target)}' +def test_locked_readers_never_see_torn_writes(folder, writers=8, per_writer=15, readers=2): + jh = fresh_helper() + target = os.path.join(folder, 'config.json') + jh.writefile({'seed': 0}, target, silent=True) + stop = threading.Event() + torn = [] + + def reader(_n): + while not stop.is_set(): + if not jh.readfile(target, silent=True, lock=True, as_type='dict'): + torn.append(1) + + def writer(n): + for i in range(per_writer): + jh.writefile({'writer': n, 'i': i, 'blob': 'x' * (40000 + 1000 * n)}, target, silent=True, atomic=False) + + pool = [threading.Thread(target=reader, args=(r,)) for r in range(readers)] + for t in pool: + t.start() + run_threads(writer, writers) + stop.set() + for t in pool: + t.join() + assert torn == [], f'{len(torn)} reads returned nothing' + with open(target, encoding='utf8') as f: + json.load(f) + assert jh.log.errors == [], jh.log.errors + + +def test_legacy_lock_file_removed(folder): + jh = fresh_helper() + target = os.path.join(folder, 'legacy.json') + with open(target + '.lock', 'w', encoding='utf8'): + pass + jh.writefile({'x': 1}, target, silent=True) + assert not os.path.exists(target + '.lock'), 'legacy lock file kept' + assert jh.log.errors == [], jh.log.errors + + def test_concurrent_inserts_keep_every_save(folder, threads=4, per_thread=40): jh = fresh_helper() target = os.path.join(folder, 'cache.json') @@ -153,6 +206,126 @@ def test_concurrent_inserts_keep_every_save(folder, threads=4, per_thread=40): assert len(json.load(f)) == threads * per_thread, 'last save is incomplete' +def test_default_writes_never_tear(folder, writers=4, per_writer=20, readers=2): + jh = fresh_helper() + target = os.path.join(folder, 'default.json') + jh.writefile({'seed': 0}, target, silent=True) + stop = threading.Event() + torn = [] + + def reader(_n): + while not stop.is_set(): + if not jh.readfile(target, silent=True, as_type='dict'): + torn.append(1) + time.sleep(0.001) # real readers do not spin; a spinning reader starves the writer's retry on Windows + + def writer(n): + for i in range(per_writer): + jh.writefile({'writer': n, 'i': i, 'blob': 'x' * (40000 + 1000 * n)}, target, silent=True) + + pool = [threading.Thread(target=reader, args=(r,)) for r in range(readers)] + for t in pool: + t.start() + run_threads(writer, writers) + stop.set() + for t in pool: + t.join() + assert torn == [], f'{len(torn)} unlocked reads returned nothing' + assert jh.log.errors == [], jh.log.errors + assert leftovers(folder, target) == [], f'temp files left: {leftovers(folder, target)}' + + +def test_atomic_write_keeps_symlink(folder): + jh = fresh_helper() + real = os.path.join(folder, 'real.json') + link = os.path.join(folder, 'link.json') + jh.writefile({'v': 0}, real, silent=True) + try: + os.symlink(real, link) + except OSError as e: + log.info(f' SKIP symlink not available: {e}') + return + jh.writefile({'v': 1}, link, silent=True) + assert os.path.islink(link), 'symlink was replaced by a file' + assert jh.readfile(real, silent=True, as_type='dict') == {'v': 1}, 'target of the symlink not updated' + assert sorted(os.listdir(folder)) == ['link.json', 'real.json'], os.listdir(folder) + assert jh.log.errors == [], jh.log.errors + + +def test_read_waits_out_refused_open(folder): + jh = fresh_helper() + target = os.path.join(folder, 'refused.json') + jh.writefile({'x': 1}, target, silent=True) + calls = [] + real_open = open + + def refuse_twice(*args, **kwargs): + calls.append(1) + if len(calls) <= 2: + raise PermissionError(13, 'replace in flight') + return real_open(*args, **kwargs) + + jh.open = refuse_twice # module globals shadow the builtin inside read_bytes + try: + assert jh.readfile(target, silent=True, as_type='dict') == {'x': 1} + finally: + del jh.open + assert len(calls) == 3, f'{len(calls)} open attempts' + assert jh.log.errors == [], jh.log.errors + + +def test_empty_file_is_reported(folder): + jh = fresh_helper() + target = os.path.join(folder, 'empty.json') + with open(target, 'w', encoding='utf8'): + pass + assert jh.readfile(target, as_type='dict') == {} + assert jh.readfile(target, silent=True, as_type='list') == [] + assert len(jh.log.warnings) == 1 and 'empty' in jh.log.warnings[0], jh.log.warnings + assert jh.log.errors == [], jh.log.errors + + +def test_hash_cache_saves_under_concurrent_adds(folder, adders=4, per_adder=200): + jh = fresh_helper() + from modules import hashes # pylint: disable=import-outside-toplevel + saved_filename = hashes.cache_filename + hashes.cache_filename = os.path.join(folder, 'cache.json') + hashes.cache('hashes').clear() + hashes.cache('hashes-addnet').clear() + stop = threading.Event() + + def adder(n): + for i in range(per_adder): + hashes.cache('hashes').add_hash(f'checkpoint/{n}-{i}', 1.0, 'a' * 64) + if i % 10 == 9: + hashes.save_cache() + + def churn(_n): # a second store whose size keeps changing while the saves run, bounded so the snapshots stay small + i = 0 + while not stop.is_set(): + store = hashes.cache('hashes-addnet') + if i % 500 == 499: + store.clear() + else: + store.add_hash(f'lora/{i % 500}', 1.0, 'b' * 64) + i += 1 + + thread = threading.Thread(target=churn, args=(0,)) + thread.start() + try: + run_threads(adder, adders) + finally: + stop.set() + thread.join() + hashes.cache_filename = saved_filename + assert jh.log.errors == [], jh.log.errors + with open(os.path.join(folder, 'cache.json'), encoding='utf8') as f: + on_disk = json.load(f) + assert len(on_disk['hashes']) == adders * per_adder, f'{len(on_disk["hashes"])} of {adders * per_adder} hashes on disk' + hashes.cache('hashes').clear() + hashes.cache('hashes-addnet').clear() + + def test_roundtrip(folder): jh = fresh_helper() target = os.path.join(folder, 'roundtrip.json') @@ -170,8 +343,15 @@ def run_all(): test_atomic_save_replaces_target, test_failed_atomic_save_leaves_no_temp, test_atomic_save_waits_for_open_target, - test_locking_survives_overlap, + test_overlapping_locked_access, + test_locked_readers_never_see_torn_writes, + test_legacy_lock_file_removed, test_concurrent_inserts_keep_every_save, + test_default_writes_never_tear, + test_atomic_write_keeps_symlink, + test_read_waits_out_refused_open, + test_empty_file_is_reported, + test_hash_cache_saves_under_concurrent_adds, test_roundtrip, ] passed = 0 @@ -192,7 +372,6 @@ def run_all(): if __name__ == '__main__': - import time t0 = time.time() ok = run_all() log.warning(f'Total time: {time.time() - t0:.2f}s')