run.go

package internal

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/kitproj/kit/internal/proc"
	"github.com/kitproj/kit/internal/types"
	"github.com/kitproj/kit/internal/util"
	"github.com/pkg/browser"
	"k8s.io/utils/strings/slices"
)

var poisonPill = struct{}{}

func RunSubgraph(ctx context.Context, cancel context.CancelFunc, port int, openBrowser bool, logger *log.Logger, wf *types.Workflow, taskNames []string, tasksToSkip []string) error {

	// check that the task names are valid
	for _, name := range taskNames {
		if _, ok := wf.Tasks[name]; !ok {
			return fmt.Errorf("task %q not found in workflow", name)
		}
	}

	// validate all tasks
	for name, task := range wf.Tasks {
		if err := task.Validate(); err != nil {
			return fmt.Errorf("task %q is invalid: %w", name, err)
		}
	}

	// check skipped tasks are valid
	for _, name := range tasksToSkip {
		if _, ok := wf.Tasks[name]; !ok {
			return fmt.Errorf("skipped task %q not found in workflow", name)
		}
	}

	// name is last part of pwd
	pwd := os.Getenv("PWD")
	name := filepath.Base(pwd)

	dag := NewDAG[bool](name)
	for name, t := range wf.Tasks {
		dag.AddNode(name, true)
		for _, dependency := range t.Dependencies {
			dag.AddEdge(dependency, name)
		}
	}
	visited := dag.Subgraph(taskNames)

	taskByName := wf.Tasks
	subgraph := NewDAG[*TaskNode](name)
	for name := range visited {
		task := taskByName[name]

		logFile := filepath.Join("logs", fmt.Sprintf("%s.log", name))
		if task.Log != "" {
			logFile = task.Log
		}

		subgraph.AddNode(name, &TaskNode{
			Name:    name,
			logFile: logFile,
			Task:    task,
			Phase:   "pending",
			cancel:  func() {},
			mu:      &sync.Mutex{}})
		for _, parent := range dag.Parents[name] {
			subgraph.AddEdge(parent, name)
		}
	}

	events := make(chan any, len(subgraph.Nodes)*2)

	// schedule the tasks in the subgraph that are ready to run , this is done by sending the task name to the events channel of any task that does not have any parents
	for taskName := range subgraph.Nodes {
		if len(subgraph.Parents[taskName]) == 0 {
			events <- taskName
		}
	}

	if len(subgraph.Nodes) == 0 {
		logger.Println("no tasks to run")
		return nil
	}

	// create logs directory
	if err := os.MkdirAll("logs", 0755); err != nil && !errors.Is(err, os.ErrExist) {
		return fmt.Errorf("failed to create logs directory: %w", err)
	}

	// start a file watcher for each task
	for _, node := range subgraph.Nodes {

		// start watching files for changes
		watcher, err := fsnotify.NewWatcher()
		if err != nil {
			return fmt.Errorf("failed to create watcher: %w", err)
		}
		for _, source := range node.Task.Watch {
			watchPath := filepath.Join(node.Task.WorkingDir, source)
			info, err := os.Stat(watchPath)
			if err != nil {
				return fmt.Errorf("failed to stat %q: %w", source, err)
			}
			if info.IsDir() {
				// Walk the directory tree to add the directory and all subdirectories
				err := filepath.Walk(watchPath, func(path string, info os.FileInfo, err error) error {
					if err != nil {
						return err
					}
					if info.IsDir() {
						if err := watcher.Add(path); err != nil {
							return fmt.Errorf("failed to watch %q: %w", path, err)
						}
					}
					return nil
				})
				if err != nil {
					return fmt.Errorf("failed to walk %q: %w", source, err)
				}
			} else {
				// It's a file, watch it directly
				if err := watcher.Add(watchPath); err != nil {
					return fmt.Errorf("failed to watch %q: %w", source, err)
				}
			}
		}
		defer watcher.Close()

		go func() {
			debounceTimer := time.AfterFunc(0, func() {})
			defer debounceTimer.Stop()
			for {
				select {
				case <-ctx.Done():
					return
				case event := <-watcher.Events:
					// Handle file writes
					if event.Op&fsnotify.Write == fsnotify.Write {
						debounceTimer.Stop()
						debounceTimer = time.AfterFunc(100*time.Millisecond, func() {
							logger.Printf("[%s] %s changed, re-running\n", node.Name, event.Name)
							events <- node.Name
						})
					}
					// Handle new directories being created - add them to the watcher
					if event.Op&fsnotify.Create == fsnotify.Create {
						if info, err := os.Stat(event.Name); err == nil && info.IsDir() {
							_ = watcher.Add(event.Name)
						}
						debounceTimer.Stop()
						debounceTimer = time.AfterFunc(100*time.Millisecond, func() {
							logger.Printf("[%s] %s created, re-running\n", node.Name, event.Name)
							events <- node.Name
						})
					}
				}
			}
		}()
	}

	semaphores := util.NewSemaphores(wf.Semaphores)

	wg := &sync.WaitGroup{}

	statusEvents := make(chan *TaskNode, 100)

	if port > 0 {
		go StartServer(ctx, port, wg, subgraph, statusEvents)
		if openBrowser {
			if err := browser.OpenURL(fmt.Sprintf("http://localhost:%d", port)); err != nil {
				return fmt.Errorf("failed to open browser: %v", err)
			}
		}
	}

	stallTimers := map[string]*time.Timer{}
	for name, taskNode := range subgraph.Nodes {
		stalledTime := taskNode.Task.GetStalledTimeout()
		stallTimers[name] = time.AfterFunc(stalledTime, func() {
			if taskNode.Phase == "starting" || taskNode.Phase == "running" {
				// we suffix the message with "starting" so we can differentiate between a task that is starting and one that is running, later on we can change the message to "output received"
				// and restore the phase to "running" or "starting"
				taskNode.Message = fmt.Sprintf("no output for %s or more while %s", stalledTime, taskNode.Phase)
				taskNode.Phase = "stalled"
				logger.Printf("[%s] %s\n", taskNode.Name, taskNode.Message)
				statusEvents <- taskNode
			}
		})
	}

	allRunning := false
	graphCompleted := false

	for {
		select {
		case <-ctx.Done():

			logger.Println("waiting for all tasks to complete")

			wg.Wait()

			// if any task failed, we will return an error
			var failures []string
			for _, node := range subgraph.Nodes {

				color := 30
				faint := 0
				switch node.Phase {
				case "failed":
					// red
					color = 31
					faint = 1
					failures = append(failures, node.Name)
				case "pending", "waiting":
					faint = 2
				}

				logger.Printf("\033[%d;%dm[%s] (%s) %s\033[0m\n", faint, color, node.Name, node.Phase, node.Message)
			}

			if len(failures) > 0 {
				runLifecycleHook(context.Background(), types.Task{}, wf.Lifecycle.GetOnFailureHook(), os.Stdout, logger)
				return fmt.Errorf("failed tasks: %v", failures)
			}

			if graphCompleted {
				runLifecycleHook(context.Background(), types.Task{}, wf.Lifecycle.GetOnSuccessHook(), os.Stdout, logger)
			}
			return nil
		case event := <-events:
			switch x := event.(type) {
			// if we get the poison pill, we should see if any job tasks are failed, if so we must exist
			// if all jobs are either succeeded or skipped, we can exit
			case struct{}:
				remainingTasks := map[string]bool{}
				pendingTasks := map[string]bool{}
				for _, x := range taskNames {
					remainingTasks[x] = true
					pendingTasks[x] = true
				}

				for _, node := range subgraph.Nodes {
					// Check if task should cause immediate exit
					if node.Phase == "failed" && node.Task.GetRestartPolicy() == "Never" {
						logger.Printf("🚫 exiting because task %q failed and should not be restarted", node.Name)
						cancel()
						continue
					}

					// Check if task is complete and should be removed from tracking
					isComplete := false
					switch node.Phase {
					case "succeeded", "skipped":
						isComplete = true
					case "running", "stalled":
						// Services are complete when running (ready to serve)
						// Jobs are only complete when succeeded
						isComplete = node.Task.GetType() == types.TaskTypeService
					}

					if isComplete {
						delete(remainingTasks, node.Name)
						if node.Task.GetRestartPolicy() != "Always" {
							delete(pendingTasks, node.Name)
						}
					}
				}

				if len(pendingTasks) == 0 {
					logger.Println("✅ exiting because all requested tasks completed and none should be restarted")
					graphCompleted = true
					cancel()
				} else if len(remainingTasks) == 0 {
					if !allRunning {
						allRunning = true
						logger.Println("🔵 all requested tasks are running:")
						// print a list of running tasks, and their ports
						for _, node := range subgraph.Nodes {
							if (node.Phase == "running" || node.Phase == "stalled") && node.Task.Ports != nil {
								for _, port := range node.Task.Ports {
									logger.Printf(" - %s: http://localhost:%d\n", node.Name, port.HostPort)
								}
							}
						}
					}
				} else {
					allRunning = false
				}

			// if the event is a string, it is the name of the task to run
			case string:
				taskName := x

				// we will only execute this task, if its parents are "succeeded" or "skipped" or ("running" and the task is a service)
				blocked := false
				for _, parentName := range subgraph.Parents[taskName] {
					parent := subgraph.Nodes[parentName]
					if parent.blocked() {
						logger.Printf("task %q is blocked by %q (%s): %s\n", taskName, parentName, parent.Phase, parent.Message)
						blocked = true
					}
				}

				if blocked {
					continue
				}

				// we might already be pending, waiting, starting or running this task, so we don't want to start it again
				node := subgraph.Nodes[taskName]

				node.cancel()
				allRunning = false

				// each task is executed in a separate goroutine
				wg.Add(1)

				go func(node *TaskNode) {

					// lock the task, so we do not run two instances of it at the same time
					node.mu.Lock()

					ctx, cancel := context.WithCancel(ctx)
					defer cancel()

					node.cancel = cancel

					// send a poison pill to indicate that we've finish and the main loop must check to see if we need to exit
					defer func() { events <- poisonPill }()
					defer wg.Done()
					defer node.mu.Unlock()

					t := node.Task

					var out io.Writer = &logWriter{
						logger: logger,
						prefixSuffixProvider: func() (string, string) {
							return fmt.Sprintf("%s[%s] (%s)  ", color(node.Name), node.Name, node.Phase), "\033[0m"
						},
					}

					logger := log.New(out, "", 0)

					setNodeStatus := func(node *TaskNode, phase string, message string) {
						node.Phase = phase
						node.Message = message
						stallTimers[node.Name].Reset(node.Task.GetStalledTimeout())
						logger.Println(node.Message)
						statusEvents <- node
						events <- poisonPill
					}

					setNodeStatus(node, "waiting", "")

					queueChildren := func() {
						for _, child := range subgraph.Children[node.Name] {
							// only queue tasks in the subgraph
							if _, ok := subgraph.Nodes[child]; ok {
								logger.Printf("queuing %q\n", child)
								events <- child
							}
						}
					}

					// if the task can be skipped, lets exit early
					if t.Skip() || slices.Contains(tasksToSkip, node.Name) {
						setNodeStatus(node, "skipped", "")
						queueChildren()
						return
					}

					// if the task needs a mutex, lets wait for it
					if t.Mutex != "" {
						mu := util.GetMutex(t.Mutex)
						setNodeStatus(node, "waiting", "waiting for mutex")
						mu.Lock()
						setNodeStatus(node, "waiting", "acquired mutex")
						defer mu.Unlock()
					}

					// if the task needs a semaphore, lets wait for it
					if t.Semaphore != "" {
						sema := semaphores.Get(t.Semaphore)
						setNodeStatus(node, "waiting", "waiting for semaphore")
						if err := sema.Acquire(ctx, 1); err != nil {
							setNodeStatus(node, "failed", fmt.Sprintf("failed to acquire semaphore: %v", err))
							return
						}
						setNodeStatus(node, "waiting", "acquired semaphore")
						defer sema.Release(1)
					}

					p := proc.New(taskName, t, logger, types.Spec(*wf))

					if probe := t.GetLivenessProbe(); probe != nil {
						liveFunc := func(live bool, err error) {
							if !live {
								setNodeStatus(node, "failed", fmt.Sprintf("liveness probe failed: %v", err))
								cancel()
							}
						}
						go probeLoop(ctx, *probe, liveFunc)
					}
					if probe := t.GetReadinessProbe(); probe != nil {
						readyFunc := func(ready bool, err error) {
							if ready {
								setNodeStatus(node, "running", "readiness probe succeeded")
								queueChildren()
							} else {
								setNodeStatus(node, "failed", fmt.Sprintf("readiness probe failed: %v", err))
								cancel()
							}
						}
						go probeLoop(ctx, *probe, readyFunc)
					}

					if t.GetType() == types.TaskTypeService {
						if t.Ports != nil {
							setNodeStatus(node, "starting", "service starting")
						} else {
							setNodeStatus(node, "running", "no ports to expose")
							queueChildren()
						}
					} else {
						// non a service, must be a job
						setNodeStatus(node, "running", "job running")
					}

					restart := func() {
						select {
						case <-ctx.Done():
						case <-time.After(3 * time.Second):
							logger.Println("restarting")
							cancel()
							events <- node.Name
						}
					}

					file, err := os.Create(node.logFile)
					if err != nil {
						setNodeStatus(node, "failed", fmt.Sprintf("failed to create log file: %v", err))
						return
					}
					defer file.Close()

					// if the task has a log file, we will write to that file, we sync after each write
					// so when we tail the log file, we see the output immediately
					buf := funcWriter(func(p []byte) (int, error) {
						stallTimers[node.Name].Reset(node.Task.GetStalledTimeout())
						if node.Phase == "stalled" {
							if strings.HasSuffix(node.Message, "starting") {
								setNodeStatus(node, "starting", "output received")
							} else {
								setNodeStatus(node, "running", "output received")
							}
						}
						n, err := file.Write(p)
						if err != nil {
							return n, err
						}
						if err := file.Sync(); err != nil {
							return n, err
						}
						return n, nil
					})

					if t.Log != "" {
						out = buf
					} else {
						out = io.MultiWriter(out, buf)
					}

					go func() {
						ticker := time.NewTicker(20 * time.Second)
						defer ticker.Stop()

						for {
							select {
							case <-ctx.Done():
								return
							case <-ticker.C:
								if node.Phase != "running" && node.Phase != "stalled" {
									continue
								}
								metrics, err := p.GetMetrics(ctx)
								if err != nil {
									logger.Printf("failed to get metrics: %v", err)
									continue
								}
								node.Metrics = metrics
								statusEvents <- node
							}
						}
					}()

					err = p.Run(ctx, out, out)
					// if the task was cancelled, we don't want to restart it, this is normal exit
					if errors.Is(ctx.Err(), context.Canceled) {
						setNodeStatus(node, "cancelled", "")
						return
					}

					if err != nil {
						setNodeStatus(node, "failed", fmt.Sprint(err))
						runLifecycleHook(ctx, t, t.GetOnFailureHook(), out, logger)
						if t.GetRestartPolicy() != "Never" {
							restart()
						}
						return
					}

					setNodeStatus(node, "succeeded", "")
					runLifecycleHook(ctx, t, t.GetOnSuccessHook(), out, logger)
					if t.GetRestartPolicy() == "Always" {
						restart()
					}
					queueChildren()

				}(node)
			default:
				panic(fmt.Sprintf("unexpected event: %v", event))
			}
		}
	}
}

// runLifecycleHook runs the given lifecycle hook command, logging any errors.
// It is a best-effort operation: if the hook command fails, the error is logged
// but does not affect the task's outcome.
func runLifecycleHook(ctx context.Context, t types.Task, hook *types.LifecycleHook, out io.Writer, logger *log.Logger) {
	if hook == nil {
		return
	}
	cmd := hook.GetCommand()
	if len(cmd) == 0 {
		return
	}
	environ, err := types.Environ(types.Spec{}, t)
	if err != nil {
		logger.Printf("lifecycle hook: failed to get environment: %v", err)
		return
	}
	c := exec.CommandContext(ctx, cmd[0], cmd[1:]...)
	c.Dir = t.WorkingDir
	c.Stdout = out
	c.Stderr = out
	c.Env = append(environ, os.Environ()...)
	if err := c.Run(); err != nil {
		logger.Printf("lifecycle hook failed: %v", err)
	}
}