forked from flashmob/go-guerrilla
-
Notifications
You must be signed in to change notification settings - Fork 0
/
hook.go
191 lines (170 loc) · 4.24 KB
/
hook.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
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
package log
import (
"bufio"
log "github.com/sirupsen/logrus"
"io"
"io/ioutil"
"os"
"strings"
"sync"
)
// custom logrus hook
// hookMu ensures all io operations are synced. Always on exported functions
var hookMu sync.Mutex
// LoggerHook extends the log.Hook interface by adding Reopen() and Rename()
type LoggerHook interface {
log.Hook
Reopen() error
}
type LogrusHook struct {
w io.Writer
// file descriptor, can be re-opened
fd *os.File
// filename to the file descriptor
fname string
// txtFormatter that doesn't use colors
plainTxtFormatter *log.TextFormatter
mu sync.Mutex
}
// newLogrusHook creates a new hook. dest can be a file name or one of the following strings:
// "stderr" - log to stderr, lines will be written to os.Stdout
// "stdout" - log to stdout, lines will be written to os.Stdout
// "off" - no log, lines will be written to ioutil.Discard
func NewLogrusHook(dest string) (LoggerHook, error) {
hookMu.Lock()
defer hookMu.Unlock()
hook := LogrusHook{fname: dest}
err := hook.setup(dest)
return &hook, err
}
type OutputOption int
const (
OutputStderr OutputOption = 1 + iota
OutputStdout
OutputOff
OutputNull
OutputFile
)
var outputOptions = [...]string{
"stderr",
"stdout",
"off",
"",
"file",
}
func (o OutputOption) String() string {
return outputOptions[o-1]
}
func parseOutputOption(str string) OutputOption {
switch str {
case "stderr":
return OutputStderr
case "stdout":
return OutputStdout
case "off":
return OutputOff
case "":
return OutputNull
}
return OutputFile
}
// Setup sets the hook's writer w and file descriptor fd
// assumes the hook.fd is closed and nil
func (hook *LogrusHook) setup(dest string) error {
out := parseOutputOption(dest)
if out == OutputNull || out == OutputStderr {
hook.w = os.Stderr
} else if out == OutputStdout {
hook.w = os.Stdout
} else if out == OutputOff {
hook.w = ioutil.Discard
} else {
if _, err := os.Stat(dest); err == nil {
// file exists open the file for appending
if err := hook.openAppend(dest); err != nil {
return err
}
} else {
// create the file
if err := hook.openCreate(dest); err != nil {
return err
}
}
}
// disable colors when writing to file
if hook.fd != nil {
hook.plainTxtFormatter = &log.TextFormatter{DisableColors: true}
}
return nil
}
// openAppend opens the dest file for appending. Default to os.Stderr if it can't open dest
func (hook *LogrusHook) openAppend(dest string) (err error) {
fd, err := os.OpenFile(dest, os.O_APPEND|os.O_WRONLY, 0644)
if err != nil {
log.WithError(err).Error("Could not open log file for appending")
hook.w = os.Stderr
hook.fd = nil
return
}
hook.w = bufio.NewWriter(fd)
hook.fd = fd
return
}
// openCreate creates a new dest file for appending. Default to os.Stderr if it can't open dest
func (hook *LogrusHook) openCreate(dest string) (err error) {
fd, err := os.OpenFile(dest, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
if err != nil {
log.WithError(err).Error("Could not create log file")
hook.w = os.Stderr
hook.fd = nil
return
}
hook.w = bufio.NewWriter(fd)
hook.fd = fd
return
}
// Fire implements the logrus Hook interface. It disables color text formatting if writing to a file
func (hook *LogrusHook) Fire(entry *log.Entry) error {
hookMu.Lock()
defer hookMu.Unlock()
line, err := entry.String()
if err == nil {
r := strings.NewReader(line)
if _, err = io.Copy(hook.w, r); err != nil {
return err
}
if wb, ok := hook.w.(*bufio.Writer); ok {
if err := wb.Flush(); err != nil {
return err
}
if hook.fd != nil {
err = hook.fd.Sync()
}
}
return err
}
return err
}
// Levels implements the logrus Hook interface
func (hook *LogrusHook) Levels() []log.Level {
return log.AllLevels
}
// Reopen closes and re-open log file descriptor, which is a special feature of this hook
func (hook *LogrusHook) Reopen() error {
hookMu.Lock()
defer hookMu.Unlock()
var err error
if hook.fd != nil {
if err = hook.fd.Close(); err != nil {
return err
}
// The file could have been re-named by an external program such as logrotate(8)
_, err := os.Stat(hook.fname)
if err != nil {
// The file doesn't exist, create a new one.
return hook.openCreate(hook.fname)
}
return hook.openAppend(hook.fname)
}
return err
}