/
scout-sniper.go
163 lines (139 loc) · 4.21 KB
/
scout-sniper.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
package main
import (
"bufio"
"exec"
"flag"
"http"
"os"
"syscall"
"time"
// "runtime"
)
const Version = "0.1"
const Seconds = 1e9
const KiloBytes = 1024
const Megabytes = KiloBytes * 1024
// const UserAgent = "scout-sniper/" + Version + " golang/" + runtime.GOOS + "-" + runtime.GOARCH
var extraInterval = flag.Int64("extra-interval", 30, "The time interval between extra checks")
var killCode = flag.Int("kill-signal", 9, "Kill signal to use when killing a process")
var httpTimeoutUrl = flag.String("http-timeout-url", "", "The url to check for HTTP timeouts")
var httpTimeoutTime = flag.Int64("http-timeout-time", 5, "The timeout for the HTTP timeout check")
var httpStatusUrl = flag.String("http-status-url", "", "The url to check for HTTP status")
var httpStatusCode = flag.Int("http-status-code", 200, "The status code for the HTTP status check")
var maxMemory = flag.Float64("max-mem", 0, "The amount of memory in megabytes that the process is allowed")
type BC chan bool
type ABC []BC
type ProcInfo map[string] string
func getProcessInformation(pid int) ProcInfo {
return map[string] string {
"foo": "bar",
}
}
func (a ABC) closeAll() {
for _, c := range(a) {
c <- true
close(c)
}
}
func timeout(timeout int64, f func(chan int)) (ok bool) {
ctr := make(chan int, 1)
cto := make(chan int, 1)
go f(ctr)
go func() {
time.Sleep(timeout * Seconds)
cto <- 1
}()
select {
case <- cto:
ok = false
case <- ctr:
ok = true
}
return ok
}
func check(interval int64, quit BC, f func()) {
for {
time.Sleep(interval * Seconds)
select {
case <- quit:
return
default:
f()
}
}
}
func setupHttpTimeoutCheck(pid int) BC {
ch := make(BC, 1)
go check(*extraInterval, ch, func() {
failed := !timeout(*httpTimeoutTime, func(ctr chan int) {
resp, _, err := http.Get(*httpTimeoutUrl)
if err == nil && resp.Body != nil {
defer resp.Body.Close()
reader := bufio.NewReader(resp.Body)
reader.ReadString('\r')
ctr <- 1
}
})
if failed {
println("HTTP timeout check failed after", *httpTimeoutTime, "seconds. Killing process", pid, "with signal", *killCode)
syscall.Kill(pid, *killCode)
}
})
return ch
}
func setupHttpStatusCheck(pid int) BC {
ch := make(BC, 1)
go check(*extraInterval, ch, func() {
failed := !timeout(*httpTimeoutTime, func(ctr chan int) {
resp, _, err := http.Get(*httpStatusUrl)
if err == nil && *httpStatusCode == resp.StatusCode {
ctr <- 1
}
})
if failed {
println("HTTP status check failed after", *httpTimeoutTime, "seconds. Killing process", pid, "with signal", *killCode)
syscall.Kill(pid, *killCode)
}
})
return ch
}
func setupMaxMemoryCheck(pid int) BC {
ch := make(BC, 1)
poller := GetPoller()
alreadyExceded := false
go check(*extraInterval, ch, func() {
kb := poller.GetMemory(pid)
if kb > *maxMemory {
if alreadyExceded {
println("Memory check failed for 2 intervals. Killing process", pid, "with signal", *killCode)
syscall.Kill(pid, *killCode)
}
alreadyExceded = true
} else {
alreadyExceded = false
}
})
return ch
}
func main() {
flag.Parse()
cwd, _ := os.Getwd()
binary, _ := exec.LookPath(flag.Args()[0])
for {
cmd, _ := exec.Run(binary, flag.Args(), nil, cwd, exec.PassThrough, exec.PassThrough, exec.PassThrough)
pid := cmd.Process.Pid
extras := make(ABC, 0)
if *httpTimeoutUrl != "" {
extras = append(extras, setupHttpTimeoutCheck(pid))
}
if *httpStatusUrl != "" {
extras = append(extras, setupHttpStatusCheck(pid))
}
if *maxMemory > 0 {
extras = append(extras, setupMaxMemoryCheck(pid))
}
cmd.Wait(os.WSTOPPED)
println("Process died, restarting.")
extras.closeAll()
}
}