scheduler.go

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package exegraph

import (
	"cmp"
	"context"
	"errors"
	"fmt"
	"log"
	"runtime"
	"runtime/debug"
	"slices"
	"sync"
	"time"

	"github.com/azure/azure-dev/cli/azd/internal/tracing"
	"github.com/azure/azure-dev/cli/azd/internal/tracing/events"
	"github.com/azure/azure-dev/cli/azd/internal/tracing/fields"
)

// ErrorPolicy controls how the scheduler handles step failures.
type ErrorPolicy int

const (
	// FailFast cancels all in-flight steps when the first step fails.
	FailFast ErrorPolicy = iota

	// ContinueOnError lets all independent steps run to completion and
	// aggregates errors at the end. Steps whose dependencies failed are skipped.
	ContinueOnError
)

// String returns the string representation of the ErrorPolicy.
func (p ErrorPolicy) String() string {
	switch p {
	case FailFast:
		return "fail_fast"
	case ContinueOnError:
		return "continue_on_error"
	default:
		return "unknown"
	}
}

// RunOptions configures the execution graph scheduler.
type RunOptions struct {
	// MaxConcurrency limits the number of steps running simultaneously.
	// Zero (the default) caps at GOMAXPROCS×2; a positive value is honored
	// as the upper bound (useful to bump or cap for IO-bound workloads).
	// Negative values are treated as zero.
	MaxConcurrency int

	// GroupConcurrency limits the number of simultaneously running steps in each
	// named [Step.ConcurrencyGroup]. Non-positive limits are ignored. Group limits
	// are enforced beneath MaxConcurrency and do not reserve worker capacity.
	GroupConcurrency map[string]int

	// ErrorPolicy determines behavior on step failure.
	ErrorPolicy ErrorPolicy

	// StepTimeout imposes a per-step deadline. When positive, each step's
	// context is wrapped with context.WithTimeout before execution. If the
	// step does not complete within the duration the context expires with
	// context.DeadlineExceeded, which the scheduler treats as a normal step
	// failure (respecting ErrorPolicy). Zero (the default) means no per-step
	// timeout is applied.
	StepTimeout time.Duration

	// OnStepStart is called (if non-nil) when a step begins execution.
	// It is invoked from worker goroutines and must be safe for concurrent use.
	OnStepStart func(stepName string)

	// OnStepDone is called (if non-nil) when a step reaches a terminal state,
	// including when it is skipped. It may be invoked from worker goroutines or
	// the scheduler coordinator and must be safe for concurrent use.
	OnStepDone func(stepName string, err error)
}

// Run executes all steps in the graph respecting dependency order, with bounded
// concurrency. It returns nil if all steps succeed, or a
// combined error describing all failures.
//
// The returned error may contain both genuine step failures and [StepSkippedError]
// values (for steps whose dependencies failed). Use [RunWithResult] when you need
// to distinguish failures from transitive skips via [RunResult.Steps].
func Run(ctx context.Context, g *Graph, opts RunOptions) (err error) {
	result := RunWithResult(ctx, g, opts)
	return result.Error
}

// RunWithResult executes all steps in the graph and returns a [RunResult]
// containing per-step timing, status, and aggregate error. This is the
// instrumented variant of [Run] for callers that need execution telemetry.
func RunWithResult(ctx context.Context, g *Graph, opts RunOptions) (result *RunResult) {
	result = &RunResult{}

	ctx, span := tracing.Start(ctx, events.ExeGraphRunEvent)
	defer func() { span.EndWithStatus(result.Error) }()

	if g == nil {
		result.Error = errors.New("exegraph: graph must not be nil")
		return result
	}

	setRunSpanAttributes(span, g, opts)

	if g.Len() == 0 {
		return result
	}

	if err := g.Validate(); err != nil {
		result.Error = fmt.Errorf("graph validation failed: %w", err)
		return result
	}

	result = execute(ctx, g, opts)
	return result
}

// execute implements an event-driven scheduler with a bounded worker pool.
// A single coordinator admits ready steps in work-conserving round-robin order
// across concurrency groups, then workers execute admitted steps and report
// completion. Capacity is never acquired by workers, so a saturated group
// cannot occupy workers while waiting for another step in that group to finish.
func execute(ctx context.Context, g *Graph, opts RunOptions) *RunResult {
	n := g.Len()
	result := &RunResult{
		Steps: make([]StepTiming, 0, n),
	}
	var timingMu sync.Mutex

	// Build in-degree map and reverse adjacency list (step → dependents).
	inDegree := make(map[string]int, n)
	dependents := make(map[string][]string, n)
	for _, s := range g.steps {
		if _, ok := inDegree[s.Name]; !ok {
			inDegree[s.Name] = 0
		}
		for _, dep := range s.DependsOn {
			inDegree[s.Name]++
			dependents[dep] = append(dependents[dep], s.Name)
		}
	}

	// Track failed steps for dependency-skip logic.
	failed := make(map[string]bool, n)
	var allErrors []error

	// Derive a cancellable context for FailFast tear-down.
	runCtx, runCancel := context.WithCancel(ctx)
	defer runCancel()

	// Worker pool size: when MaxConcurrency is explicitly set (>0), honor it
	// as the upper bound (still clamped by the step count). Otherwise default
	// to GOMAXPROCS*2 to avoid unbounded goroutine creation for the large
	// pure-CPU-bound case.
	numWorkers := min(n, runtime.GOMAXPROCS(0)*2)
	if opts.MaxConcurrency > 0 {
		numWorkers = min(n, opts.MaxConcurrency)
	}

	type stepCompletion struct {
		name  string
		err   error
		start time.Time
		end   time.Time
		// schedulerCanceled captures whether runCtx was already canceled at
		// the moment the step finished executing. This is observed by the
		// worker (not by the event loop) so a subsequent FailFast tear-down
		// from a *different* step's failure cannot mis-classify this step's
		// own per-step timeout or genuine failure as a scheduler cancellation.
		schedulerCanceled bool
	}

	workQueue := make(chan string, n)
	completions := make(chan stepCompletion, n)

	// Start bounded worker pool.
	var workerWg sync.WaitGroup
	for range numWorkers {
		workerWg.Go(func() {
			for name := range workQueue {
				// Skip execution if the run has been canceled (FailFast or parent
				// context). Without this, queued steps start with a canceled context
				// and waste time before discovering the cancellation.
				if err := runCtx.Err(); err != nil {
					// Emit lifecycle callbacks so consumers tracking step state see
					// a terminal event for every queued step (parity with the drain
					// loop's StepSkipped cascade below).
					safeNotifyStart(opts, name)
					safeNotifyDone(opts, name, err)
					now := time.Now()
					completions <- stepCompletion{
						name: name, err: err, start: now, end: now,
						schedulerCanceled: true,
					}
					continue
				}
				step := g.steps[name]
				start := time.Now()
				err := runStep(runCtx, step, opts)
				end := time.Now()
				// Snapshot runCtx state immediately after the step's Action returns.
				// This is the canonical "was this cancellation from the scheduler?"
				// signal — doing it later in the event loop risks racing with a
				// FailFast tear-down triggered by an unrelated step's failure,
				// which would incorrectly re-classify a genuine per-step timeout
				// or action failure as a scheduler cancellation.
				completions <- stepCompletion{
					name: name, err: err, start: start, end: end,
					schedulerCanceled: runCtx.Err() != nil,
				}
			}
		})
	}

	runStart := time.Now()

	// Ready steps are partitioned by concurrency group. The order within each
	// group preserves the graph's critical-path priority. The group cursor
	// rotates after every admission so a continuously ready group cannot starve
	// another group, while an uncontested group can still consume all capacity.
	priorityOrder := g.priorityOrder()
	priorityRank := make(map[string]int, len(priorityOrder))
	for rank, name := range priorityOrder {
		priorityRank[name] = rank
	}
	compareReady := func(a, b string) int {
		return cmp.Compare(priorityRank[a], priorityRank[b])
	}
	readyByGroup := make(map[string][]string)
	var groupOrder []string
	groupActive := make(map[string]int)
	groupCursor := 0
	readyCount := 0
	inflight := 0

	enqueueReady := func(names []string) {
		names = slices.Clone(names)
		slices.SortFunc(names, compareReady)

		touchedGroups := make(map[string]struct{})
		for _, name := range names {
			group := g.steps[name].ConcurrencyGroup
			if _, ok := readyByGroup[group]; !ok {
				groupOrder = append(groupOrder, group)
			}
			readyByGroup[group] = append(readyByGroup[group], name)
			touchedGroups[group] = struct{}{}
			readyCount++
		}
		for group := range touchedGroups {
			slices.SortFunc(readyByGroup[group], compareReady)
		}
	}

	groupHasCapacity := func(group string) bool {
		limit := opts.GroupConcurrency[group]
		return limit <= 0 || groupActive[group] < limit
	}

	nextReady := func() (string, bool) {
		for offset := range len(groupOrder) {
			index := (groupCursor + offset) % len(groupOrder)
			group := groupOrder[index]
			queue := readyByGroup[group]
			if len(queue) == 0 || !groupHasCapacity(group) {
				continue
			}

			name := queue[0]
			readyByGroup[group] = queue[1:]
			readyCount--
			groupCursor = (index + 1) % len(groupOrder)
			return name, true
		}
		return "", false
	}

	dispatchReady := func() {
		for inflight < numWorkers && readyCount > 0 {
			name, ok := nextReady()
			if !ok {
				return
			}
			groupActive[g.steps[name].ConcurrencyGroup]++
			inflight++
			workQueue <- name
		}
	}

	releaseCapacity := func(name string) {
		group := g.steps[name].ConcurrencyGroup
		groupActive[group]--
	}

	// Seed ready queues with zero in-degree steps, sorted by transitive
	// dependent count descending (critical-path heuristic).
	var initialReady []string
	for _, name := range priorityOrder {
		if inDegree[name] == 0 {
			initialReady = append(initialReady, name)
		}
	}
	enqueueReady(initialReady)
	dispatchReady()

	// Event loop: process completions as they arrive. Each completion may
	// unblock successors whose in-degree drops to zero.
	for inflight > 0 {
		comp := <-completions
		inflight--
		releaseCapacity(comp.name)

		status, isRealFailure := classifyStepResult(comp.err, comp.schedulerCanceled)
		timingMu.Lock()
		result.Steps = append(result.Steps, StepTiming{
			Name:     comp.name,
			Status:   status,
			Start:    comp.start,
			End:      comp.end,
			Duration: comp.end.Sub(comp.start),
			Tags:     g.steps[comp.name].Tags,
			Err:      comp.err,
		})
		timingMu.Unlock()

		if comp.err != nil {
			if isRealFailure {
				failed[comp.name] = true
				allErrors = append(allErrors, comp.err)
			}

			if isRealFailure && opts.ErrorPolicy == FailFast {
				runCancel()
				// Drain remaining in-flight work.
				for inflight > 0 {
					r := <-completions
					inflight--
					releaseCapacity(r.name)
					drainStatus, drainIsReal := classifyStepResult(r.err, r.schedulerCanceled)
					if drainIsReal {
						allErrors = append(allErrors, r.err)
					}
					timingMu.Lock()
					result.Steps = append(result.Steps, StepTiming{
						Name:     r.name,
						Status:   drainStatus,
						Start:    r.start,
						End:      r.end,
						Duration: r.end.Sub(r.start),
						Tags:     g.steps[r.name].Tags,
						Err:      r.err,
					})
					timingMu.Unlock()
				}
				break
			}
		}

		// Decrement in-degrees of dependents; enqueue newly ready steps.
		//
		// When a step succeeds, iterate dependents read-only (no allocation).
		// When a step fails under ContinueOnError, skip propagation may append
		// to the queue — use a clone to avoid aliasing the shared slice.
		//
		// SAFETY: This loop runs in the single coordinator goroutine. The
		// `queue = append(queue, ...)` pattern (growing the slice while
		// iterating by index) is safe here because there is no concurrent
		// access. Do not refactor to parallel processing without replacing
		// this with a proper work-stealing queue.
		deps := dependents[comp.name]
		needsClone := false
		var readyBatch []string // collect newly ready steps for grouped enqueue
		for i := 0; i < len(deps); i++ {
			dep := deps[i]
			inDegree[dep]--
			if inDegree[dep] == 0 {
				if hasFailedDep(g.steps[dep], failed) {
					failed[dep] = true
					skipErr := &StepSkippedError{StepName: dep}
					allErrors = append(allErrors, skipErr)
					safeNotifyDone(opts, dep, skipErr)

					// Record skipped step timing (zero duration).
					now := time.Now()
					timingMu.Lock()
					result.Steps = append(result.Steps, StepTiming{
						Name:   dep,
						Status: StepSkipped,
						Start:  now,
						End:    now,
						Tags:   g.steps[dep].Tags,
						Err:    skipErr,
					})
					timingMu.Unlock()

					// Cascade: process skipped step's dependents too.
					// Clone on first mutation to avoid aliasing the shared dependents slice.
					if !needsClone {
						deps = slices.Clone(deps)
						needsClone = true
					}
					deps = append(deps, dependents[dep]...)
				} else {
					readyBatch = append(readyBatch, dep)
				}
			}
		}

		// enqueueReady sorts each touched group by priority, so newly ready
		// steps merge into their group's queue in critical-path order.
		enqueueReady(readyBatch)
		dispatchReady()
	}

	close(workQueue)
	workerWg.Wait()

	result.TotalDuration = time.Since(runStart)

	// FailFast / parent-cancel completeness: synthesize StepSkipped timings
	// for any steps that never appeared in result.Steps. Without this, a
	// FailFast tear-down leaves the post-break unqueued downstream steps
	// invisible to result.Steps, so consumers iterating timings (e.g. log
	// lines, dashboards) see an incomplete picture — a step that was
	// structurally blocked by the failure is indistinguishable from a step
	// that silently went missing due to a scheduler bug.
	if runCtx.Err() != nil {
		seen := make(map[string]bool, len(result.Steps))
		for _, st := range result.Steps {
			seen[st.Name] = true
		}
		now := time.Now()
		var skippedSteps []StepTiming
		for _, name := range g.order {
			if seen[name] {
				continue
			}
			skipErr := &StepSkippedError{StepName: name}
			safeNotifyDone(opts, name, skipErr)
			skippedSteps = append(skippedSteps, StepTiming{
				Name:   name,
				Status: StepSkipped,
				Start:  now,
				End:    now,
				Tags:   g.steps[name].Tags,
				Err:    skipErr,
			})
		}
		timingMu.Lock()
		result.Steps = append(result.Steps, skippedSteps...)
		timingMu.Unlock()
	}

	// Post-execution completeness check: verify every step was
	// either executed (completed/failed) or skipped. A corrupted in-degree map
	// or a missed dependency edge would leave steps permanently pending with no
	// way to detect the bug at runtime. This is a read-only safety net.
	//
	// Skip when the run was aborted (FailFast or parent cancel) — unreachable
	// steps with inDegree > 0 are expected in that case.
	if runCtx.Err() == nil {
		resolved := 0
		for _, name := range g.order {
			if inDegree[name] == 0 {
				resolved++
			}
		}
		if resolved != n {
			var stuckNames []string
			for _, name := range g.order {
				if inDegree[name] > 0 {
					stuckNames = append(stuckNames, fmt.Sprintf("%s(in-degree=%d)", name, inDegree[name]))
				}
			}
			allErrors = append(allErrors, fmt.Errorf(
				"exegraph: %d of %d steps never became ready; this indicates a scheduler bug: %v",
				n-resolved, n, stuckNames,
			))
		}
	}

	// If the parent ctx was canceled (user ctrl-C, deploy timeout, etc.) and
	// no step-level failures were recorded, surface the cancellation once at
	// the run boundary so callers see a non-nil RunResult.Error / Run() error.
	// Individual canceled steps are NOT added to allErrors (see the main loop
	// above) to avoid duplicating `context.Canceled` for every in-flight step.
	if runCtx.Err() != nil && ctx.Err() != nil && len(allErrors) == 0 {
		allErrors = append(allErrors, ctx.Err())
	}

	result.Error = errors.Join(allErrors...)
	return result
}

// setRunSpanAttributes records scheduler configuration on the run span.
func setRunSpanAttributes(span tracing.Span, g *Graph, opts RunOptions) {
	span.SetAttributes(
		fields.ExeGraphStepCountKey.Int(g.Len()),
		fields.ExeGraphMaxConcurrencyKey.Int(opts.MaxConcurrency),
		fields.ExeGraphErrorPolicyKey.String(opts.ErrorPolicy.String()),
	)

	if limit, ok := opts.GroupConcurrency["package"]; ok {
		span.SetAttributes(fields.ExeGraphPackageConcurrencyKey.Int(limit))
	}
	if limit, ok := opts.GroupConcurrency["provision"]; ok {
		span.SetAttributes(fields.ExeGraphProvisionConcurrencyKey.Int(limit))
	}
	if limit, ok := opts.GroupConcurrency["deploy"]; ok {
		span.SetAttributes(fields.ExeGraphDeployConcurrencyKey.Int(limit))
	}
}

// setStepSpanAttributes records identifying attributes for step on span. The step
// name and DependsOn entries embed user-chosen identifiers from azure.yaml (service
// names, layer names) and are hashed before emission per
// docs/specs/metrics-audit/privacy-review-checklist.md. Tags are a fixed internal
// vocabulary ("provision", "deploy", "cmdhook", ...) and are emitted raw.
func setStepSpanAttributes(span tracing.Span, step *Step) {
	span.SetAttributes(fields.StringHashed(fields.ExeGraphStepNameKey, step.Name))
	if len(step.DependsOn) > 0 {
		span.SetAttributes(fields.StringSliceHashed(fields.ExeGraphStepDepsKey, step.DependsOn))
	}
	if len(step.Tags) > 0 {
		span.SetAttributes(fields.ExeGraphStepTagsKey.StringSlice(step.Tags))
	}
}

// runStep executes a single step with panic recovery, tracing, and callbacks.
func runStep(ctx context.Context, step *Step, opts RunOptions) (stepErr error) {
	ctx, span := tracing.Start(ctx, events.ExeGraphStepEvent)

	setStepSpanAttributes(span, step)

	if opts.StepTimeout > 0 {
		span.SetAttributes(fields.ExeGraphStepTimeoutKey.Int(int(opts.StepTimeout.Seconds())))
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, opts.StepTimeout)
		defer cancel()
	}

	defer func() {
		if r := recover(); r != nil {
			log.Printf("exegraph: recovered panic in step %q (this is likely an internal bug): %v", step.Name, r)
			stepErr = fmt.Errorf("step %q panicked: %v\n%s", step.Name, r, debug.Stack())
		}
		span.EndWithStatus(stepErr)
		safeNotifyDone(opts, step.Name, stepErr)
	}()

	safeNotifyStart(opts, step.Name)

	if err := step.Action(ctx); err != nil {
		return fmt.Errorf("step %q failed: %w", step.Name, err)
	}
	return nil
}

// hasFailedDep checks if any of a step's dependencies are in the failed set.
func hasFailedDep(s *Step, failed map[string]bool) bool {
	for _, dep := range s.DependsOn {
		if failed[dep] {
			return true
		}
	}
	return false
}

// classifyStepResult determines the execution status for a completed step.
// It returns the status and whether the error is a real failure (as opposed to
// a scheduler cancellation that should not count toward the failure set).
//
// A step whose own per-step timeout (StepTimeout) fired but whose parent context
// (runCtx) was NOT canceled is classified as StepFailed — the timeout is a
// genuine step-level failure, not a scheduler tear-down side effect.
func classifyStepResult(err error, schedulerCanceled bool) (StepStatus, bool) {
	if err == nil {
		return StepDone, false
	}
	if IsStepSkipped(err) {
		return StepSkipped, true
	}
	if (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) && schedulerCanceled {
		return StepSkipped, false
	}
	return StepFailed, true
}

func notifyStart(opts RunOptions, name string) {
	if opts.OnStepStart != nil {
		opts.OnStepStart(name)
	}
}

func notifyDone(opts RunOptions, name string, err error) {
	if opts.OnStepDone != nil {
		opts.OnStepDone(name, err)
	}
}

// safeNotifyStart wraps notifyStart with panic recovery. Without this, a
// panicking OnStepStart callback kills the worker goroutine and deadlocks
// the scheduler.
func safeNotifyStart(opts RunOptions, name string) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("exegraph: callback panic for %q: %v", name, r)
		}
	}()
	notifyStart(opts, name)
}

// safeNotifyDone wraps notifyDone with panic recovery. An unrecovered panic in
// a callback would kill the goroutine (and in Go, the entire process). This is
// used in the coordinator goroutine and in runStep's deferred cleanup where
// the caller-supplied OnStepDone is invoked.
func safeNotifyDone(opts RunOptions, name string, err error) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("exegraph: callback panic for %q: %v", name, r)
		}
	}()
	notifyDone(opts, name, err)
}