bee/audit/internal/platform/scenario.go

package platform

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// ScenarioJob is one step of a ScenarioSpec: either a one-shot/blocking
// command ("command") or a periodic background sampler ("sampler") that
// keeps running until every command job in the scenario has finished (or
// the scenario's overall timeout elapses).
type ScenarioJob struct {
	Name string   `json:"name"`
	Type string   `json:"type"` // "command" | "sampler"
	Cmd  []string `json:"cmd"`
	// GPUIndices, if non-empty, replaces the literal token "{{gpus}}" in Cmd
	// with a comma-joined index list — lets a scenario file say
	// ["dcgmi","diag","-r","nvbandwidth","-i","{{gpus}}"] once instead of
	// hardcoding a specific host's GPU indices into the file.
	GPUIndices []int `json:"gpu_indices,omitempty"`
	// IntervalSec is the sampling period for a "sampler" job; ignored for
	// "command" jobs.
	IntervalSec int `json:"interval_sec,omitempty"`
	// Parallel, for a "command" job, means it starts alongside whichever
	// other parallel command jobs precede it instead of waiting for them to
	// finish first. Sequential (default) command jobs run in file order,
	// each waiting for any preceding parallel batch to finish first.
	Parallel bool `json:"parallel,omitempty"`
}

// ScenarioSpec is a user-authored test scenario: what to run, on which
// GPUs, and what to sample in the background while it runs. Parsed from
// plain JSON (see ParseScenarioJSON) so authoring one needs no tooling
// beyond a text editor and no new Go dependency for this codebase.
//
// Example — reproduce a hard reboot seen only when all GPUs run nvbandwidth
// together (not on any single socket alone), while watching PSU/IPMI
// sensors for a power-delivery correlation:
//
//	{
//	  "name": "nvbandwidth-all-gpu-power-watch",
//	  "timeout_sec": 1800,
//	  "jobs": [
//	    {"name": "ipmi-sensors", "type": "sampler", "interval_sec": 2,
//	     "cmd": ["ipmitool", "sensor"]},
//	    {"name": "gpu-power", "type": "sampler", "interval_sec": 2,
//	     "cmd": ["nvidia-smi", "--query-gpu=index,power.draw,temperature.gpu,clocks.sm", "--format=csv"]},
//	    {"name": "nvbandwidth-all", "type": "command",
//	     "cmd": ["dcgmi", "diag", "-r", "nvbandwidth", "-i", "{{gpus}}"],
//	     "gpu_indices": [0, 1, 2, 3, 4]}
//	  ]
//	}
type ScenarioSpec struct {
	Name       string        `json:"name"`
	TimeoutSec int           `json:"timeout_sec,omitempty"`
	Jobs       []ScenarioJob `json:"jobs"`
}

// ParseScenarioJSON parses and validates a scenario file's contents.
func ParseScenarioJSON(data []byte) (ScenarioSpec, error) {
	var spec ScenarioSpec
	if err := json.Unmarshal(data, &spec); err != nil {
		return ScenarioSpec{}, fmt.Errorf("parse scenario: %w", err)
	}
	if strings.TrimSpace(spec.Name) == "" {
		return ScenarioSpec{}, fmt.Errorf(`scenario: "name" is required`)
	}
	if len(spec.Jobs) == 0 {
		return ScenarioSpec{}, fmt.Errorf("scenario: at least one job is required")
	}
	seen := map[string]bool{}
	for i, j := range spec.Jobs {
		if strings.TrimSpace(j.Name) == "" {
			return ScenarioSpec{}, fmt.Errorf("scenario: job %d: \"name\" is required", i)
		}
		if seen[j.Name] {
			return ScenarioSpec{}, fmt.Errorf("scenario: duplicate job name %q", j.Name)
		}
		seen[j.Name] = true
		if j.Type != "command" && j.Type != "sampler" {
			return ScenarioSpec{}, fmt.Errorf("scenario: job %q: type must be \"command\" or \"sampler\", got %q", j.Name, j.Type)
		}
		if len(j.Cmd) == 0 {
			return ScenarioSpec{}, fmt.Errorf("scenario: job %q: \"cmd\" is required", j.Name)
		}
		if j.Type == "sampler" && j.IntervalSec <= 0 {
			return ScenarioSpec{}, fmt.Errorf("scenario: sampler job %q: \"interval_sec\" must be > 0", j.Name)
		}
	}
	return spec, nil
}

// RunScenario executes a ScenarioSpec: runs its "command" jobs (sequential
// unless marked Parallel, each blocking until it exits), while every
// "sampler" job runs concurrently in the background on its own interval
// until all command jobs have finished or TimeoutSec elapses. Every job's
// own output streams live to its own file under the returned run
// directory.
//
// Command jobs are wired through the same satJobBoundaryHook/
// satSyncBracketHook seams the SAT job runner uses (see sat.go), so a
// scenario run gets the same durability treatment: a crash mid-command
// still leaves whatever printed up to that point, and — if blackbox is
// running — that evidence is requested (and, for the sync-bracket hook,
// waited on) to reach removable media before/after the command runs rather
// than only at the end of an adaptive flush period.
func (s *System) RunScenario(ctx context.Context, baseDir string, spec ScenarioSpec, logFunc func(string)) (string, error) {
	if baseDir == "" {
		baseDir = "/var/log/bee-sat"
	}
	ts := time.Now().UTC().Format("20060102-150405")
	runDir := filepath.Join(baseDir, "scenario-"+sanitizeScenarioName(spec.Name)+"-"+ts)
	if err := os.MkdirAll(runDir, 0755); err != nil {
		return "", err
	}
	verboseLog := filepath.Join(runDir, "verbose.log")

	if spec.TimeoutSec > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, time.Duration(spec.TimeoutSec)*time.Second)
		defer cancel()
	}

	sampCtx, stopSamplers := context.WithCancel(ctx)
	defer stopSamplers()
	var samplerWG sync.WaitGroup
	for _, job := range spec.Jobs {
		if job.Type != "sampler" {
			continue
		}
		job := job
		samplerWG.Add(1)
		go func() {
			defer samplerWG.Done()
			runScenarioSampler(sampCtx, runDir, verboseLog, job, logFunc)
		}()
	}

	var (
		mu      sync.Mutex
		summary strings.Builder
		cmdWG   sync.WaitGroup
	)
	fmt.Fprintf(&summary, "scenario=%s\n", spec.Name)
	fmt.Fprintf(&summary, "run_at_utc=%s\n", time.Now().UTC().Format(time.RFC3339))

	runOne := func(job ScenarioJob) {
		err := runScenarioCommand(ctx, runDir, verboseLog, job, logFunc)
		mu.Lock()
		if err != nil {
			fmt.Fprintf(&summary, "%s_status=FAILED\n", job.Name)
			fmt.Fprintf(&summary, "%s_error=%s\n", job.Name, err.Error())
		} else {
			fmt.Fprintf(&summary, "%s_status=OK\n", job.Name)
		}
		mu.Unlock()
	}

	for _, job := range spec.Jobs {
		if job.Type != "command" {
			continue
		}
		if job.Parallel {
			cmdWG.Add(1)
			go func(j ScenarioJob) {
				defer cmdWG.Done()
				runOne(j)
			}(job)
			continue
		}
		// A sequential job waits for any parallel batch launched ahead of it
		// to finish first, so scenario-file order stays intuitive: parallel
		// jobs run together, the next sequential job starts only once they
		// (and any earlier sequential job) are done.
		cmdWG.Wait()
		runOne(job)
	}
	cmdWG.Wait()

	stopSamplers()
	samplerWG.Wait()

	if err := os.WriteFile(filepath.Join(runDir, "summary.txt"), []byte(summary.String()), 0644); err != nil {
		return "", err
	}
	return runDir, nil
}

func runScenarioCommand(ctx context.Context, runDir, verboseLog string, job ScenarioJob, logFunc func(string)) error {
	cmd := substituteGPUIndices(job.Cmd, job.GPUIndices)
	livePath := filepath.Join(runDir, sanitizeScenarioName(job.Name)+".log")

	if satSyncBracketHook != nil {
		if err := satSyncBracketHook(job.Name, "before"); err != nil && logFunc != nil {
			logFunc(fmt.Sprintf("%s: blackbox sync wait (before) did not complete cleanly: %v", job.Name, err))
		}
	}
	_, err := runSATCommandCtx(ctx, verboseLog, job.Name, cmd, nil, logFunc, livePath)
	if satJobBoundaryHook != nil {
		satJobBoundaryHook(job.Name)
	}
	if satSyncBracketHook != nil {
		if syncErr := satSyncBracketHook(job.Name, "after"); syncErr != nil && logFunc != nil {
			logFunc(fmt.Sprintf("%s: blackbox sync wait (after) did not complete cleanly: %v", job.Name, syncErr))
		}
	}
	return err
}

// runScenarioSampler runs job.Cmd once immediately (so even a scenario that
// finishes very quickly gets a baseline reading) and then every
// job.IntervalSec until ctx is done, appending each timestamped sample to
// its own file under runDir.
func runScenarioSampler(ctx context.Context, runDir, verboseLog string, job ScenarioJob, logFunc func(string)) {
	path := filepath.Join(runDir, sanitizeScenarioName(job.Name)+".log")
	f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		if logFunc != nil {
			logFunc(fmt.Sprintf("%s: could not open sampler output %s: %v", job.Name, path, err))
		}
		return
	}
	defer f.Close()

	cmd := substituteGPUIndices(job.Cmd, job.GPUIndices)
	sample := func() {
		out, cmdErr := satExecCommand(cmd[0], cmd[1:]...).CombinedOutput()
		ts := time.Now().UTC().Format(time.RFC3339Nano)
		fmt.Fprintf(f, "=== %s ===\n", ts)
		if cmdErr != nil {
			fmt.Fprintf(f, "error: %v\n", cmdErr)
		}
		_, _ = f.Write(out)
		if len(out) == 0 || out[len(out)-1] != '\n' {
			_, _ = f.Write([]byte("\n"))
		}
		_ = f.Sync()
	}

	sample()
	ticker := time.NewTicker(time.Duration(job.IntervalSec) * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			appendSATVerboseLog(verboseLog, fmt.Sprintf("[%s] sampler %s stopped", time.Now().UTC().Format(time.RFC3339), job.Name))
			return
		case <-ticker.C:
			sample()
		}
	}
}

// substituteGPUIndices replaces the literal token "{{gpus}}" in each cmd
// argument with a comma-joined GPU index list. Returns cmd unchanged
// (same backing values, no allocation) when gpuIndices is empty.
func substituteGPUIndices(cmd []string, gpuIndices []int) []string {
	if len(gpuIndices) == 0 {
		return cmd
	}
	joined := joinIndexList(gpuIndices)
	out := make([]string, len(cmd))
	for i, arg := range cmd {
		out[i] = strings.ReplaceAll(arg, "{{gpus}}", joined)
	}
	return out
}

// sanitizeScenarioName converts a scenario/job name into a safe filename
// component: lowercase alphanumerics, '-', and '_' pass through; anything
// else (spaces, punctuation) becomes '-'.
func sanitizeScenarioName(name string) string {
	var b strings.Builder
	for _, r := range strings.ToLower(strings.TrimSpace(name)) {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9', r == '-', r == '_':
			b.WriteRune(r)
		default:
			b.WriteRune('-')
		}
	}
	s := b.String()
	if s == "" {
		return "scenario"
	}
	return s
}