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 func main() {
  20         reporter := crunchstat.Reporter{
  21                 Logger: log.New(os.Stderr, "crunchstat: ", 0),
  22         }
  23
  24         flag.StringVar(&reporter.CgroupRoot, "cgroup-root", "", "Root of cgroup tree")
  25         flag.StringVar(&reporter.CgroupParent, "cgroup-parent", "", "Name of container parent under cgroup")
  26         flag.StringVar(&reporter.CIDFile, "cgroup-cid", "", "Path to container id file")
  27         pollMsec := flag.Int64("poll", 1000, "Reporting interval, in milliseconds")
  28
  29         flag.Parse()
  30
  31         if reporter.CgroupRoot == "" {
  32                 reporter.Logger.Fatal("error: must provide -cgroup-root")
  33         }
  34         reporter.PollPeriod = time.Duration(*pollMsec) * time.Millisecond
  35
  36         reporter.Start()
  37         err := runCommand(flag.Args(), reporter.Logger)
  38         reporter.Stop()
  39
  40         if err, ok := err.(*exec.ExitError); ok {
  41                 // The program has exited with an exit code != 0
  42
  43                 // This works on both Unix and Windows. Although
  44                 // package syscall is generally platform dependent,
  45                 // WaitStatus is defined for both Unix and Windows and
  46                 // in both cases has an ExitStatus() method with the
  47                 // same signature.
  48                 if status, ok := err.Sys().(syscall.WaitStatus); ok {
  49                         os.Exit(status.ExitStatus())
  50                 } else {
  51                         reporter.Logger.Fatalln("ExitError without WaitStatus:", err)
  52                 }
  53         } else if err != nil {
  54                 reporter.Logger.Fatalln("error in cmd.Wait:", err)
  55         }
  56 }
  57
  58 func runCommand(argv []string, logger *log.Logger) error {
  59         cmd := exec.Command(argv[0], argv[1:]...)
  60
  61         logger.Println("Running", argv)
  62
  63         // Child process will use our stdin and stdout pipes
  64         // (we close our copies below)
  65         cmd.Stdin = os.Stdin
  66         cmd.Stdout = os.Stdout
  67
  68         // Forward SIGINT and SIGTERM to child process
  69         sigChan := make(chan os.Signal, 1)
  70         go func(sig <-chan os.Signal) {
  71                 catch := <-sig
  72                 if cmd.Process != nil {
  73                         cmd.Process.Signal(catch)
  74                 }
  75                 logger.Println("notice: caught signal:", catch)
  76         }(sigChan)
  77         signal.Notify(sigChan, syscall.SIGTERM)
  78         signal.Notify(sigChan, syscall.SIGINT)
  79
  80         // Funnel stderr through our channel
  81         stderr_pipe, err := cmd.StderrPipe()
  82         if err != nil {
  83                 logger.Fatalln("error in StderrPipe:", err)
  84         }
  85
  86         // Run subprocess
  87         if err := cmd.Start(); err != nil {
  88                 logger.Fatalln("error in cmd.Start:", err)
  89         }
  90
  91         // Close stdin/stdout in this (parent) process
  92         os.Stdin.Close()
  93         os.Stdout.Close()
  94
  95         copyPipeToChildLog(stderr_pipe, log.New(os.Stderr, "", 0))
  96
  97         return cmd.Wait()
  98 }
  99
 100 func copyPipeToChildLog(in io.ReadCloser, logger *log.Logger) {
 101         reader := bufio.NewReaderSize(in, MaxLogLine)
 102         var prefix string
 103         for {
 104                 line, isPrefix, err := reader.ReadLine()
 105                 if err == io.EOF {
 106                         break
 107                 } else if err != nil {
 108                         logger.Fatal("error reading child stderr:", err)
 109                 }
 110                 var suffix string
 111                 if isPrefix {
 112                         suffix = "[...]"
 113                 }
 114                 logger.Print(prefix, string(line), suffix)
 115                 // Set up prefix for following line
 116                 if isPrefix {
 117                         prefix = "[...]"
 118                 } else {
 119                         prefix = ""
 120                 }
 121         }
 122         in.Close()
 123 }