services/crunchstat/crunchstat.go

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