09a99e5df2a5b72eb38671df4e250aa9d109d983
[arvados.git] / services / crunch / crunchstat / src / arvados.org / crunchstat / crunchstat.go
1 package main
2
3 import (
4         "bufio"
5         "flag"
6         "fmt"
7         "io"
8         "io/ioutil"
9         "log"
10         "os"
11         "os/exec"
12         "os/signal"
13         "strings"
14         "syscall"
15         "time"
16 )
17
18 func ReadLineByLine(inp io.ReadCloser, out chan string, finish chan bool) {
19         s := bufio.NewScanner(inp)
20         for s.Scan() {
21                 out <- s.Text()
22         }
23         finish <- true
24 }
25
26 func OutputChannel(stdout chan string, stderr chan string) {
27         for {
28                 select {
29                 case s, ok := <-stdout:
30                         if ok {
31                                 fmt.Fprintln(os.Stdout, s)
32                         } else {
33                                 return
34                         }
35                 case s, ok := <-stderr:
36                         if ok {
37                                 fmt.Fprintln(os.Stderr, s)
38                         } else {
39                                 return
40                         }
41                 }
42         }
43 }
44
45 func FindStat(cgroup_path string, statgroup string, stat string) string {
46         path := fmt.Sprintf("%s/%s.%s", cgroup_path, statgroup, stat)
47         if _, err := os.Stat(path); err != nil {
48                 return path
49         }
50         path = fmt.Sprintf("%s/%s/%s.%s", cgroup_path, statgroup, statgroup, stat)
51         if _, err := os.Stat(path); err != nil {
52                 return path
53         }
54         return ""
55 }
56
57 func PollCgroupStats(cgroup_path string, stderr chan string, poll int64) {
58         //var last_usage int64 = 0
59         var last_user int64 = 0
60         var last_sys int64 = 0
61         var last_cpucount int64 = 0
62
63         type Disk struct {
64                 last_read  int64
65                 next_read  int64
66                 last_write int64
67                 next_write int64
68         }
69
70         disk := make(map[string]*Disk)
71
72         //cpuacct_usage := FindStat(cgroup_path, "cpuacct", "usage")
73         cpuacct_stat := FindStat(cgroup_path, "cpuacct", "stat")
74         blkio_io_service_bytes := FindStat(cgroup_path, "blkio", "io_service_bytes")
75         cpuset_cpus := FindStat(cgroup_path, "cpuset", "cpus")
76         memory_stat := FindStat(cgroup_path, "memory", "stat")
77
78         var elapsed int64 = poll
79
80         for {
81                 /*{
82                         c, _ := os.Open(cpuacct_usage)
83                         b, _ := ioutil.ReadAll(c)
84                         var next int64
85                         fmt.Sscanf(string(b), "%d", &next)
86                         if last_usage != 0 {
87                                 stderr <- fmt.Sprintf("crunchstat: cpuacct.usage %v", (next-last_usage)/10000000)
88                         }
89                         //fmt.Printf("usage %d %d %d %d%%\n", last_usage, next, next-last_usage, (next-last_usage)/10000000)
90                         last_usage = next
91                         c.Close()
92                 }*/
93                 var cpus int64 = 0
94                 if cpuset_cpus != "" {
95                         c, _ := os.Open(cpuset_cpus)
96                         b, _ := ioutil.ReadAll(c)
97                         sp := strings.Split(string(b), ",")
98                         for _, v := range sp {
99                                 var min, max int64
100                                 n, _ := fmt.Sscanf(v, "%d-%d", &min, &max)
101                                 if n == 2 {
102                                         cpus += (max - min) + 1
103                                 } else {
104                                         cpus += 1
105                                 }
106                         }
107
108                         if cpus != last_cpucount {
109                                 stderr <- fmt.Sprintf("crunchstat: cpuset.cpus %v", cpus)
110                         }
111                         last_cpucount = cpus
112
113                         c.Close()
114                 }
115                 if cpus == 0 {
116                         cpus = 1
117                 }
118                 if cpuacct_stat != "" {
119                         c, _ := os.Open(cpuacct_stat)
120                         b, _ := ioutil.ReadAll(c)
121                         var next_user int64
122                         var next_sys int64
123                         fmt.Sscanf(string(b), "user %d\nsystem %d", &next_user, &next_sys)
124                         c.Close()
125
126                         if last_user != 0 {
127                                 user_diff := next_user - last_user
128                                 sys_diff := next_sys - last_sys
129                                 // Assume we're reading stats based on 100
130                                 // jiffies per second.  Because the ellaspsed
131                                 // time is in milliseconds, we need to boost
132                                 // that to 1000 jiffies per second, then boost
133                                 // it by another 100x to get a percentage, then
134                                 // finally divide by the actual elapsed time
135                                 // and the number of cpus to get average load
136                                 // over the polling period.
137                                 user_pct := (user_diff * 10 * 100) / (elapsed * cpus)
138                                 sys_pct := (sys_diff * 10 * 100) / (elapsed * cpus)
139
140                                 stderr <- fmt.Sprintf("crunchstat: cpuacct.stat user %v", user_pct)
141                                 stderr <- fmt.Sprintf("crunchstat: cpuacct.stat sys %v", sys_pct)
142                         }
143
144                         /*fmt.Printf("user %d %d %d%%\n", last_user, next_user, next_user-last_user)
145                         fmt.Printf("sys %d %d %d%%\n", last_sys, next_sys, next_sys-last_sys)
146                         fmt.Printf("sum %d%%\n", (next_user-last_user)+(next_sys-last_sys))*/
147                         last_user = next_user
148                         last_sys = next_sys
149                 }
150                 if blkio_io_service_bytes != "" {
151                         c, _ := os.Open(blkio_io_service_bytes)
152                         b := bufio.NewScanner(c)
153                         var device, op string
154                         var next int64
155                         for b.Scan() {
156                                 if _, err := fmt.Sscanf(string(b.Text()), "%s %s %d", &device, &op, &next); err == nil {
157                                         if disk[device] == nil {
158                                                 disk[device] = new(Disk)
159                                         }
160                                         if op == "Read" {
161                                                 disk[device].last_read = disk[device].next_read
162                                                 disk[device].next_read = next
163                                                 if disk[device].last_read > 0 {
164                                                         stderr <- fmt.Sprintf("crunchstat: blkio.io_service_bytes %s read %v", device, disk[device].next_read-disk[device].last_read)
165                                                 }
166                                         }
167                                         if op == "Write" {
168                                                 disk[device].last_write = disk[device].next_write
169                                                 disk[device].next_write = next
170                                                 if disk[device].last_write > 0 {
171                                                         stderr <- fmt.Sprintf("crunchstat: blkio.io_service_bytes %s write %v", device, disk[device].next_write-disk[device].last_write)
172                                                 }
173                                         }
174                                 }
175                         }
176                         c.Close()
177                 }
178
179                 if memory_stat != "" {
180                         c, _ := os.Open(memory_stat)
181                         b := bufio.NewScanner(c)
182                         var stat string
183                         var val int64
184                         for b.Scan() {
185                                 if _, err := fmt.Sscanf(string(b.Text()), "%s %d", &stat, &val); err == nil {
186                                         if stat == "rss" {
187                                                 stderr <- fmt.Sprintf("crunchstat: memory.stat rss %v", val)
188                                         }
189                                 }
190                         }
191                         c.Close()
192                 }
193
194                 bedtime := time.Now()
195                 time.Sleep(time.Duration(poll) * time.Millisecond)
196                 morning := time.Now()
197                 elapsed = morning.Sub(bedtime).Nanoseconds() / int64(time.Millisecond)
198         }
199 }
200
201 func main() {
202
203         var (
204                 cgroup_path    string
205                 cgroup_parent  string
206                 cgroup_cidfile string
207                 wait           int64
208                 poll           int64
209         )
210
211         flag.StringVar(&cgroup_path, "cgroup-path", "", "Direct path to cgroup")
212         flag.StringVar(&cgroup_parent, "cgroup-parent", "", "Path to parent cgroup")
213         flag.StringVar(&cgroup_cidfile, "cgroup-cid", "", "Path to container id file")
214         flag.Int64Var(&wait, "wait", 5, "Maximum time (in seconds) to wait for cid file to show up")
215         flag.Int64Var(&poll, "poll", 1000, "Polling frequency, in milliseconds")
216
217         flag.Parse()
218
219         logger := log.New(os.Stderr, "crunchstat: ", 0)
220
221         if cgroup_path == "" && cgroup_cidfile == "" {
222                 logger.Fatal("Must provide either -cgroup-path or -cgroup-cid")
223         }
224
225         // Make output channel
226         stdout_chan := make(chan string)
227         stderr_chan := make(chan string)
228         finish_chan := make(chan bool)
229         defer close(stdout_chan)
230         defer close(stderr_chan)
231         defer close(finish_chan)
232
233         go OutputChannel(stdout_chan, stderr_chan)
234
235         var cmd *exec.Cmd
236
237         if len(flag.Args()) > 0 {
238                 // Set up subprocess
239                 cmd = exec.Command(flag.Args()[0], flag.Args()[1:]...)
240
241                 logger.Print("Running ", flag.Args())
242
243                 // Forward SIGINT and SIGTERM to inner process
244                 term := make(chan os.Signal, 1)
245                 go func(sig <-chan os.Signal) {
246                         catch := <-sig
247                         if cmd.Process != nil {
248                                 cmd.Process.Signal(catch)
249                         }
250                         logger.Print("caught signal:", catch)
251                 }(term)
252                 signal.Notify(term, syscall.SIGTERM)
253                 signal.Notify(term, syscall.SIGINT)
254
255                 // Funnel stdout and stderr from subprocess to output channels
256                 stdout_pipe, err := cmd.StdoutPipe()
257                 if err != nil {
258                         logger.Fatal(err)
259                 }
260                 go ReadLineByLine(stdout_pipe, stdout_chan, finish_chan)
261
262                 stderr_pipe, err := cmd.StderrPipe()
263                 if err != nil {
264                         logger.Fatal(err)
265                 }
266                 go ReadLineByLine(stderr_pipe, stderr_chan, finish_chan)
267
268                 // Run subprocess
269                 if err := cmd.Start(); err != nil {
270                         logger.Fatal(err)
271                 }
272         }
273
274         // Read the cid file
275         if cgroup_cidfile != "" {
276                 // wait up to 'wait' seconds for the cid file to appear
277                 var i time.Duration
278                 for i = 0; i < time.Duration(wait)*time.Second; i += (100 * time.Millisecond) {
279                         f, err := os.Open(cgroup_cidfile)
280                         if err == nil {
281                                 cid, err2 := ioutil.ReadAll(f)
282                                 if err2 == nil && len(cid) > 0 {
283                                         cgroup_path = string(cid)
284                                         f.Close()
285                                         break
286                                 }
287                         }
288                         time.Sleep(100 * time.Millisecond)
289                 }
290                 if cgroup_path == "" {
291                         logger.Printf("Could not read cid file %s", cgroup_cidfile)
292                 }
293         }
294
295         // add the parent prefix
296         if cgroup_parent != "" {
297                 cgroup_path = fmt.Sprintf("%s/%s", cgroup_parent, cgroup_path)
298         }
299
300         logger.Print("Using cgroup ", cgroup_path)
301
302         go PollCgroupStats(cgroup_path, stderr_chan, poll)
303
304         // Wait for each of stdout and stderr to drain
305         <-finish_chan
306         <-finish_chan
307
308         if err := cmd.Wait(); err != nil {
309                 if exiterr, ok := err.(*exec.ExitError); ok {
310                         // The program has exited with an exit code != 0
311
312                         // This works on both Unix and Windows. Although package
313                         // syscall is generally platform dependent, WaitStatus is
314                         // defined for both Unix and Windows and in both cases has
315                         // an ExitStatus() method with the same signature.
316                         if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
317                                 os.Exit(status.ExitStatus())
318                         }
319                 } else {
320                         logger.Fatalf("cmd.Wait: %v", err)
321                 }
322         }
323 }