feat(01-01): implement WorkflowStateManager with atomic write
GREEN phase: Implement state management with: - Atomic write pattern (temp file + os.replace) - Backup rotation (workflow_state.json -> .bak) - Backup fallback on corruption - get_or_create idempotent creation Fix test to properly create backup before corruption test. All 20 tests pass. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,152 @@
|
|||||||
|
"""
|
||||||
|
WorkflowStateManager - 步骤级状态持久化管理器
|
||||||
|
|
||||||
|
提供工作流状态的原子写入、备份轮转和崩溃恢复功能。
|
||||||
|
|
||||||
|
Purpose: 为工作流提供可靠的状态管理,支持断点续传和崩溃恢复
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from kb_pipeline.schemas import WorkflowState
|
||||||
|
|
||||||
|
|
||||||
|
class WorkflowStateManager:
|
||||||
|
"""工作流状态管理器
|
||||||
|
|
||||||
|
管理工作流状态的持久化,包括:
|
||||||
|
- 原子写入(使用临时文件 + os.replace)
|
||||||
|
- 备份轮转(workflow_state.json -> workflow_state.json.bak)
|
||||||
|
- 主文件损坏时从备份恢复
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
STATE_FILENAME: 状态文件名
|
||||||
|
content_dir: 内容目录路径
|
||||||
|
state_path: 状态文件完整路径
|
||||||
|
backup_path: 备份文件完整路径
|
||||||
|
"""
|
||||||
|
|
||||||
|
STATE_FILENAME = "workflow_state.json"
|
||||||
|
|
||||||
|
def __init__(self, content_dir: Path):
|
||||||
|
"""初始化状态管理器
|
||||||
|
|
||||||
|
Args:
|
||||||
|
content_dir: 内容目录路径,状态文件将保存在此目录下
|
||||||
|
"""
|
||||||
|
self.content_dir = Path(content_dir)
|
||||||
|
self.state_path = self.content_dir / self.STATE_FILENAME
|
||||||
|
self.backup_path = self.content_dir / f"{self.STATE_FILENAME}.bak"
|
||||||
|
|
||||||
|
def load(self) -> Optional[WorkflowState]:
|
||||||
|
"""加载工作流状态
|
||||||
|
|
||||||
|
加载策略:
|
||||||
|
1. 尝试加载主文件 (workflow_state.json)
|
||||||
|
2. 如果主文件损坏或不存在,尝试加载备份文件
|
||||||
|
3. 如果都不存在,返回 None
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
WorkflowState 实例,如果文件不存在则返回 None
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: 如果文件格式无效且备份也无法恢复
|
||||||
|
"""
|
||||||
|
def _try_load(path: Path) -> Optional[WorkflowState]:
|
||||||
|
"""尝试从指定路径加载状态"""
|
||||||
|
if not path.exists():
|
||||||
|
return None
|
||||||
|
|
||||||
|
try:
|
||||||
|
with open(path, 'r', encoding='utf-8') as f:
|
||||||
|
data = json.load(f)
|
||||||
|
return WorkflowState.model_validate(data)
|
||||||
|
except (json.JSONDecodeError, ValueError):
|
||||||
|
return None
|
||||||
|
|
||||||
|
# 先尝试主文件
|
||||||
|
state = _try_load(self.state_path)
|
||||||
|
if state is not None:
|
||||||
|
return state
|
||||||
|
|
||||||
|
# 主文件失败,尝试备份
|
||||||
|
state = _try_load(self.backup_path)
|
||||||
|
return state
|
||||||
|
|
||||||
|
def save(self, state: WorkflowState) -> None:
|
||||||
|
"""原子保存工作流状态
|
||||||
|
|
||||||
|
保存策略(原子写入 + 备份轮转):
|
||||||
|
1. 创建父目录(如果不存在)
|
||||||
|
2. 如果状态文件存在,将其重命名为备份文件
|
||||||
|
3. 写入临时文件
|
||||||
|
4. 原子替换(os.replace)
|
||||||
|
|
||||||
|
原子写入确保:
|
||||||
|
- 读取者永远不会看到部分写入的文件
|
||||||
|
- 崩溃时要么是旧文件,要么是新文件
|
||||||
|
|
||||||
|
Args:
|
||||||
|
state: 要保存的工作流状态
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
OSError: 文件操作失败(权限、磁盘空间等)
|
||||||
|
"""
|
||||||
|
# 1. 创建父目录
|
||||||
|
self.content_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
# 2. 备份轮转:如果状态文件存在,移动到备份
|
||||||
|
if self.state_path.exists():
|
||||||
|
# 删除旧备份
|
||||||
|
if self.backup_path.exists():
|
||||||
|
self.backup_path.unlink()
|
||||||
|
# 将当前状态移动到备份
|
||||||
|
self.state_path.rename(self.backup_path)
|
||||||
|
|
||||||
|
# 3. 写入临时文件(在相同目录下,确保原子替换有效)
|
||||||
|
fd, temp_path = tempfile.mkstemp(dir=self.content_dir, suffix='.tmp')
|
||||||
|
try:
|
||||||
|
with os.fdopen(fd, 'w', encoding='utf-8') as f:
|
||||||
|
f.write(state.model_dump_json(indent=2))
|
||||||
|
# 4. 原子替换
|
||||||
|
os.replace(temp_path, self.state_path)
|
||||||
|
except Exception:
|
||||||
|
# 清理临时文件
|
||||||
|
if os.path.exists(temp_path):
|
||||||
|
os.unlink(temp_path)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def get_or_create(
|
||||||
|
self,
|
||||||
|
content_id: str,
|
||||||
|
platform: str,
|
||||||
|
source_url: str
|
||||||
|
) -> WorkflowState:
|
||||||
|
"""获取现有状态或创建新状态
|
||||||
|
|
||||||
|
幂等操作:如果状态已存在,返回现有状态;否则创建新状态。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
content_id: 内容唯一标识
|
||||||
|
platform: 平台名称
|
||||||
|
source_url: 原始链接
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
WorkflowState 实例
|
||||||
|
"""
|
||||||
|
existing = self.load()
|
||||||
|
if existing is not None:
|
||||||
|
return existing
|
||||||
|
|
||||||
|
# 创建新状态
|
||||||
|
return WorkflowState(
|
||||||
|
content_id=content_id,
|
||||||
|
platform=platform,
|
||||||
|
source_url=source_url
|
||||||
|
)
|
||||||
@@ -281,7 +281,7 @@ class TestBackupFallback:
|
|||||||
state_file = tmp_path / "workflow_state.json"
|
state_file = tmp_path / "workflow_state.json"
|
||||||
backup_file = tmp_path / "workflow_state.json.bak"
|
backup_file = tmp_path / "workflow_state.json.bak"
|
||||||
|
|
||||||
# Create valid backup
|
# Create valid backup by saving twice (backup is created on second save)
|
||||||
valid_state = WorkflowState(
|
valid_state = WorkflowState(
|
||||||
content_id="backup-valid-001",
|
content_id="backup-valid-001",
|
||||||
platform="bilibili",
|
platform="bilibili",
|
||||||
@@ -289,11 +289,22 @@ class TestBackupFallback:
|
|||||||
)
|
)
|
||||||
manager.save(valid_state)
|
manager.save(valid_state)
|
||||||
|
|
||||||
|
# Save again to create backup
|
||||||
|
valid_state2 = WorkflowState(
|
||||||
|
content_id="backup-valid-002",
|
||||||
|
platform="bilibili",
|
||||||
|
source_url="https://www.bilibili.com/video/BV888"
|
||||||
|
)
|
||||||
|
manager.save(valid_state2)
|
||||||
|
|
||||||
|
# Now backup exists with "backup-valid-001", main has "backup-valid-002"
|
||||||
|
assert backup_file.exists()
|
||||||
|
|
||||||
# Corrupt main file
|
# Corrupt main file
|
||||||
with open(state_file, 'w', encoding='utf-8') as f:
|
with open(state_file, 'w', encoding='utf-8') as f:
|
||||||
f.write("{ invalid json content }")
|
f.write("{ invalid json content }")
|
||||||
|
|
||||||
# Should fall back to backup
|
# Should fall back to backup (which has backup-valid-001)
|
||||||
loaded = manager.load()
|
loaded = manager.load()
|
||||||
|
|
||||||
assert loaded is not None
|
assert loaded is not None
|
||||||
|
|||||||
Reference in New Issue
Block a user