Merge pull request #5088 from CalamitousFelicitousness/refactor/json-helpers-path-lock

Refactor/json helpers path lock

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:

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}')

requirements.txt

@@ -14,7 +14,6 @@ psutil
 pyyaml
 toml
 voluptuous
-fasteners
 limits
 orjson
 ftfy

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')