feat(checkpoint): 新增文件检查点持久化与恢复功能

This commit is contained in:
unclesam committed 2026-06-11 02:11:05 +08:00
1 parent 2bbd2b7b04
commit 33cd062d15
2 files changed
+128

No files matched your search

+62
View File
@@ -0,0 +1,62 @@
package GopherGraph
import (
"context"
"encoding/json"
"fmt"
"os"
"path/filepath"
)
// Checkpointer 定义了持久化状态的接口,S 是工作流的强类型状态。
type Checkpointer[S any] interface {
// Save 保存当前线程快照。
Save(ctx context.Context, threadID string, thread *Thread[S]) error
// Load 读取之前保存的线程快照,用于恢复执行。
Load(ctx context.Context, threadID string) (*Thread[S], error)
}
// FileCheckpointer 实现了基于本地 JSON 文件的状态持久化。
type FileCheckpointer[S any] struct {
dir string // 存放状态文件的文件夹路径
}
// NewFileCheckpointer 创建一个文件检查点管理器。
func NewFileCheckpointer[S any](dir string) (*FileCheckpointer[S], error) {
// 确保文件夹存在,不存在则创建
if err := os.MkdirAll(dir, 0755); err != nil {
return nil, fmt.Errorf("failed to create checkpoint dir: %w", err)
}
return &FileCheckpointer[S]{dir: dir}, nil
}
// Save 将 Thread 序列化为 JSON 并写入本地文件
func (fc *FileCheckpointer[S]) Save(ctx context.Context, threadID string, thread *Thread[S]) error {
path := filepath.Join(fc.dir, threadID+".json")
// 使用 MarshalIndent 方便人工阅读生成的 JSON 文件
data, err := json.MarshalIndent(thread, "", " ")
if err != nil {
return fmt.Errorf("failed to marshal thread: %w", err)
}
return os.WriteFile(path, data, 0644)
}
// Load 从本地文件中读取 JSON 并反序列化为 Thread。
func (fc *FileCheckpointer[S]) Load(ctx context.Context, threadID string) (*Thread[S], error) {
path := filepath.Join(fc.dir, threadID+".json")
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("failed to read thread file: %w", err)
}
var thread Thread[S]
if err := json.Unmarshal(data, &thread); err != nil {
return nil, fmt.Errorf("failed to unmarshal thread: %w", err)
}
return &thread, nil
}
+66
View File
@@ -251,3 +251,69 @@ func TestParallelExecution(t *testing.T) {
}
}
}
// TestFileCheckpointer 测试文件检查点管理器对工作流进度的持久化与恢复
func TestFileCheckpointer(t *testing.T) {
// 创建测试临时文件夹(测试结束后 Go 会自动清理)
tmpDir := t.TempDir()
fc, err := NewFileCheckpointer[TestState](tmpDir)
if err != nil {
t.Fatalf("创建文件检查点管理器失败: %v", err)
}
g := NewGraph[TestState]()
g.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) {
s.Value = 100
return s, nil
})
g.AddNode("B", func(ctx context.Context, s TestState) (TestState, error) {
s.Value += 200
return s, nil
})
g.AddEdge("A", "B")
g.AddInterrupt("B") // 暂停在 B 之前
cg, err := g.Compile()
if err != nil {
t.Fatalf("编译失败: %v", err)
}
// 1. 运行并触发中断
thread, err := cg.Start(context.Background(), "A", TestState{Value: 0})
if err != nil {
t.Fatalf("运行失败: %v", err)
}
if !thread.IsPaused || thread.State.Value != 100 {
t.Fatalf("未成功触发中断,当前值: %d", thread.State.Value)
}
// 2. 模拟人工介入前:将线程快照(内存数据)写入磁盘
sessionID := "test-session-123"
ctx := context.Background()
if err := fc.Save(ctx, sessionID, thread); err != nil {
t.Fatalf("保存状态失败: %v", err)
}
// 3. 从磁盘重新读取状态(模拟程序崩掉重启后,重新从文件加载)
loadedThread, err := fc.Load(ctx, sessionID)
if err != nil {
t.Fatalf("读取状态失败: %v", err)
}
// 校验反序列化出来的状态是否完好无损
if loadedThread.NextNode != "B" || !loadedThread.IsPaused || loadedThread.State.Value != 100 {
t.Fatalf("加载出的状态不符合预期: %+v", loadedThread)
}
// 4. 使用从磁盘加载出来的线程指针,传入新状态恢复执行
thread, err = cg.Resume(ctx, loadedThread, TestState{Value: 500})
if err != nil {
t.Fatalf("恢复执行失败: %v", err)
}
// 验证最终执行结果: 500 + 200 = 700
if !thread.IsFinished || thread.State.Value != 700 {
t.Errorf("执行结果错误,最终值: %d, 是否结束: %t", thread.State.Value, thread.IsFinished)
}
}