forked from vitessio/vitess
-
Notifications
You must be signed in to change notification settings - Fork 0
/
exit.go
79 lines (65 loc) · 1.92 KB
/
exit.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
// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
/*
Package exit provides an alternative to os.Exit(int) that executes deferred
functions before exiting.
Defer a call to exit.Recover() or exit.RecoverAll() at the beginning of main().
Use exit.Return(int) to initiate an exit.
func main() {
defer exit.Recover()
...
f()
}
func f() {
exit.Return(123)
}
This package should only be used from the main goroutine.
*/
package exit
import (
"os"
log "github.com/golang/glog"
"github.com/youtube/vitess/go/tb"
)
type exitCode int
var (
exitFunc = os.Exit // can be faked out for testing
)
// Recover should be deferred as the first line of main(). It recovers the
// panic initiated by Return and converts it to a call to os.Exit. Any
// functions deferred after Recover in the main goroutine should be executed
// prior to exiting. Recover will re-panic anything other than the panic it
// expects from Return.
func Recover() {
doRecover(recover(), false)
}
// RecoverAll can be deferred instead of Recover as the first line of main().
// Instead of re-panicking, RecoverAll will absorb any panic and convert it to
// an error log entry with a stack trace, followed by a call to os.Exit(255).
func RecoverAll() {
doRecover(recover(), true)
}
func doRecover(err interface{}, recoverAll bool) {
if err == nil {
return
}
switch code := err.(type) {
case exitCode:
exitFunc(int(code))
default:
if recoverAll {
log.Errorf("panic: %v", tb.Errorf("%v", err))
exitFunc(255)
} else {
panic(err)
}
}
}
// Return initiates a panic that sends the return code to the deferred Recover,
// executing other deferred functions along the way. When the panic reaches
// Recover, the return code will be passed to os.Exit. This should only be
// called from the main goroutine.
func Return(code int) {
panic(exitCode(code))
}