27db6ef197
The per-bucket lock was inserted without synchronization. Two first writers could take different locks and overwrite each other's preset file. Lock creation now happens under a guard lock. Co-authored-by: AE <by-ae@users.noreply.github.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com>
139 lines
3.9 KiB
Python
139 lines
3.9 KiB
Python
"""File-based character/object/scene storage in ComfyUI user data."""
|
|
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
|
|
import folder_paths
|
|
|
|
_STORE_SUBDIR = "byae-ideogram-tools"
|
|
_BUCKETS = frozenset({"characters", "objects", "scenes"})
|
|
# Allow Unicode letters, numbers, spaces, underscores, and hyphens
|
|
# Must start with a letter or number (Unicode-aware)
|
|
_SAFE_NAME = re.compile(r"^[\w][\w\- ]{0,63}$", re.UNICODE)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# File locks for concurrent write protection (one lock per bucket).
|
|
# The guard lock covers creating the per-bucket lock; without it two threads
|
|
# can each install a different lock and then write the same file together.
|
|
_BUCKET_LOCKS = {}
|
|
_BUCKET_LOCKS_GUARD = threading.Lock()
|
|
|
|
|
|
def _store_dir() -> str:
|
|
path = os.path.join(folder_paths.get_user_directory(), _STORE_SUBDIR)
|
|
os.makedirs(path, exist_ok=True)
|
|
return path
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _bucket_lock(bucket: str):
|
|
"""Thread-safe lock for bucket writes to prevent concurrent write corruption."""
|
|
with _BUCKET_LOCKS_GUARD:
|
|
lock = _BUCKET_LOCKS.get(bucket)
|
|
if lock is None:
|
|
lock = threading.Lock()
|
|
_BUCKET_LOCKS[bucket] = lock
|
|
lock.acquire()
|
|
try:
|
|
yield
|
|
finally:
|
|
lock.release()
|
|
|
|
|
|
def _bucket_path(bucket: str) -> str:
|
|
if bucket not in _BUCKETS:
|
|
raise ValueError(f"Unknown store bucket: {bucket}")
|
|
return os.path.join(_store_dir(), f"{bucket}.json")
|
|
|
|
|
|
def _validate_name(name: str) -> str:
|
|
name = (name or "").strip()
|
|
if not _SAFE_NAME.match(name):
|
|
raise ValueError("Invalid name (use letters, numbers, spaces, _ or -)")
|
|
return name
|
|
|
|
|
|
def _read_json(path: str) -> dict:
|
|
if not os.path.isfile(path):
|
|
return {}
|
|
try:
|
|
with open(path, encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
return data if isinstance(data, dict) else {}
|
|
except (json.JSONDecodeError, UnicodeDecodeError) as e:
|
|
# Move corrupt file aside and treat bucket as empty
|
|
timestamp = int(time.time())
|
|
backup_path = f"{path}.corrupt-{timestamp}"
|
|
logger.warning(
|
|
"ByAE Ideogram: corrupt store file %s (%s), moving to %s",
|
|
path, e, backup_path
|
|
)
|
|
try:
|
|
os.rename(path, backup_path)
|
|
except OSError as rename_err:
|
|
logger.error("ByAE Ideogram: failed to backup corrupt file: %s", rename_err)
|
|
return {}
|
|
|
|
|
|
def _write_json(path: str, data: dict) -> None:
|
|
# Use a temporary file with a unique name to avoid conflicts during concurrent writes
|
|
temp_fd, tmp = tempfile.mkstemp(
|
|
dir=os.path.dirname(path),
|
|
prefix=".tmp_",
|
|
suffix=".json"
|
|
)
|
|
try:
|
|
with os.fdopen(temp_fd, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=2, ensure_ascii=False)
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
os.replace(tmp, path)
|
|
except Exception:
|
|
# Clean up temp file on write failure
|
|
try:
|
|
os.close(temp_fd)
|
|
except (OSError, ValueError):
|
|
pass
|
|
if os.path.exists(tmp):
|
|
try:
|
|
os.remove(tmp)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
|
|
|
|
def load_bucket(bucket: str) -> dict:
|
|
return _read_json(_bucket_path(bucket))
|
|
|
|
|
|
def load_all() -> dict:
|
|
return {bucket: load_bucket(bucket) for bucket in _BUCKETS}
|
|
|
|
|
|
def save_entry(bucket: str, name: str, data) -> None:
|
|
name = _validate_name(name)
|
|
path = _bucket_path(bucket)
|
|
with _bucket_lock(bucket):
|
|
store = _read_json(path)
|
|
store[name] = data
|
|
_write_json(path, store)
|
|
|
|
|
|
def delete_entry(bucket: str, name: str) -> bool:
|
|
name = _validate_name(name)
|
|
path = _bucket_path(bucket)
|
|
with _bucket_lock(bucket):
|
|
store = _read_json(path)
|
|
if name not in store:
|
|
return False
|
|
del store[name]
|
|
_write_json(path, store)
|
|
return True
|