feat(checkpoint): 新增文件检查点持久化与恢复功能
This commit is contained in:
1 parent
2bbd2b7b04
commit
33cd062d15
2 files changed
+128
No files matched your search
@@ -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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user