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