oko_public/oko-run/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"gopkg.in/yaml.v3"
)

type Config struct {
	Nodes []Node `yaml:"nodes"`
}

type Node struct {
	Name   string                 `yaml:"name"`
	Binary string                 `yaml:"binary"`
	Flags  map[string]interface{} `yaml:"flags"`
}

func main() {
	configPath := flag.String("config", "config.yaml", "path to pipeline configuration")
	flag.Parse()

	absConfig, err := filepath.Abs(*configPath)
	if err != nil {
		log.Fatalf("abs config path: %v", err)
	}
	rootDir := filepath.Dir(absConfig)

	data, err := os.ReadFile(absConfig)
	if err != nil {
		log.Fatalf("read config: %v", err)
	}

	var cfg Config
	if err := yaml.Unmarshal(data, &cfg); err != nil {
		log.Fatalf("parse config: %v", err)
	}

	if len(cfg.Nodes) == 0 {
		log.Fatal("no nodes defined in config")
	}

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	for _, node := range cfg.Nodes {
		go runNode(ctx, rootDir, node)
	}

	<-ctx.Done()
	log.Println("shutting down pipeline")
}

func runNode(ctx context.Context, rootDir string, node Node) {
	binaryDir := filepath.Join(rootDir, node.Binary)
	binaryPath := filepath.Join(binaryDir, node.Binary)

	if _, err := os.Stat(binaryPath); err != nil {
		log.Printf("building %s (%s)", node.Name, node.Binary)
		build := exec.CommandContext(ctx, "go", "build", "-o", node.Binary, ".")
		build.Dir = binaryDir
		build.Stdout = os.Stdout
		build.Stderr = os.Stderr
		if err := build.Run(); err != nil {
			log.Printf("%s build failed: %v", node.Name, err)
			return
		}
	}

	args := buildArgs(node.Flags)

	// Restart a crashed node until the pipeline is told to shut down. Each
	// restart is preceded by a short backoff so a crash loop does not spin.
	for ctx.Err() == nil {
		cmd := exec.Command(binaryPath, args...)
		cmd.Dir = binaryDir
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr

		// Deliver SIGTERM (not SIGKILL) on shutdown so each node drains and
		// closes its V4L handles / socket cleanly.
		stop := make(chan struct{})
		go func() {
			select {
			case <-ctx.Done():
				cmd.Process.Signal(syscall.SIGTERM)
			case <-stop:
			}
		}()

		log.Printf("starting %s (%s)", node.Name, node.Binary)
		err := cmd.Run()
		close(stop)
		if ctx.Err() != nil {
			return
		}
		log.Printf("%s exited unexpectedly (%v); restarting in 2s", node.Name, err)
		select {
		case <-ctx.Done():
			return
		case <-time.After(2 * time.Second):
		}
	}
}

func buildArgs(flags map[string]interface{}) []string {
	var args []string
	for key, val := range flags {
		switch v := val.(type) {
		case string:
			args = append(args, fmt.Sprintf("--%s=%s", key, v))
		case int:
			args = append(args, fmt.Sprintf("--%s=%d", key, v))
		case float64:
			args = append(args, fmt.Sprintf("--%s=%v", key, v))
		case bool:
			if v {
				args = append(args, fmt.Sprintf("--%s", key))
			}
		case []interface{}:
			var parts []string
			for _, item := range v {
				parts = append(parts, fmt.Sprint(item))
			}
			args = append(args, fmt.Sprintf("--%s=%s", key, strings.Join(parts, ",")))
		default:
			args = append(args, fmt.Sprintf("--%s=%v", key, v))
		}
	}
	return args
}