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"
|
||||
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(
|
||||
content_id="backup-valid-001",
|
||||
platform="bilibili",
|
||||
@@ -289,11 +289,22 @@ class TestBackupFallback:
|
||||
)
|
||||
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
|
||||
with open(state_file, 'w', encoding='utf-8') as f:
|
||||
f.write("{ invalid json content }")
|
||||
|
||||
# Should fall back to backup
|
||||
# Should fall back to backup (which has backup-valid-001)
|
||||
loaded = manager.load()
|
||||
|
||||
assert loaded is not None
|
||||
|
||||
Reference in New Issue
Block a user