296 lines
7.7 KiB
Go
296 lines
7.7 KiB
Go
package taskrun
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"chat-agent-system/internal/bizdev"
|
|
"chat-agent-system/internal/stream"
|
|
)
|
|
|
|
// memIdemStore 最小内存 store：支持 idempotency_key 唯一占位。
|
|
type memIdemStore struct {
|
|
mu sync.Mutex
|
|
byID map[string]*RunRecord
|
|
byKey map[string]string
|
|
events map[string][]StoredEvent
|
|
}
|
|
|
|
func newMemIdemStore() *memIdemStore {
|
|
return &memIdemStore{
|
|
byID: make(map[string]*RunRecord),
|
|
byKey: make(map[string]string),
|
|
events: make(map[string][]StoredEvent),
|
|
}
|
|
}
|
|
|
|
func (m *memIdemStore) Available() bool { return true }
|
|
func (m *memIdemStore) InsertRun(_ context.Context, rec *RunRecord) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if rec.IdempotencyKey != "" {
|
|
if _, ok := m.byKey[rec.IdempotencyKey]; ok {
|
|
return ErrIdempotencyConflict
|
|
}
|
|
m.byKey[rec.IdempotencyKey] = rec.RunID
|
|
}
|
|
cpy := *rec
|
|
m.byID[rec.RunID] = &cpy
|
|
return nil
|
|
}
|
|
func (m *memIdemStore) UpdateRun(_ context.Context, rec *RunRecord) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
cpy := *rec
|
|
m.byID[rec.RunID] = &cpy
|
|
return nil
|
|
}
|
|
func (m *memIdemStore) GetByIdempotencyKey(_ context.Context, key string) (*RunRecord, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
id, ok := m.byKey[key]
|
|
if !ok {
|
|
return nil, nil
|
|
}
|
|
rec := m.byID[id]
|
|
if rec == nil {
|
|
return nil, nil
|
|
}
|
|
cpy := *rec
|
|
return &cpy, nil
|
|
}
|
|
func (m *memIdemStore) GetRun(_ context.Context, runID string) (*RunRecord, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
rec := m.byID[runID]
|
|
if rec == nil {
|
|
return nil, errors.New("run not found")
|
|
}
|
|
cpy := *rec
|
|
return &cpy, nil
|
|
}
|
|
func (m *memIdemStore) AppendEvent(_ context.Context, runID string, seq int64, ev stream.StreamEvent) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.events[runID] = append(m.events[runID], StoredEvent{Seq: seq, Event: ev})
|
|
return nil
|
|
}
|
|
func (m *memIdemStore) ListEvents(_ context.Context, runID string, afterSeq int64) ([]StoredEvent, error) {
|
|
return nil, nil
|
|
}
|
|
func (m *memIdemStore) AcquireLease(context.Context, string, string, time.Duration) error {
|
|
return nil
|
|
}
|
|
func (m *memIdemStore) HeartbeatLease(context.Context, string, string, time.Duration) error {
|
|
return nil
|
|
}
|
|
func (m *memIdemStore) ReleaseLease(context.Context, string, string) error { return nil }
|
|
func (m *memIdemStore) ResetForResume(context.Context, string, int) error { return nil }
|
|
func (m *memIdemStore) MarkOrphanedRunning(context.Context) (int, error) { return 0, nil }
|
|
func (m *memIdemStore) ListActiveBySession(_ context.Context, sessionID string) ([]*RunRecord, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
out := make([]*RunRecord, 0)
|
|
for _, rec := range m.byID {
|
|
if rec == nil || rec.SessionID != sessionID {
|
|
continue
|
|
}
|
|
if rec.Status != "pending" && rec.Status != "running" {
|
|
continue
|
|
}
|
|
cpy := *rec
|
|
out = append(out, &cpy)
|
|
}
|
|
return out, nil
|
|
}
|
|
func (m *memIdemStore) LatestBySession(_ context.Context, sessionID string) (*RunRecord, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
var best *RunRecord
|
|
for _, rec := range m.byID {
|
|
if rec == nil || rec.SessionID != sessionID {
|
|
continue
|
|
}
|
|
if best == nil || rec.UpdatedAt.After(best.UpdatedAt) {
|
|
cpy := *rec
|
|
best = &cpy
|
|
}
|
|
}
|
|
return best, nil
|
|
}
|
|
|
|
func TestClaimRun_ConcurrentSameKey_OneOwner(t *testing.T) {
|
|
store := newMemIdemStore()
|
|
mgr := NewManager(nil, nil)
|
|
mgr.SetStore(store)
|
|
|
|
const n = 16
|
|
var owners atomic.Int32
|
|
var losers atomic.Int32
|
|
var wg sync.WaitGroup
|
|
key := BuildTaskRunIdempotencyKey("sess-conc", "turn-conc-1")
|
|
runIDs := make(chan string, n)
|
|
|
|
wg.Add(n)
|
|
for i := 0; i < n; i++ {
|
|
go func(i int) {
|
|
defer wg.Done()
|
|
rec := &RunRecord{
|
|
RunID: "",
|
|
SessionID: "sess-conc",
|
|
ClientTurnID: "turn-conc-1",
|
|
IdempotencyKey: key,
|
|
UserID: 1,
|
|
TaskKey: "intake",
|
|
Status: "pending",
|
|
}
|
|
claimed, created, err := mgr.ClaimRun(context.Background(), rec)
|
|
if err != nil {
|
|
t.Errorf("claim: %v", err)
|
|
return
|
|
}
|
|
if created {
|
|
owners.Add(1)
|
|
runIDs <- claimed.RunID
|
|
} else {
|
|
losers.Add(1)
|
|
runIDs <- claimed.RunID
|
|
}
|
|
}(i)
|
|
}
|
|
wg.Wait()
|
|
close(runIDs)
|
|
|
|
if owners.Load() != 1 {
|
|
t.Fatalf("want exactly 1 claim owner, got %d", owners.Load())
|
|
}
|
|
if losers.Load() != n-1 {
|
|
t.Fatalf("want %d losers, got %d", n-1, losers.Load())
|
|
}
|
|
first := ""
|
|
for id := range runIDs {
|
|
if first == "" {
|
|
first = id
|
|
} else if id != first {
|
|
t.Fatalf("all claimants must share same run_id, got %s vs %s", first, id)
|
|
}
|
|
}
|
|
if len(store.byID) != 1 {
|
|
t.Fatalf("store must have 1 run, got %d", len(store.byID))
|
|
}
|
|
}
|
|
|
|
func TestClaimThenStart_Concurrent_OneExecutor(t *testing.T) {
|
|
store := newMemIdemStore()
|
|
mgr := NewManager(nil, nil)
|
|
mgr.SetStore(store)
|
|
|
|
var execCalls atomic.Int32
|
|
mgr.SetExecutor(func(ctx context.Context, runID string, req bizdev.TaskRunRequest, st *stream.Streamer) (*bizdev.TaskRunResult, error) {
|
|
execCalls.Add(1)
|
|
time.Sleep(20 * time.Millisecond)
|
|
return &bizdev.TaskRunResult{Success: true, WorkDir: t.TempDir(), TaskKey: "intake"}, nil
|
|
}, nil, nil)
|
|
|
|
const n = 12
|
|
var chatWrites atomic.Int32
|
|
var starts atomic.Int32
|
|
key := BuildTaskRunIdempotencyKey("sess-exec", "turn-exec-1")
|
|
var wg sync.WaitGroup
|
|
wg.Add(n)
|
|
for i := 0; i < n; i++ {
|
|
go func() {
|
|
defer wg.Done()
|
|
rec := &RunRecord{
|
|
SessionID: "sess-exec",
|
|
ClientTurnID: "turn-exec-1",
|
|
IdempotencyKey: key,
|
|
UserID: 7,
|
|
TaskKey: "intake",
|
|
Status: "pending",
|
|
}
|
|
claimed, created, err := mgr.ClaimRun(context.Background(), rec)
|
|
if err != nil {
|
|
t.Errorf("claim: %v", err)
|
|
return
|
|
}
|
|
if !created {
|
|
return
|
|
}
|
|
// 模拟 handler：仅 owner 写 1 user + 1 assistant
|
|
chatWrites.Add(2)
|
|
claimed.AssistantMsgID = 42
|
|
_ = mgr.BindAssistantMsgID(context.Background(), claimed.RunID, 42)
|
|
starts.Add(1)
|
|
_, err = mgr.StartClaimedRun(context.Background(), claimed, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, "")
|
|
if err != nil {
|
|
t.Errorf("start: %v", err)
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
|
|
deadline := time.Now().Add(3 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
if execCalls.Load() >= 1 {
|
|
break
|
|
}
|
|
time.Sleep(20 * time.Millisecond)
|
|
}
|
|
if starts.Load() != 1 || chatWrites.Load() != 2 {
|
|
t.Fatalf("want 1 start + 2 chat writes, got starts=%d chats=%d", starts.Load(), chatWrites.Load())
|
|
}
|
|
if execCalls.Load() != 1 {
|
|
t.Fatalf("want 1 executor call, got %d", execCalls.Load())
|
|
}
|
|
if len(store.byID) != 1 {
|
|
t.Fatalf("want 1 persisted run, got %d", len(store.byID))
|
|
}
|
|
}
|
|
|
|
type failUpdateStore struct {
|
|
*memIdemStore
|
|
updateErr error
|
|
}
|
|
|
|
func (f *failUpdateStore) UpdateRun(ctx context.Context, rec *RunRecord) error {
|
|
if f.updateErr != nil {
|
|
return f.updateErr
|
|
}
|
|
return f.memIdemStore.UpdateRun(ctx, rec)
|
|
}
|
|
|
|
func TestStartClaimedRun_UpdateRunFailureDoesNotStartWorker(t *testing.T) {
|
|
base := newMemIdemStore()
|
|
store := &failUpdateStore{memIdemStore: base, updateErr: errors.New("persist boom")}
|
|
mgr := NewManager(nil, nil)
|
|
mgr.SetStore(store)
|
|
|
|
rec := &RunRecord{
|
|
SessionID: "sess-upd-fail",
|
|
ClientTurnID: "turn-upd-fail",
|
|
IdempotencyKey: BuildTaskRunIdempotencyKey("sess-upd-fail", "turn-upd-fail"),
|
|
UserID: 1,
|
|
TaskKey: "intake",
|
|
Status: "pending",
|
|
}
|
|
claimed, created, err := mgr.ClaimRun(context.Background(), rec)
|
|
if err != nil || !created {
|
|
t.Fatalf("claim: created=%v err=%v", created, err)
|
|
}
|
|
_, err = mgr.StartClaimedRun(context.Background(), claimed, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, "")
|
|
if err == nil {
|
|
t.Fatal("expected StartClaimedRun to fail when UpdateRun fails")
|
|
}
|
|
mgr.mu.Lock()
|
|
_, live := mgr.runs[claimed.RunID]
|
|
mgr.mu.Unlock()
|
|
if live {
|
|
t.Fatal("worker must not start when UpdateRun persistence fails")
|
|
}
|
|
}
|