311 lines
10 KiB
Python
311 lines
10 KiB
Python
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Dict
|
|
|
|
|
|
_LOGGER = logging.getLogger("ComfyUI.SimpleIterator.Core")
|
|
|
|
|
|
def _sha1(parts) -> str:
|
|
"""对输入片段做稳定哈希，生成可持久化的短键。"""
|
|
digest = hashlib.sha1()
|
|
for part in parts:
|
|
digest.update(str(part).encode("utf-8", errors="ignore"))
|
|
digest.update(b"\x00")
|
|
return digest.hexdigest()
|
|
|
|
|
|
def stable_scope(kind: str, unique_id: str, source_key: str) -> str:
|
|
"""基于节点与数据源信息生成状态作用域键。"""
|
|
return _sha1((kind, unique_id or "default", source_key))
|
|
|
|
|
|
def format_output_filename(path: Path, with_ext: bool) -> str:
|
|
"""根据开关返回文件名（带后缀或不带后缀）。"""
|
|
return path.name if with_ext else path.stem
|
|
|
|
|
|
class IteratorStateStore:
|
|
"""定义迭代游标的持久化存取契约。
|
|
|
|
典型调用场景:
|
|
- 节点执行前查询当前游标位置（`peek_cursor`）。
|
|
- 节点执行时原子领取当前索引并推进游标（`claim_index`）。
|
|
"""
|
|
|
|
_lock = threading.Lock()
|
|
_loaded = False
|
|
_state: Dict[str, Dict[str, int]] = {}
|
|
_state_file = Path(__file__).resolve().parent / ".iterator_state.json"
|
|
_config_file = Path(__file__).resolve().parent / "iterator_config.json"
|
|
_env_ttl_key = "SIMPLE_ITERATOR_STATE_TTL_SECONDS"
|
|
_env_max_key = "SIMPLE_ITERATOR_STATE_MAX_ENTRIES"
|
|
_default_ttl_seconds = 30 * 24 * 60 * 60
|
|
_default_max_entries = 2000
|
|
_ttl_seconds = _default_ttl_seconds
|
|
_max_entries = _default_max_entries
|
|
|
|
@classmethod
|
|
def _parse_int(
|
|
cls,
|
|
*,
|
|
value,
|
|
default: int,
|
|
min_value: int,
|
|
source: str,
|
|
key: str,
|
|
) -> int:
|
|
"""将配置值解析为整数并做范围保护。"""
|
|
try:
|
|
parsed = int(value)
|
|
except (TypeError, ValueError):
|
|
_LOGGER.warning(
|
|
"Invalid %s '%s' from %s; fallback to default=%s",
|
|
key,
|
|
value,
|
|
source,
|
|
default,
|
|
)
|
|
return default
|
|
if parsed < min_value:
|
|
_LOGGER.warning(
|
|
"Out-of-range %s=%s from %s; require >= %s, fallback to default=%s",
|
|
key,
|
|
parsed,
|
|
source,
|
|
min_value,
|
|
default,
|
|
)
|
|
return default
|
|
return parsed
|
|
|
|
@classmethod
|
|
def _load_gc_settings(cls) -> None:
|
|
"""加载状态清理配置（默认值 < JSON 文件 < 环境变量）。"""
|
|
ttl_value = cls._default_ttl_seconds
|
|
max_value = cls._default_max_entries
|
|
|
|
file_settings = {}
|
|
if cls._config_file.exists():
|
|
try:
|
|
raw = json.loads(cls._config_file.read_text(encoding="utf-8"))
|
|
if isinstance(raw, dict):
|
|
file_settings = raw
|
|
else:
|
|
_LOGGER.warning(
|
|
"Invalid config root in %s; expected object, got %s. Using defaults.",
|
|
cls._config_file,
|
|
type(raw).__name__,
|
|
)
|
|
except (json.JSONDecodeError, OSError) as err:
|
|
_LOGGER.warning(
|
|
"Failed to read config file %s: %s. Using defaults.",
|
|
cls._config_file,
|
|
err,
|
|
)
|
|
|
|
if "state_ttl_seconds" in file_settings:
|
|
ttl_value = cls._parse_int(
|
|
value=file_settings["state_ttl_seconds"],
|
|
default=cls._default_ttl_seconds,
|
|
min_value=0,
|
|
source=f"file:{cls._config_file}",
|
|
key="state_ttl_seconds",
|
|
)
|
|
if "state_max_entries" in file_settings:
|
|
max_value = cls._parse_int(
|
|
value=file_settings["state_max_entries"],
|
|
default=cls._default_max_entries,
|
|
min_value=1,
|
|
source=f"file:{cls._config_file}",
|
|
key="state_max_entries",
|
|
)
|
|
|
|
env_ttl = os.getenv(cls._env_ttl_key)
|
|
if env_ttl is not None:
|
|
ttl_value = cls._parse_int(
|
|
value=env_ttl,
|
|
default=ttl_value,
|
|
min_value=0,
|
|
source=f"env:{cls._env_ttl_key}",
|
|
key="state_ttl_seconds",
|
|
)
|
|
env_max = os.getenv(cls._env_max_key)
|
|
if env_max is not None:
|
|
max_value = cls._parse_int(
|
|
value=env_max,
|
|
default=max_value,
|
|
min_value=1,
|
|
source=f"env:{cls._env_max_key}",
|
|
key="state_max_entries",
|
|
)
|
|
|
|
cls._ttl_seconds = ttl_value
|
|
cls._max_entries = max_value
|
|
|
|
@classmethod
|
|
def _ensure_loaded(cls) -> None:
|
|
if cls._loaded:
|
|
return
|
|
cls._load_gc_settings()
|
|
if cls._state_file.exists():
|
|
try:
|
|
cls._state = json.loads(cls._state_file.read_text(encoding="utf-8"))
|
|
except (json.JSONDecodeError, OSError):
|
|
cls._state = {}
|
|
else:
|
|
cls._state = {}
|
|
pruned = cls._prune_state()
|
|
cls._loaded = True
|
|
if pruned:
|
|
cls._save()
|
|
|
|
@classmethod
|
|
def _save(cls) -> None:
|
|
cls._prune_state()
|
|
payload = json.dumps(cls._state, ensure_ascii=False, indent=2)
|
|
cls._state_file.write_text(payload, encoding="utf-8")
|
|
|
|
@classmethod
|
|
def _prune_state(cls, now_ts: int | None = None) -> bool:
|
|
"""按保留策略清理状态条目，避免状态文件无限增长。
|
|
|
|
语义:
|
|
- 先删除超过 TTL 未更新的条目，再执行最大条目数裁剪。
|
|
- 无效条目（非字典）会被直接移除。
|
|
|
|
返回:
|
|
- `True`: 状态有变更（发生清理）。
|
|
- `False`: 状态无变更。
|
|
"""
|
|
if not isinstance(cls._state, dict):
|
|
cls._state = {}
|
|
return True
|
|
|
|
changed = False
|
|
now_value = int(time.time()) if now_ts is None else int(now_ts)
|
|
ttl = max(int(cls._ttl_seconds), 0)
|
|
max_entries = max(int(cls._max_entries), 1)
|
|
|
|
# 语义: TTL 清理优先于容量裁剪，先清掉长期无活跃条目。
|
|
if ttl > 0:
|
|
expired_keys = []
|
|
for key, entry in cls._state.items():
|
|
if not isinstance(entry, dict):
|
|
expired_keys.append(key)
|
|
continue
|
|
updated_at = int(entry.get("updated_at", 0) or 0)
|
|
if (now_value - updated_at) > ttl:
|
|
expired_keys.append(key)
|
|
if expired_keys:
|
|
for key in expired_keys:
|
|
cls._state.pop(key, None)
|
|
changed = True
|
|
else:
|
|
invalid_keys = [key for key, entry in cls._state.items() if not isinstance(entry, dict)]
|
|
if invalid_keys:
|
|
for key in invalid_keys:
|
|
cls._state.pop(key, None)
|
|
changed = True
|
|
|
|
# 说明: 当条目超过上限时，按 updated_at 新到旧保留最近 N 条。
|
|
if len(cls._state) > max_entries:
|
|
ranked = sorted(
|
|
cls._state.items(),
|
|
key=lambda item: int(item[1].get("updated_at", 0) or 0),
|
|
reverse=True,
|
|
)
|
|
cls._state = dict(ranked[:max_entries])
|
|
changed = True
|
|
|
|
return changed
|
|
|
|
@classmethod
|
|
def peek_cursor(cls, scope_key: str) -> int:
|
|
"""读取当前作用域的游标位置（不推进状态）。
|
|
|
|
参数:
|
|
- `scope_key`: 由节点类型、节点唯一标识、数据源键组合出的作用域键。
|
|
|
|
返回:
|
|
- 当前游标值，不存在时返回 `0`。
|
|
"""
|
|
with cls._lock:
|
|
cls._ensure_loaded()
|
|
entry = cls._state.get(scope_key, {})
|
|
return int(entry.get("cursor", 0))
|
|
|
|
@classmethod
|
|
def claim_index(
|
|
cls,
|
|
*,
|
|
scope_key: str,
|
|
total: int,
|
|
reset: bool,
|
|
loop_mode: str,
|
|
) -> int:
|
|
"""原子领取当前索引并按策略推进游标。
|
|
|
|
参数:
|
|
- `scope_key`: 迭代状态作用域键。
|
|
- `total`: 当前可迭代元素总数，必须大于 `0`。
|
|
- `reset`: 沿触发复位信号；仅在 `False -> True` 边沿时从 `0` 重新开始。
|
|
- `loop_mode`: 迭代耗尽行为，支持 `loop/stop/hold_last`。
|
|
|
|
返回:
|
|
- 本次应输出的索引（推进前索引）。
|
|
"""
|
|
if total <= 0:
|
|
raise ValueError("total must be greater than zero")
|
|
|
|
with cls._lock:
|
|
cls._ensure_loaded()
|
|
entry = cls._state.get(scope_key, {})
|
|
prev_reset = bool(entry.get("last_reset", False))
|
|
reset_edge = bool(reset) and (not prev_reset)
|
|
cursor = 0 if reset_edge else int(entry.get("cursor", 0))
|
|
|
|
if cursor < 0:
|
|
cursor = 0
|
|
|
|
# 语义: 当游标越界时，先根据 loop_mode 决定“当前输出索引”如何收敛。
|
|
# 优先级: reset 结果 > loop_mode 越界策略 > 默认顺序推进。
|
|
if cursor >= total:
|
|
if loop_mode == "loop":
|
|
cursor = 0
|
|
elif loop_mode == "hold_last":
|
|
cursor = total - 1
|
|
elif loop_mode == "stop":
|
|
raise RuntimeError(
|
|
"Iterator exhausted. Set loop_mode=loop/hold_last or enable reset."
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported loop_mode: {loop_mode}")
|
|
|
|
emit_index = cursor
|
|
|
|
# 语义: 输出后计算“下一次执行”的游标。
|
|
# 说明: stop 模式会把游标推进到 total 之外，用于下次触发耗尽错误。
|
|
if loop_mode == "loop":
|
|
next_cursor = (cursor + 1) % total
|
|
elif loop_mode == "hold_last":
|
|
next_cursor = min(cursor + 1, total - 1)
|
|
elif loop_mode == "stop":
|
|
next_cursor = cursor + 1
|
|
else:
|
|
raise ValueError(f"Unsupported loop_mode: {loop_mode}")
|
|
|
|
cls._state[scope_key] = {
|
|
"cursor": next_cursor,
|
|
"total": total,
|
|
"last_reset": bool(reset),
|
|
"updated_at": int(time.time()),
|
|
}
|
|
cls._save()
|
|
return emit_index
|