orca/internal/engine/executor.go

package engine

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os/exec"
	"sync"
	"time"

	"github.com/google/uuid"

	"git.cloudinit.dev/coreci/orca/internal/model"
	"git.cloudinit.dev/coreci/orca/internal/store"
)

type Executor struct {
	jobs  *store.JobRepo
	tasks *store.TaskRepo
	log   *slog.Logger
	mu    sync.Mutex
}

func NewExecutor(jobs *store.JobRepo, tasks *store.TaskRepo, log *slog.Logger) *Executor {
	if log == nil {
		log = slog.Default()
	}
	return &Executor{jobs: jobs, tasks: tasks, log: log}
}

// Submit is the dispatch-friendly entry point (v0.2 P02). It parses
// the spec bytes as a minimal TaskSpec and runs a single task under
// a fresh job. Returns the job ID. This is intentionally simpler
// than the v0.1 Run() entry point — the cross-node dispatch wire
// format is a flat task (one process), not a multi-task job.
//
// The spec format is a JSON object with at least:
//
//	{ "name": "...", "command": "...", "args": [...], "env": [...] }
//
// All fields except command are optional.
func (e *Executor) Submit(ctx context.Context, specBytes []byte) (string, error) {
	type wireSpec struct {
		Name    string   `json:"name"`
		Command string   `json:"command"`
		Args    []string `json:"args"`
		Env     []string `json:"env"`
	}
	var ws wireSpec
	if err := json.Unmarshal(specBytes, &ws); err != nil {
		return "", fmt.Errorf("Executor.Submit: parse: %w", err)
	}
	if ws.Command == "" {
		return "", errors.New("Executor.Submit: spec.command is required")
	}
	if ws.Name == "" {
		ws.Name = "dispatched"
	}
	job := &model.Job{
		ID:     uuid.NewString(),
		Spec:   string(specBytes),
		Status: model.JobStatusPending,
	}
	ts := TaskSpec{
		Name:    ws.Name,
		Command: ws.Command,
		Args:    ws.Args,
		Env:     ws.Env,
	}
	if err := e.Run(ctx, job, []TaskSpec{ts}); err != nil {
		return job.ID, err
	}
	return job.ID, nil
}

// Status returns the current state of a job for the Status dispatch
// endpoint. The returned string is one of: "pending", "running",
// "complete", "failed", "stopped". Maps to model.JobStatus* values.
func (e *Executor) Status(ctx context.Context, jobID string) (string, error) {
	if e.jobs == nil {
		return "", errors.New("Executor.Status: nil job repo")
	}
	j, err := e.jobs.Get(ctx, jobID)
	if err != nil {
		return "", err
	}
	return string(j.Status), nil
}

type TaskSpec struct {
	Name    string
	Command string
	Args    []string
	Env     []string
}

func (e *Executor) Run(ctx context.Context, job *model.Job, specs []TaskSpec) error {
	e.mu.Lock()
	defer e.mu.Unlock()

	// Insert the job first so tasks can reference it via foreign key.
	if err := e.jobs.Insert(ctx, job); err != nil {
		return err
	}
	if err := e.jobs.UpdateStatus(ctx, job.ID, model.JobStatusRunning, 0); err != nil {
		return err
	}

	var (
		wg          sync.WaitGroup
		failedCount int
		exitCode    int
		mu          sync.Mutex
	)

	for _, ts := range specs {
		wg.Add(1)
		go func(ts TaskSpec) {
			defer wg.Done()
			if err := e.runOne(ctx, job, ts); err != nil {
				mu.Lock()
				failedCount++
				e.log.Error("task failed",
					slog.String("job_id", job.ID),
					slog.String("task", ts.Name),
					slog.String("error", err.Error()))
				mu.Unlock()
			}
		}(ts)
	}
	wg.Wait()

	if failedCount > 0 {
		exitCode = 1
		if err := e.jobs.UpdateStatus(ctx, job.ID, model.JobStatusFailed, exitCode); err != nil {
			return err
		}
		return fmt.Errorf("%d/%d tasks failed", failedCount, len(specs))
	}

	if err := e.jobs.UpdateStatus(ctx, job.ID, model.JobStatusComplete, 0); err != nil {
		return err
	}
	return nil
}

func (e *Executor) runOne(ctx context.Context, job *model.Job, ts TaskSpec) error {
	task := &model.Task{
		ID:      uuid.NewString(),
		JobID:   job.ID,
		Command: ts.Command,
		Args:    ts.Args,
		Env:     ts.Env,
		Status:  model.TaskStatusPending,
	}
	if err := e.tasks.Insert(ctx, task); err != nil {
		return err
	}

	cmd := exec.CommandContext(ctx, ts.Command, ts.Args...)
	cmd.Env = append(cmd.Environ(), ts.Env...)
	// WaitDelay (Go 1.25+) bounds the time spent waiting on a child process
	// that fails to exit after the context is canceled.
	cmd.WaitDelay = 5 * time.Second

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Start(); err != nil {
		_ = e.tasks.UpdateKilled(ctx, task.ID)
		return fmt.Errorf("start: %w", err)
	}

	if err := e.tasks.UpdateRunning(ctx, task.ID, cmd.Process.Pid); err != nil {
		e.log.Warn("update running failed", slog.String("error", err.Error()))
	}

	e.log.Info("task started",
		slog.String("job_id", job.ID),
		slog.String("task", ts.Name),
		slog.Int("pid", cmd.Process.Pid))

	done := make(chan error, 1)
	go func() { done <- cmd.Wait() }()

	select {
	case err := <-done:
		exitCode := 0
		if err != nil {
			if ee, ok := err.(*exec.ExitError); ok {
				exitCode = ee.ExitCode()
			} else {
				exitCode = 1
			}
		}
		_ = e.tasks.UpdateDone(ctx, task.ID, exitCode, stdout.String(), stderr.String())
		if err != nil {
			return err
		}
		return nil
	case <-ctx.Done():
		// WaitDelay (set above) gives the process a grace period to exit
		// cleanly before being killed.
		_ = e.tasks.UpdateKilled(ctx, task.ID)
		return ctx.Err()
	}
}