-
Notifications
You must be signed in to change notification settings - Fork 1
/
timeout.go
102 lines (96 loc) · 2.91 KB
/
timeout.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
// Command timeout runs a command, stopping it after a duration.
//
// Usage: timeout [-s signal] [-g grace] <duration> <command> [args...]
//
// Timeout runs <command> with [args...] signalling it with <signal> (default
// SIGTERM) after the given <duration>. If the process has not exited by the
// time the <grace> period passes (default 3s), the process is killed.
//
// The exit code of timeout is <command>'s exit code if it exited on its own.
// If the command could not be run or the timeout expires, the exit code of
// timeout is 1.
//
// <duration> and <grace> are parsed as time.Duration strings (see
// https://golang.org/pkg/time/#ParseDuration).
//
// <command> is executed directly with [args...] as provided. If <command> does
// not contain any path separators, the search path is used to locate it. No
// shell is used to run <command>.
package main
import (
"context"
"errors"
"flag"
"fmt"
"os"
"os/exec"
"syscall"
"time"
)
var (
grace = flag.Duration("g", 3*time.Second, "Grace period before sending SIGKILL")
sig = flag.Int("s", int(syscall.SIGTERM), "Signal to send on timeout")
)
func main() {
flag.Parse()
code, err := timeout(flag.Args())
if err != nil {
fmt.Fprintln(os.Stderr, err)
}
os.Exit(code)
}
func timeout(args []string) (int, error) {
if len(args) < 2 {
return 1, errors.New("usage: timeout [-s signal] [-g grace] <duration> <command> [args...]")
}
d, err := time.ParseDuration(args[0])
if err != nil {
return 1, fmt.Errorf("invalid duration: %v", err)
}
return Timeout(d, args[1:]...)
}
// Timeout runs a command, signalling it after a given duration. If the process
// has not exited by the time the grace period passes, the process is killed.
// The command's exit code and an error (if any) are returned. If the command
// could not be run, or the timeout expires, the exit code returned is 1.
func Timeout(d time.Duration, argv ...string) (int, error) {
cmd, err := exec.LookPath(argv[0])
if err != nil {
return 1, err
}
// Prevent stdin, stdout and stderr from being closed by StartProcess
attr := &os.ProcAttr{Files: []*os.File{os.Stdin, os.Stdout, os.Stderr}}
p, err := os.StartProcess(cmd, argv, attr)
if err != nil {
return 1, err
}
ctx, cancel := context.WithCancel(context.Background())
isTimeout := make(chan bool)
go func() {
timedOut := callAfter(ctx, d, func() { _ = p.Signal(syscall.Signal(*sig)) })
if timedOut {
callAfter(ctx, *grace, func() { _ = p.Kill() })
}
isTimeout <- timedOut
}()
ps, err := p.Wait()
cancel()
if <-isTimeout {
return 1, fmt.Errorf("timeout (%s)", d)
}
if err != nil {
return 1, err
}
return ps.ExitCode(), nil
}
// callAfter calls function f after duration d and returns true. If the context
// is cancelled before duration d, return false without invoking f.
func callAfter(ctx context.Context, d time.Duration, f func()) bool {
select {
case <-time.After(d):
f()
return true
case <-ctx.Done():
return false
}
}