services/crunchstat/crunchstat.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "flag"
   6         "io"
   7         "log"
   8         "os"
   9         "os/exec"
  10         "os/signal"
  11         "syscall"
  12         "time"
  13
  14         "git.curoverse.com/arvados.git/lib/crunchstat"
  15 )
  16
  17 const MaxLogLine = 1 << 14 // Child stderr lines >16KiB will be split
  18
  19 var signalOnDeadPPID int
  20 var ppidCheckInterval = time.Second
  21
  22 func main() {
  23         reporter := crunchstat.Reporter{
  24                 Logger: log.New(os.Stderr, "crunchstat: ", 0),
  25         }
  26
  27         flag.StringVar(&reporter.CgroupRoot, "cgroup-root", "", "Root of cgroup tree")
  28         flag.StringVar(&reporter.CgroupParent, "cgroup-parent", "", "Name of container parent under cgroup")
  29         flag.StringVar(&reporter.CIDFile, "cgroup-cid", "", "Path to container id file")
  30         flag.IntVar(&signalOnDeadPPID, "signal-on-dead-ppid", 15, "Signal to send child if crunchstat's parent process disappears")
  31         flag.DurationVar(&ppidCheckInterval, "ppid-check-interval", ppidCheckInterval, "Time between checks for parent process disappearance")
  32         pollMsec := flag.Int64("poll", 1000, "Reporting interval, in milliseconds")
  33
  34         flag.Parse()
  35
  36         if reporter.CgroupRoot == "" {
  37                 reporter.Logger.Fatal("error: must provide -cgroup-root")
  38         }
  39         reporter.PollPeriod = time.Duration(*pollMsec) * time.Millisecond
  40
  41         reporter.Start()
  42         err := runCommand(flag.Args(), reporter.Logger)
  43         reporter.Stop()
  44
  45         if err, ok := err.(*exec.ExitError); ok {
  46                 // The program has exited with an exit code != 0
  47
  48                 // This works on both Unix and Windows. Although
  49                 // package syscall is generally platform dependent,
  50                 // WaitStatus is defined for both Unix and Windows and
  51                 // in both cases has an ExitStatus() method with the
  52                 // same signature.
  53                 if status, ok := err.Sys().(syscall.WaitStatus); ok {
  54                         os.Exit(status.ExitStatus())
  55                 } else {
  56                         reporter.Logger.Fatalln("ExitError without WaitStatus:", err)
  57                 }
  58         } else if err != nil {
  59                 reporter.Logger.Fatalln("error in cmd.Wait:", err)
  60         }
  61 }
  62
  63 func runCommand(argv []string, logger *log.Logger) error {
  64         cmd := exec.Command(argv[0], argv[1:]...)
  65
  66         logger.Println("Running", argv)
  67
  68         // Child process will use our stdin and stdout pipes
  69         // (we close our copies below)
  70         cmd.Stdin = os.Stdin
  71         cmd.Stdout = os.Stdout
  72
  73         // Forward SIGINT and SIGTERM to child process
  74         sigChan := make(chan os.Signal, 1)
  75         go func(sig <-chan os.Signal) {
  76                 catch := <-sig
  77                 if cmd.Process != nil {
  78                         cmd.Process.Signal(catch)
  79                 }
  80                 logger.Println("notice: caught signal:", catch)
  81         }(sigChan)
  82         signal.Notify(sigChan, syscall.SIGTERM)
  83         signal.Notify(sigChan, syscall.SIGINT)
  84
  85         // Kill our child proc if our parent process disappears
  86         if signalOnDeadPPID != 0 {
  87                 go sendSignalOnDeadPPID(signalOnDeadPPID, os.Getppid(), cmd, logger)
  88         }
  89
  90         // Funnel stderr through our channel
  91         stderr_pipe, err := cmd.StderrPipe()
  92         if err != nil {
  93                 logger.Fatalln("error in StderrPipe:", err)
  94         }
  95
  96         // Run subprocess
  97         if err := cmd.Start(); err != nil {
  98                 logger.Fatalln("error in cmd.Start:", err)
  99         }
 100
 101         // Close stdin/stdout in this (parent) process
 102         os.Stdin.Close()
 103         os.Stdout.Close()
 104
 105         copyPipeToChildLog(stderr_pipe, log.New(os.Stderr, "", 0))
 106
 107         return cmd.Wait()
 108 }
 109
 110 func sendSignalOnDeadPPID(signum, ppidOrig int, cmd *exec.Cmd, logger *log.Logger) {
 111         for _ = range time.NewTicker(ppidCheckInterval).C {
 112                 ppid := os.Getppid()
 113                 if ppid == ppidOrig {
 114                         continue
 115                 }
 116                 if cmd.Process == nil {
 117                         // Child process isn't running yet
 118                         continue
 119                 }
 120                 logger.Printf("notice: crunchstat ppid changed from %d to %d -- killing child pid %d with signal %d", ppidOrig, ppid, cmd.Process.Pid, signum)
 121                 err := cmd.Process.Signal(syscall.Signal(signum))
 122                 if err != nil {
 123                         logger.Printf("error: sending signal: %d", err)
 124                         continue
 125                 }
 126                 break
 127         }
 128 }
 129
 130 func copyPipeToChildLog(in io.ReadCloser, logger *log.Logger) {
 131         reader := bufio.NewReaderSize(in, MaxLogLine)
 132         var prefix string
 133         for {
 134                 line, isPrefix, err := reader.ReadLine()
 135                 if err == io.EOF {
 136                         break
 137                 } else if err != nil {
 138                         logger.Fatal("error reading child stderr:", err)
 139                 }
 140                 var suffix string
 141                 if isPrefix {
 142                         suffix = "[...]"
 143                 }
 144                 logger.Print(prefix, string(line), suffix)
 145                 // Set up prefix for following line
 146                 if isPrefix {
 147                         prefix = "[...]"
 148                 } else {
 149                         prefix = ""
 150                 }
 151         }
 152         in.Close()
 153 }