235 lines
8.7 KiB
Python
235 lines
8.7 KiB
Python
"""
|
|
Core cache management for persistent contexts.
|
|
"""
|
|
import os
|
|
import time
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
from .persistent_context import PersistentContext
|
|
from .storage import StorageManager
|
|
from .cleanup import CleanupManager
|
|
from ..config import get_config
|
|
|
|
|
|
class PersistentContextCache:
|
|
def __init__(self):
|
|
self._data: dict[str, PersistentContext] = {}
|
|
self._last_save_time = time.time()
|
|
|
|
# Thread pool for async disk operations
|
|
self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="persistent_context_saver")
|
|
self._save_lock = threading.Lock()
|
|
self._pending_save = False
|
|
self._needs_another_save = False
|
|
self._save_future = None
|
|
|
|
# Create cache directory if it doesn't exist
|
|
os.makedirs(self._get_cache_dir(), exist_ok=True)
|
|
|
|
# Load existing cache from disk
|
|
self._load()
|
|
|
|
def _get_cache_dir(self) -> str:
|
|
"""Get cache directory from config (real-time)"""
|
|
return get_config().get_cache_directory_path()
|
|
|
|
def _get_auto_save(self) -> bool:
|
|
"""Get auto_save setting from config (real-time)"""
|
|
config = get_config().get_persistent_context_config()
|
|
return config.get('auto_save', False)
|
|
|
|
def _get_max_cache_size_bytes(self) -> int:
|
|
"""Get max cache size from config (real-time)"""
|
|
config = get_config().get_persistent_context_config()
|
|
return config.get('max_cache_size_mb', 100) * 1024 * 1024
|
|
|
|
def _get_old_data_threshold_seconds(self) -> float:
|
|
"""Get old data threshold from config (real-time)"""
|
|
config = get_config().get_persistent_context_config()
|
|
return config.get('old_data_threshold_hours', 24) * 3600
|
|
|
|
def _get_absolute_max_cache_size_bytes(self) -> int:
|
|
"""Get absolute max cache size from config (real-time)"""
|
|
config = get_config().get_persistent_context_config()
|
|
return config.get('absolute_max_cache_size_mb', 200) * 1024 * 1024
|
|
|
|
def _get_max_context_size_bytes(self) -> int:
|
|
"""Get max context size from config (real-time)"""
|
|
config = get_config().get_persistent_context_config()
|
|
return config.get('max_context_size_mb', 50) * 1024 * 1024
|
|
|
|
def _get_max_key_length(self) -> int:
|
|
"""Get max key length from config (real-time)"""
|
|
config = get_config().get_persistent_context_config()
|
|
return config.get('max_key_length', 256)
|
|
|
|
def _validate_key(self, key: str):
|
|
"""Validate key length and format
|
|
|
|
Raises:
|
|
ValueError: If key is invalid
|
|
"""
|
|
if not key:
|
|
raise ValueError("Context key cannot be empty")
|
|
|
|
max_key_length = self._get_max_key_length()
|
|
if max_key_length > 0 and len(key) > max_key_length:
|
|
raise ValueError(f"Context key length ({len(key)}) exceeds maximum allowed length ({max_key_length})")
|
|
|
|
def get_context(self, key: str) -> PersistentContext:
|
|
self._validate_key(key)
|
|
self.update_context_access_time(key)
|
|
if key in self._data:
|
|
return self._data[key]
|
|
raise KeyError(f"Context with key '{key}' not found")
|
|
|
|
def has_context(self, key: str) -> bool:
|
|
try:
|
|
self._validate_key(key)
|
|
except ValueError:
|
|
return False
|
|
self.update_context_access_time(key)
|
|
return key in self._data
|
|
|
|
def remove_context(self, key: str):
|
|
self._validate_key(key)
|
|
self.update_context_access_time(key)
|
|
if key in self._data:
|
|
del self._data[key]
|
|
# Delete the file from disk
|
|
StorageManager.delete_context_file(self._get_cache_dir(), key)
|
|
|
|
def create_context(self, key: str, value: any) -> PersistentContext:
|
|
self._validate_key(key)
|
|
context = PersistentContext(value, key=key, cache=self)
|
|
self._data[key] = context
|
|
# Trigger an async save after creating a new context (if auto_save is enabled)
|
|
if self._get_auto_save():
|
|
self.save()
|
|
return context
|
|
|
|
def resolve_contexts_by_value_type(self, value_type: type) -> list[PersistentContext]:
|
|
"""Get all contexts whose value is an instance of the specified type
|
|
|
|
Args:
|
|
value_type: The type to filter by (e.g., str, int, dict, list, or custom classes)
|
|
|
|
Returns:
|
|
List of PersistentContext objects whose values match the specified type
|
|
"""
|
|
result = []
|
|
for context in self._data.values():
|
|
if isinstance(context.get_value(), value_type):
|
|
result.append(context)
|
|
return result
|
|
|
|
def update_context_access_time(self, key: str):
|
|
try:
|
|
self._validate_key(key)
|
|
except ValueError:
|
|
return
|
|
if key in self._data:
|
|
self._data[key].update_access_time()
|
|
|
|
def _load(self):
|
|
"""Load cache from disk - scan directory for individual context files"""
|
|
cache_dir = self._get_cache_dir()
|
|
self._data = StorageManager.load_contexts(cache_dir, self)
|
|
|
|
def _do_save(self):
|
|
"""Internal method to perform the actual save operation - save only dirty contexts"""
|
|
try:
|
|
saved_count = 0
|
|
removed_count = 0
|
|
skipped_count = 0
|
|
oversized_count = 0
|
|
|
|
cache_dir = self._get_cache_dir()
|
|
max_context_size = self._get_max_context_size_bytes()
|
|
|
|
# Only save dirty contexts
|
|
for key, context in list(self._data.items()):
|
|
# Skip if not dirty
|
|
if not context.is_dirty():
|
|
skipped_count += 1
|
|
continue
|
|
|
|
success, status = StorageManager.save_context(cache_dir, key, context, max_context_size)
|
|
|
|
if status == 'saved':
|
|
saved_count += 1
|
|
elif status == 'removed_empty':
|
|
removed_count += 1
|
|
# Remove from memory
|
|
del self._data[key]
|
|
elif status == 'oversized':
|
|
oversized_count += 1
|
|
|
|
self._last_save_time = time.time()
|
|
|
|
if saved_count > 0:
|
|
print(f"[comfyui-easytoolkit] Saved {saved_count} context(s) to {cache_dir}")
|
|
if removed_count > 0:
|
|
print(f"[comfyui-easytoolkit] Removed {removed_count} context(s) with empty values")
|
|
if oversized_count > 0:
|
|
print(f"[comfyui-easytoolkit] Skipped {oversized_count} oversized context(s)")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"[comfyui-easytoolkit] Error during save operation: {e}")
|
|
return False
|
|
|
|
def _cleanup_old_data(self):
|
|
"""Clean up old data when cache size exceeds threshold"""
|
|
CleanupManager.cleanup_old_data(
|
|
cache_dir=self._get_cache_dir(),
|
|
contexts_dict=self._data,
|
|
max_cache_size=self._get_max_cache_size_bytes(),
|
|
absolute_max_cache_size=self._get_absolute_max_cache_size_bytes(),
|
|
old_data_threshold=self._get_old_data_threshold_seconds(),
|
|
remove_context_callback=self.remove_context
|
|
)
|
|
|
|
def _save_task(self):
|
|
"""Task to run in thread pool for async saves"""
|
|
with self._save_lock:
|
|
self._pending_save = True
|
|
self._needs_another_save = False
|
|
|
|
success = self._do_save()
|
|
|
|
# After successful save, check if cleanup is needed
|
|
if success:
|
|
self._cleanup_old_data()
|
|
|
|
with self._save_lock:
|
|
self._pending_save = False
|
|
# If another save was requested while we were saving, do it now
|
|
if self._needs_another_save:
|
|
self._save_future = self._executor.submit(self._save_task)
|
|
else:
|
|
self._save_future = None
|
|
|
|
return success
|
|
|
|
def save(self, force_sync=False):
|
|
"""Save cache to disk asynchronously (or synchronously if force_sync=True)"""
|
|
if force_sync:
|
|
# Synchronous save
|
|
success = self._do_save()
|
|
if success:
|
|
self._cleanup_old_data()
|
|
return success
|
|
|
|
# Asynchronous save
|
|
with self._save_lock:
|
|
if self._pending_save:
|
|
# A save is already in progress, mark that we need another one
|
|
self._needs_another_save = True
|
|
else:
|
|
# Submit a new save task
|
|
self._save_future = self._executor.submit(self._save_task)
|
|
|