Files
Agent-System/internal/taskrun/manager.go
T
weichen 265616e997 feat(stream): 打通小队过程事件到持久化 detail 帧
pi-native 与 runner NDJSON 把推理和工具调用投影为 detail,任务流始终落库;非 200 的流式尝试回落到一次性 /v1/run,避免改变既有错误分类。
2026-09-28 10:30:52 +00:00

849 lines
23 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package taskrun
import (
"context"
"errors"
"fmt"
"log/slog"
"sync"
"time"
"chat-agent-system/internal/bizdev"
"chat-agent-system/internal/ctxmgr"
"chat-agent-system/internal/errcode"
"chat-agent-system/internal/projectworkflow"
"chat-agent-system/internal/stream"
)
// Executor 执行具体任务逻辑(由 server 注入)。
type Executor func(ctx context.Context, runID string, req bizdev.TaskRunRequest, streamer *stream.Streamer) (*bizdev.TaskRunResult, error)
// ResumeExecutor 续跑既有工作区。
type ResumeExecutor func(ctx context.Context, runID string, workDir string, streamer *stream.Streamer) (*bizdev.TaskRunResult, error)
// StageExecutor 执行成熟度阶段任务序列。
type StageExecutor func(ctx context.Context, runID string, req bizdev.StageRunRequest, streamer *stream.Streamer) (*bizdev.TaskRunResult, error)
// Finalizer 任务成功后投影运行记录与聊天消息(由 server 注入)。
type Finalizer func(ctx context.Context, record *RunRecord, result *bizdev.TaskRunResult) (terminalStatus string, err error)
// Manager 后台任务运行观察者:SSE 断线不取消 worker;可选 PG 持久化支持重启重放。
type Manager struct {
mu sync.RWMutex
runs map[string]*activeRun
store Store
workerID string
exec Executor
resume ResumeExecutor
stage StageExecutor
finalize Finalizer
cmds *Commands
pwSvc *projectworkflow.Service
ctxMgr *ctxmgr.ContextManager
idemMu sync.Mutex
idemLocks map[string]*idemLockEntry // keyed by idempotency_key; TTL-evicted in cleanupLoop
}
type idemLockEntry struct {
mu sync.Mutex
lastUsed time.Time
}
const runRetention = time.Hour
const idemLockRetention = time.Hour
var (
ErrRunNotResumable = errors.New("taskrun: run is not resumable")
ErrRunAlreadyActive = errors.New("taskrun: run already active in memory")
)
func NewManager(pwSvc *projectworkflow.Service, ctxMgr *ctxmgr.ContextManager) *Manager {
m := &Manager{
runs: make(map[string]*activeRun),
store: noopStore{},
workerID: newWorkerID(),
pwSvc: pwSvc,
ctxMgr: ctxMgr,
idemLocks: make(map[string]*idemLockEntry),
}
go m.cleanupLoop()
return m
}
func (m *Manager) lockIdempotencyKey(key string) func() {
if key == "" {
return func() {}
}
m.idemMu.Lock()
lk, ok := m.idemLocks[key]
if !ok {
lk = &idemLockEntry{}
m.idemLocks[key] = lk
}
lk.lastUsed = time.Now()
m.idemMu.Unlock()
lk.mu.Lock()
return func() {
lk.mu.Unlock()
m.idemMu.Lock()
lk.lastUsed = time.Now()
m.idemMu.Unlock()
}
}
// SetStore 注入持久化存储(PG);启动时标记 lease 过期的 running 任务为 interrupted。
func (m *Manager) SetStore(store Store) {
if m == nil {
return
}
if store == nil {
store = noopStore{}
}
m.store = store
if store.Available() {
if n, err := store.MarkOrphanedRunning(context.Background()); err != nil {
slog.Warn("taskrun_orphan_recovery_failed", "error", err)
} else if n > 0 {
slog.Info("taskrun_orphans_marked_interrupted", "count", n)
}
}
}
func (m *Manager) cleanupLoop() {
ticker := time.NewTicker(10 * time.Minute)
defer ticker.Stop()
for range ticker.C {
m.evictStaleRuns()
m.evictStaleIdemLocks()
}
}
func (m *Manager) evictStaleRuns() {
cutoff := time.Now().Add(-runRetention)
m.mu.Lock()
defer m.mu.Unlock()
for id, ar := range m.runs {
ar.mu.RLock()
status := ar.record.Status
updated := ar.record.UpdatedAt
ar.mu.RUnlock()
if updated.Before(cutoff) && status != "running" {
delete(m.runs, id)
}
}
}
func (m *Manager) evictStaleIdemLocks() {
cutoff := time.Now().Add(-idemLockRetention)
m.idemMu.Lock()
defer m.idemMu.Unlock()
for key, lk := range m.idemLocks {
if lk.lastUsed.After(cutoff) {
continue
}
// Skip locks currently held; TryLock avoids blocking cleanup on a live claim.
if !lk.mu.TryLock() {
continue
}
if !lk.lastUsed.After(cutoff) {
delete(m.idemLocks, key)
}
lk.mu.Unlock()
}
}
func (m *Manager) SetExecutor(exec Executor, resume ResumeExecutor, stage StageExecutor) {
m.exec = exec
m.resume = resume
m.stage = stage
}
func (m *Manager) SetFinalizer(f Finalizer) {
m.finalize = f
}
func (m *Manager) SetCommands(cmds *Commands) {
m.cmds = cmds
if cmds != nil {
m.finalize = cmds.Finalize
}
}
func (m *Manager) CreateAndStart(parent context.Context, runID string, record *RunRecord, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest, resumeWorkDir string) (*RunRecord, error) {
if record == nil {
record = &RunRecord{}
}
if runID != "" {
record.RunID = runID
}
claimed, created, err := m.ClaimRun(parent, record)
if err != nil {
return nil, err
}
if !created {
return claimed, nil
}
return m.StartClaimedRun(parent, claimed, taskReq, stageReq, resumeWorkDir)
}
// ClaimRun 原子占位:在写聊天 / 启 worker 之前确认本请求是否为 owner。
// created=false 表示已有同 idempotency_key 的 run,调用方不得再写消息。
func (m *Manager) ClaimRun(parent context.Context, record *RunRecord) (*RunRecord, bool, error) {
if record == nil {
return nil, false, fmt.Errorf("taskrun: record required")
}
if record.IdempotencyKey == "" && record.SessionID != "" && record.ClientTurnID != "" {
record.IdempotencyKey = BuildTaskRunIdempotencyKey(record.SessionID, record.ClientTurnID)
}
unlock := m.lockIdempotencyKey(record.IdempotencyKey)
defer unlock()
if record.IdempotencyKey != "" && m.store != nil && m.store.Available() {
if existing, err := m.store.GetByIdempotencyKey(parent, record.IdempotencyKey); err != nil {
return nil, false, err
} else if existing != nil {
return existing, false, nil
}
}
if record.RunID == "" {
record.RunID = fmt.Sprintf("run_%d", time.Now().UnixNano())
}
if record.Status == "" {
record.Status = "pending"
}
now := time.Now()
if record.CreatedAt.IsZero() {
record.CreatedAt = now
}
record.UpdatedAt = now
if record.Attempt <= 0 {
record.Attempt = 1
}
if m.store != nil && m.store.Available() {
if err := m.store.InsertRun(parent, record); err != nil {
if errors.Is(err, ErrIdempotencyConflict) && record.IdempotencyKey != "" {
existing, lookupErr := m.store.GetByIdempotencyKey(parent, record.IdempotencyKey)
if lookupErr != nil {
return nil, false, lookupErr
}
if existing != nil {
return existing, false, nil
}
}
return nil, false, err
}
}
cpy := *record
return &cpy, true, nil
}
// BindAssistantMsgID 在启 worker 前回填 assistant 消息 id,覆盖快速完成竞态。
func (m *Manager) BindAssistantMsgID(ctx context.Context, runID string, assistantMsgID int64) error {
if runID == "" || assistantMsgID <= 0 {
return nil
}
m.mu.RLock()
ar, ok := m.runs[runID]
m.mu.RUnlock()
if ok {
ar.mu.Lock()
ar.record.AssistantMsgID = assistantMsgID
ar.record.UpdatedAt = time.Now()
rec := *ar.record
ar.mu.Unlock()
if m.store != nil && m.store.Available() {
return m.store.UpdateRun(ctx, &rec)
}
return nil
}
if m.store != nil && m.store.Available() {
rec, err := m.store.GetRun(ctx, runID)
if err != nil {
return err
}
rec.AssistantMsgID = assistantMsgID
rec.UpdatedAt = time.Now()
return m.store.UpdateRun(ctx, rec)
}
return nil
}
// StartClaimedRun 仅对 ClaimRun 成功的 owner 启 worker;禁止在 Insert 冲突后调用。
func (m *Manager) StartClaimedRun(parent context.Context, record *RunRecord, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest, resumeWorkDir string) (*RunRecord, error) {
if record == nil || record.RunID == "" {
return nil, fmt.Errorf("taskrun: claimed record required")
}
record.Status = "running"
record.WorkerID = m.workerID
now := time.Now()
record.UpdatedAt = now
leaseUntil := now.Add(defaultLeaseTTL)
record.LeaseExpiresAt = &leaseUntil
if m.store != nil && m.store.Available() {
if err := m.store.UpdateRun(parent, record); err != nil {
return nil, fmt.Errorf("taskrun: persist running status before start: %w", err)
}
if err := m.store.AcquireLease(parent, record.RunID, m.workerID, defaultLeaseTTL); err != nil {
slog.Warn("taskrun_lease_acquire_failed", "run_id", record.RunID, "error", err)
}
}
workerCtx, cancel := context.WithCancel(context.WithoutCancel(parent))
ar := &activeRun{
record: record,
subs: make(map[int]chan StoredEvent),
cancel: cancel,
done: make(chan struct{}),
}
m.mu.Lock()
m.runs[record.RunID] = ar
m.mu.Unlock()
if m.pwSvc != nil && record.ProjectID > 0 {
startedBy := record.UserID
_, _ = m.pwSvc.CreateWorkflowRun(workerCtx, &projectworkflow.WorkflowRun{
RunID: record.RunID,
ProjectID: record.ProjectID,
Status: "running",
StartedBy: &startedBy,
StageKey: record.Stage,
TaskKey: record.TaskKey,
InputManifest: map[string]any{
"session_id": record.SessionID,
"client_turn_id": record.ClientTurnID,
"assistant_msg_id": record.AssistantMsgID,
},
})
}
m.runGuarded("taskrun_worker", func() {
m.runWorker(workerCtx, ar, taskReq, stageReq, resumeWorkDir)
})
cpy := *record
return &cpy, nil
}
func (m *Manager) runGuarded(name string, fn func()) {
go func() {
defer func() {
if rec := recover(); rec != nil {
slog.Error("taskrun_goroutine_panic", "name", name, "recover", rec)
}
}()
fn()
}()
}
func (m *Manager) runWorker(ctx context.Context, ar *activeRun, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest, resumeWorkDir string) {
defer close(ar.done)
hbCtx, hbCancel := context.WithCancel(ctx)
defer hbCancel()
if m.store != nil && m.store.Available() {
m.runGuarded("taskrun_heartbeat", func() {
m.heartbeatLoop(hbCtx, ar.record.RunID)
})
}
taskReq.RunID = ar.record.RunID
taskReq.SessionID = ar.record.SessionID
taskReq.ProjectID = ar.record.ProjectID
taskReq.UserID = ar.record.UserID
taskReq.ClientTurnID = ar.record.ClientTurnID
lastPhaseID := ""
listener := func(ev stream.StreamEvent) {
if ev.Type == stream.EventPhase && ev.Phase != nil && ev.Phase.PhaseID != "" {
lastPhaseID = ev.Phase.PhaseID
}
m.appendEvent(ar, ev)
}
// includeDetails=true:任务通道的 detail(思考链/工具调用过程)全量持久化并推送
// (2026-09-28 过程流式设计 C.6——"刷新后过程消失"违背透明可追溯原则;
// 认知负荷由前端默认折叠保证,旧前端忽略 detail 字段无破坏)。
ctxWithS, streamer := stream.WithListener(ctx, true, listener)
var taskResult *bizdev.TaskRunResult
var execErr error
start := time.Now()
if ar.record.Stage != "" && m.stage != nil {
stageReq.RunID = ar.record.RunID
stageReq.SessionID = ar.record.SessionID
stageReq.ProjectID = ar.record.ProjectID
stageReq.UserID = ar.record.UserID
stageReq.ClientTurnID = ar.record.ClientTurnID
taskResult, execErr = m.stage(ctxWithS, ar.record.RunID, stageReq, streamer)
} else if resumeWorkDir != "" && m.resume != nil {
taskResult, execErr = m.resume(ctxWithS, ar.record.RunID, resumeWorkDir, streamer)
} else if m.exec != nil {
taskResult, execErr = m.exec(ctxWithS, ar.record.RunID, taskReq, streamer)
} else {
execErr = fmt.Errorf("task executor not configured")
}
workspaceID := ""
if taskResult != nil && taskResult.WorkDir != "" {
workspaceID = taskResult.WorkDir
} else if resumeWorkDir != "" {
workspaceID = resumeWorkDir
}
if workspaceID != "" {
ar.mu.Lock()
ar.record.WorkspaceID = workspaceID
ar.record.UpdatedAt = time.Now()
ar.mu.Unlock()
m.persistRun(ar)
}
if execErr != nil {
if ctx.Err() == context.Canceled {
if m.cmds != nil {
_, _ = m.cmds.Cancel(context.Background(), ar.record)
}
info := stream.ClassifyError(fmt.Errorf("task cancelled"), "task_cancel", ar.record.RunID, ar.record.SessionID, workspaceID)
m.appendEvent(ar, stream.StreamEvent{Type: stream.EventError, ErrorInfo: &info})
m.finishRun(ar, "cancelled", workspaceID, execErr.Error(), nil, &info)
return
}
classifyWS := workspaceID
if classifyWS == "" && resumeWorkDir != "" {
classifyWS = resumeWorkDir
}
info := stream.ClassifyError(execErr, lastPhaseID, ar.record.RunID, ar.record.SessionID, classifyWS)
m.appendEvent(ar, stream.StreamEvent{Type: stream.EventError, ErrorInfo: &info})
durationMs := time.Since(start).Milliseconds()
evCopy := m.copyEvents(ar)
if m.cmds != nil {
failRes := taskResult
if failRes == nil {
failRes = &bizdev.TaskRunResult{Success: false, ErrorReason: execErr.Error(), WorkDir: workspaceID}
} else {
failRes.Success = false
if failRes.ErrorReason == "" {
failRes.ErrorReason = execErr.Error()
}
}
_, _ = m.cmds.FinalizeWithEvents(context.Background(), ar.record, failRes, evCopy, durationMs)
}
m.finishRun(ar, "failed", workspaceID, execErr.Error(), taskResult, &info)
return
}
// RUN-002: 先 finalize 落库,成功后再发 done 终态事件
terminalStatus := "completed"
durationMs := time.Since(start).Milliseconds()
evCopy := m.copyEvents(ar)
if m.cmds != nil && taskResult != nil {
if ts, finErr := m.cmds.FinalizeWithEvents(ctx, ar.record, taskResult, evCopy, durationMs); finErr != nil {
info := stream.ClassifyError(finErr, lastPhaseID, ar.record.RunID, ar.record.SessionID, workspaceID)
m.appendEvent(ar, stream.StreamEvent{Type: stream.EventError, ErrorInfo: &info})
m.finishRun(ar, "finalizing_failed", workspaceID, finErr.Error(), taskResult, &info)
return
} else if ts != "" {
terminalStatus = ts
}
} else if m.finalize != nil && taskResult != nil {
if ts, finErr := m.finalize(ctx, ar.record, taskResult); finErr != nil {
info := stream.ClassifyError(finErr, lastPhaseID, ar.record.RunID, ar.record.SessionID, workspaceID)
m.appendEvent(ar, stream.StreamEvent{Type: stream.EventError, ErrorInfo: &info})
m.finishRun(ar, "finalizing_failed", workspaceID, finErr.Error(), taskResult, &info)
return
} else if ts != "" {
terminalStatus = ts
}
} else if taskResult != nil && taskResult.QualityReport != nil && !taskResult.QualityReport.Passed {
terminalStatus = "quality_failed"
}
m.appendEvent(ar, stream.StreamEvent{
Type: stream.EventDone,
Done: &stream.DonePayload{
SessionID: "task_run",
WorkspaceID: workspaceID,
DurationMs: durationMs,
},
})
m.finishRun(ar, terminalStatus, workspaceID, "", taskResult, nil)
}
func (m *Manager) copyEvents(ar *activeRun) []StoredEvent {
if ar == nil {
return nil
}
ar.mu.RLock()
defer ar.mu.RUnlock()
if len(ar.events) == 0 {
return nil
}
out := make([]StoredEvent, len(ar.events))
copy(out, ar.events)
return out
}
func (m *Manager) finishRun(ar *activeRun, status, workspaceID, errMsg string, result *bizdev.TaskRunResult, errorInfo *stream.ErrorInfo) {
ar.mu.Lock()
ar.record.Status = status
ar.record.WorkspaceID = workspaceID
ar.record.ErrorMessage = errMsg
ar.record.UpdatedAt = time.Now()
runID := ar.record.RunID
ar.mu.Unlock()
slog.Info("taskrun_status_transition",
"run_id", runID,
"status", status,
"workspace_id", workspaceID,
"error", errMsg,
)
if status == "failed" || status == "cancelled" || status == "finalizing_failed" {
code := "TASKRUN_BACKGROUND"
class := errcode.ClassFatal
requestID := ""
phaseID := ""
userMessage := errMsg
if errorInfo != nil {
if errorInfo.Code != "" {
code = errorInfo.Code
}
if errorInfo.Class != "" {
class = string(errorInfo.Class)
}
requestID = errorInfo.RequestID
phaseID = errorInfo.PhaseID
if errorInfo.UserMessage != "" {
userMessage = errorInfo.UserMessage
}
}
errcode.Report(context.Background(), errcode.Event{
Code: code,
Class: class,
RequestID: requestID,
PhaseID: phaseID,
Module: "taskrun",
Source: "async",
WorkspaceID: workspaceID,
Upstream: errMsg,
UserMessage: userMessage,
})
}
if m.store != nil && m.store.Available() {
_ = m.store.ReleaseLease(context.Background(), runID, m.workerID)
}
ar.mu.Lock()
ar.record.WorkerID = ""
ar.record.LeaseExpiresAt = nil
ar.mu.Unlock()
m.persistRun(ar)
// 失败/取消路径:若 Finalize 未执行,仍更新 workflow_runs 与占位消息 metadata
if m.cmds == nil && (status == "failed" || status == "cancelled" || status == "finalizing_failed") {
if m.pwSvc != nil && ar.record.ProjectID > 0 {
_ = m.pwSvc.UpdateWorkflowRun(context.Background(), ar.record.RunID, status, "", "", errMsg)
}
if m.ctxMgr != nil && ar.record.AssistantMsgID > 0 {
meta := map[string]any{
"run_id": ar.record.RunID,
"workspace_id": workspaceID,
"task_status": status,
}
if errorInfo != nil {
meta["error_info"] = errorInfo
}
title := ""
if result != nil {
title = result.ArtifactTitle
}
if snap := BuildPhaseSnapshot(m.copyEvents(ar), 0, status, title); snap != nil {
meta["task_phase_snapshot"] = snap
}
content := "任务执行失败。"
if status == "cancelled" {
content = "任务已取消。"
} else if status == "finalizing_failed" {
content = "任务已生成,但结果落库失败,请联系管理员或稍后重试。"
}
if errMsg != "" {
content += "\n\n" + errMsg
}
_ = m.ctxMgr.FinalizeAssistantTaskMessage(context.Background(), ar.record.AssistantMsgID, content, meta, status)
}
}
_ = result // success path persistence handled by Finalizer
}
func (m *Manager) persistRun(ar *activeRun) {
if m.store == nil || !m.store.Available() || ar == nil {
return
}
ar.mu.RLock()
rec := *ar.record
ar.mu.RUnlock()
if err := m.store.UpdateRun(context.Background(), &rec); err != nil {
slog.Warn("taskrun_update_failed", "run_id", rec.RunID, "error", err)
}
}
func (m *Manager) appendEvent(ar *activeRun, ev stream.StreamEvent) {
ar.mu.Lock()
ar.nextSeq++
stored := StoredEvent{Seq: ar.nextSeq, Event: ev}
ar.events = append(ar.events, stored)
runID := ar.record.RunID
seq := ar.nextSeq
// 持锁发送:与 Unsubscribe 的 close 互斥,避免向已关闭 channel 发送 panic。
for _, ch := range ar.subs {
select {
case ch <- stored:
default:
}
}
ar.mu.Unlock()
if m.store != nil && m.store.Available() {
if err := m.store.AppendEvent(context.Background(), runID, seq, ev); err != nil {
slog.Warn("taskrun_event_persist_failed", "run_id", runID, "seq", seq, "error", err)
}
}
}
func (m *Manager) GetRun(runID string) (*RunRecord, error) {
m.mu.RLock()
ar, ok := m.runs[runID]
m.mu.RUnlock()
if ok {
ar.mu.RLock()
defer ar.mu.RUnlock()
cpy := *ar.record
return &cpy, nil
}
if m.store != nil && m.store.Available() {
return m.store.GetRun(context.Background(), runID)
}
return nil, fmt.Errorf("run not found")
}
// FindByIdempotencyKey 查找既有幂等 run(内存优先,再查 store)。
func (m *Manager) FindByIdempotencyKey(ctx context.Context, key string) (*RunRecord, error) {
if key == "" {
return nil, nil
}
m.mu.RLock()
for _, ar := range m.runs {
ar.mu.RLock()
match := ar.record != nil && ar.record.IdempotencyKey == key
var cpy RunRecord
if match {
cpy = *ar.record
}
ar.mu.RUnlock()
if match {
m.mu.RUnlock()
return &cpy, nil
}
}
m.mu.RUnlock()
if m.store != nil && m.store.Available() {
return m.store.GetByIdempotencyKey(ctx, key)
}
return nil, nil
}
func (m *Manager) Subscribe(runID string, afterSeq int64) (<-chan StoredEvent, []StoredEvent, error) {
m.mu.RLock()
ar, ok := m.runs[runID]
m.mu.RUnlock()
if ok {
ar.mu.Lock()
defer ar.mu.Unlock()
replay := make([]StoredEvent, 0)
for _, ev := range ar.events {
if ev.Seq > afterSeq {
replay = append(replay, ev)
}
}
ch := make(chan StoredEvent, 32)
ar.subSeq++
id := ar.subSeq
ar.subs[id] = ch
return ch, replay, nil
}
if m.store != nil && m.store.Available() {
if _, err := m.store.GetRun(context.Background(), runID); err != nil {
return nil, nil, fmt.Errorf("run not found")
}
events, err := m.store.ListEvents(context.Background(), runID, afterSeq)
if err != nil {
return nil, nil, err
}
ch := make(chan StoredEvent)
close(ch)
return ch, events, nil
}
return nil, nil, fmt.Errorf("run not found")
}
func (m *Manager) Unsubscribe(runID string, ch <-chan StoredEvent) {
m.mu.RLock()
ar, ok := m.runs[runID]
m.mu.RUnlock()
if !ok {
return
}
ar.mu.Lock()
defer ar.mu.Unlock()
for id, sub := range ar.subs {
if sub == ch {
delete(ar.subs, id)
close(sub)
break
}
}
}
// ResumeInterrupted 在用户确认后续跑 interrupted 任务(需有效 workspace 路径)。
func (m *Manager) ResumeInterrupted(parent context.Context, record *RunRecord, resumeWorkDir string, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest) (*RunRecord, error) {
if record == nil || record.RunID == "" {
return nil, fmt.Errorf("taskrun: run_id required")
}
if resumeWorkDir == "" {
return nil, fmt.Errorf("taskrun: workspace required")
}
if record.Status != "interrupted" {
return nil, ErrRunNotResumable
}
m.mu.RLock()
_, exists := m.runs[record.RunID]
m.mu.RUnlock()
if exists {
return nil, ErrRunAlreadyActive
}
nextAttempt := record.Attempt + 1
if nextAttempt <= 0 {
nextAttempt = 2
}
if m.store != nil && m.store.Available() {
if err := m.store.ResetForResume(parent, record.RunID, nextAttempt); err != nil {
return nil, err
}
}
now := time.Now()
leaseUntil := now.Add(defaultLeaseTTL)
cpy := *record
cpy.Status = "running"
cpy.ErrorMessage = ""
cpy.Attempt = nextAttempt
cpy.WorkerID = m.workerID
cpy.LeaseExpiresAt = &leaseUntil
cpy.UpdatedAt = now
if m.store != nil && m.store.Available() {
_ = m.store.UpdateRun(parent, &cpy)
if err := m.store.AcquireLease(parent, cpy.RunID, m.workerID, defaultLeaseTTL); err != nil {
slog.Warn("taskrun_resume_lease_failed", "run_id", cpy.RunID, "error", err)
}
}
workerCtx, cancel := context.WithCancel(context.WithoutCancel(parent))
ar := &activeRun{
record: &cpy,
subs: make(map[int]chan StoredEvent),
cancel: cancel,
done: make(chan struct{}),
}
m.mu.Lock()
m.runs[cpy.RunID] = ar
m.mu.Unlock()
m.runGuarded("taskrun_worker_resume", func() {
m.runWorker(workerCtx, ar, taskReq, stageReq, resumeWorkDir)
})
return &cpy, nil
}
func (m *Manager) Cancel(runID string) error {
rec, _ := m.GetRun(runID)
if m.cmds != nil && rec != nil {
if _, err := m.cmds.Cancel(context.Background(), rec); err != nil && !errors.Is(err, ErrTerminalConflict) {
return err
}
}
m.mu.RLock()
ar, ok := m.runs[runID]
m.mu.RUnlock()
if ok {
ar.cancel()
return nil
}
if rec != nil {
return nil
}
return fmt.Errorf("run not found")
}
// WaitRunDone 等待指定 run 的 worker 退出(runWorker 关闭 ar.done)。
// 若 run 从未启动或已从内存表移除,立即返回 nil。
func (m *Manager) WaitRunDone(ctx context.Context, runID string) error {
if m == nil || runID == "" {
return nil
}
m.mu.RLock()
ar, ok := m.runs[runID]
m.mu.RUnlock()
if !ok || ar == nil || ar.done == nil {
return nil
}
select {
case <-ar.done:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
// WaitIdle 等待当前内存中全部 active run 的 worker 退出。
func (m *Manager) WaitIdle(ctx context.Context) error {
if m == nil {
return nil
}
m.mu.RLock()
dones := make([]<-chan struct{}, 0, len(m.runs))
for _, ar := range m.runs {
if ar != nil && ar.done != nil {
dones = append(dones, ar.done)
}
}
m.mu.RUnlock()
for _, done := range dones {
select {
case <-done:
case <-ctx.Done():
return ctx.Err()
}
}
return nil
}
func (m *Manager) heartbeatLoop(ctx context.Context, runID string) {
ticker := time.NewTicker(defaultHeartbeatPeriod)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
if err := m.store.HeartbeatLease(context.Background(), runID, m.workerID, defaultLeaseTTL); err != nil {
slog.Warn("taskrun_heartbeat_failed", "run_id", runID, "error", err)
}
}
}
}