142 lines
3.2 KiB
Go
142 lines
3.2 KiB
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
|
|
}
|