A self-hosted orchestration system crafted for power users. composia.xyz
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293package agent
import ( "context" "encoding/json" "errors" "fmt" "log" "os" "path/filepath" "sync" "time"
"connectrpc.com/connect" agentv1 "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/agent/v1" "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/agent/v1/agentv1connect" "forgejo.alexma.top/alexma233/composia/internal/core/config" "forgejo.alexma.top/alexma233/composia/internal/core/task" "google.golang.org/protobuf/types/known/timestamppb")
const taskLeaseRenewInterval = 20 * time.Second
type taskExecutionContextKey struct{}
type persistedTaskExecution struct { TaskID string `json:"task_id"` ExecutionID string `json:"execution_id"` Acknowledged bool `json:"acknowledged"` Completion *completion `json:"completion,omitempty"` BackupResults []persistedBackupResult `json:"backup_results,omitempty"`}
type persistedBackupResult struct { ServiceName string `json:"service_name"` DataName string `json:"data_name"` ArtifactRef string `json:"artifact_ref,omitempty"` Status task.Status `json:"status"` StartedAt time.Time `json:"started_at"` FinishedAt time.Time `json:"finished_at"` ErrorSummary string `json:"error_summary,omitempty"`}
type completion struct { Status task.Status `json:"status"` ErrorSummary string `json:"error_summary,omitempty"` FinishedAt time.Time `json:"finished_at"`}
type taskExecutionRuntime struct { mu sync.Mutex path string state persistedTaskExecution}
func taskExecutionStatePath(stateDir string) string { return filepath.Join(stateDir, "task-execution.json")}
func withTaskExecution(ctx context.Context, runtime *taskExecutionRuntime) context.Context { return context.WithValue(ctx, taskExecutionContextKey{}, runtime)}
func taskExecutionFromContext(ctx context.Context) *taskExecutionRuntime { runtime, _ := ctx.Value(taskExecutionContextKey{}).(*taskExecutionRuntime) return runtime}
func taskExecutionID(ctx context.Context) string { if runtime := taskExecutionFromContext(ctx); runtime != nil { return runtime.state.ExecutionID } return ""}
func taskExecutionTaskID(ctx context.Context) string { if runtime := taskExecutionFromContext(ctx); runtime != nil { return runtime.state.TaskID } return ""}
func persistTaskExecution(path string, state persistedTaskExecution) error { content, err := json.Marshal(state) if err != nil { return fmt.Errorf("encode task execution state: %w", err) } tempPath := path + ".tmp" file, err := os.OpenFile(tempPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600) //nolint:gosec if err != nil { return fmt.Errorf("write task execution state: %w", err) } if _, err := file.Write(content); err != nil { _ = file.Close() return fmt.Errorf("write task execution state: %w", err) } if err := file.Sync(); err != nil { _ = file.Close() return fmt.Errorf("sync task execution state: %w", err) } if err := file.Close(); err != nil { return fmt.Errorf("close task execution state: %w", err) } if err := os.Rename(tempPath, path); err != nil { return fmt.Errorf("activate task execution state: %w", err) } dir, err := os.Open(filepath.Dir(path)) //nolint:gosec if err != nil { return fmt.Errorf("open task execution state directory: %w", err) } defer func() { _ = dir.Close() }() if err := dir.Sync(); err != nil { return fmt.Errorf("sync task execution state directory: %w", err) } return nil}
func loadTaskExecution(path string) (persistedTaskExecution, error) { content, err := os.ReadFile(path) //nolint:gosec if errors.Is(err, os.ErrNotExist) { return persistedTaskExecution{}, nil } if err != nil { return persistedTaskExecution{}, fmt.Errorf("read task execution state: %w", err) } var state persistedTaskExecution if err := json.Unmarshal(content, &state); err != nil { return persistedTaskExecution{}, fmt.Errorf("decode task execution state: %w", err) } return state, nil}
func acknowledgeTaskExecution(ctx context.Context, client agentv1connect.AgentTaskServiceClient, taskID, executionID string) error { callCtx, cancel := context.WithTimeout(ctx, taskReportTimeout) defer cancel() _, err := client.AcknowledgeTask(callCtx, connect.NewRequest(&agentv1.AcknowledgeTaskRequest{TaskId: taskID, ExecutionId: executionID})) if err != nil { return fmt.Errorf("acknowledge task execution: %w", err) } return nil}
func isRejectedTaskExecution(err error) bool { code := connect.CodeOf(err) return code == connect.CodeFailedPrecondition || code == connect.CodeNotFound}
func recoverTaskExecution(ctx context.Context, cfg *config.AgentConfig, taskClient agentv1connect.AgentTaskServiceClient, reportClient agentv1connect.AgentReportServiceClient) error { path := taskExecutionStatePath(cfg.StateDir) state, err := loadTaskExecution(path) if err != nil || state.TaskID == "" { return err } if state.Completion == nil && !state.Acknowledged { if err := acknowledgeTaskExecution(ctx, taskClient, state.TaskID, state.ExecutionID); err != nil { if isRejectedTaskExecution(err) { return removeTaskExecutionState(path) } return err } else { state.Acknowledged = true if err := persistTaskExecution(path, state); err != nil { return err } } } runtime := &taskExecutionRuntime{path: path, state: state} for len(runtime.state.BackupResults) > 0 { if err := sendPersistedBackupResult(ctx, reportClient, runtime.state.TaskID, runtime.state.ExecutionID, runtime.state.BackupResults[0]); err != nil { return err } if err := runtime.removeFirstBackupResult(); err != nil { return err } } state = runtime.state if state.Completion == nil { state.Completion = &completion{Status: task.StatusFailed, ErrorSummary: "agent restarted during an accepted task execution; outcome may be partial", FinishedAt: time.Now().UTC()} runtime.state = state if err := persistTaskExecution(path, state); err != nil { return err } } request := &agentv1.ReportTaskStateRequest{TaskId: state.TaskID, ExecutionId: state.ExecutionID, Status: protoAgentTaskStatus(state.Completion.Status), ErrorSummary: state.Completion.ErrorSummary, FinishedAt: timestamppb.New(state.Completion.FinishedAt)} callCtx, cancel := context.WithTimeout(ctx, taskReportTimeout) defer cancel() if _, err := reportClient.ReportTaskState(callCtx, connect.NewRequest(request)); err != nil { if isRejectedTaskExecution(err) { archivePath, archiveErr := archiveTaskExecutionState(path) if archiveErr != nil { return fmt.Errorf("archive conflicting task completion: %w", archiveErr) } log.Printf("task completion conflicted with controller state and was archived: task_id=%s archive=%s", state.TaskID, archivePath) return nil } return fmt.Errorf("recover task completion: %w", err) } return removeTaskExecutionState(path)}
func archiveTaskExecutionState(path string) (string, error) { archivePath := filepath.Join(filepath.Dir(path), fmt.Sprintf("task-execution-conflict-%d.json", time.Now().UTC().UnixNano())) if err := os.Rename(path, archivePath); err != nil { return "", err } dir, err := os.Open(filepath.Dir(path)) //nolint:gosec if err != nil { return "", err } defer func() { _ = dir.Close() }() if err := dir.Sync(); err != nil { return "", err } return archivePath, nil}
func removeTaskExecutionState(path string) error { err := os.Remove(path) if errors.Is(err, os.ErrNotExist) { return nil } return err}
func renewTaskLease(ctx context.Context, client agentv1connect.AgentTaskServiceClient, taskID, executionID string) { ticker := time.NewTicker(taskLeaseRenewInterval) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: callCtx, cancel := context.WithTimeout(ctx, taskReportTimeout) _, err := client.RenewTaskLease(callCtx, connect.NewRequest(&agentv1.RenewTaskLeaseRequest{TaskId: taskID, ExecutionId: executionID})) cancel() if err != nil && ctx.Err() == nil { log.Printf("renew task lease failed: task_id=%s error=%v", taskID, err) } } }}
func (runtime *taskExecutionRuntime) storeCompletion(value completion) error { runtime.mu.Lock() defer runtime.mu.Unlock() if len(runtime.state.BackupResults) > 0 { return errors.New("task completion cannot be persisted before backup results are acknowledged") } if runtime.state.Completion != nil { if runtime.state.Completion.Status != value.Status { return errors.New("task completion already persisted with a different status") } return nil } runtime.state.Completion = &value return persistTaskExecution(runtime.path, runtime.state)}
func (runtime *taskExecutionRuntime) storeBackupResult(value persistedBackupResult) error { runtime.mu.Lock() defer runtime.mu.Unlock() for _, existing := range runtime.state.BackupResults { if existing.DataName == value.DataName { if existing.Status != value.Status || existing.ArtifactRef != value.ArtifactRef { return errors.New("backup result already persisted with different content") } return nil } } runtime.state.BackupResults = append(runtime.state.BackupResults, value) return persistTaskExecution(runtime.path, runtime.state)}
func (runtime *taskExecutionRuntime) removeFirstBackupResult() error { runtime.mu.Lock() defer runtime.mu.Unlock() if len(runtime.state.BackupResults) == 0 { return nil } runtime.state.BackupResults = runtime.state.BackupResults[1:] return persistTaskExecution(runtime.path, runtime.state)}
func (runtime *taskExecutionRuntime) clear() error { runtime.mu.Lock() defer runtime.mu.Unlock() err := os.Remove(runtime.path) if errors.Is(err, os.ErrNotExist) { return nil } return err}