/
new.go
122 lines (109 loc) · 2.66 KB
/
new.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
package logwriter
import (
"fmt"
"io"
"log"
"os"
"path"
"path/filepath"
"strings"
"time"
)
type OpenOption struct {
// Path to file or directory.
// If empty string specified, discard all writes operations.
// If "-" is specified, send to stderr.
// If existing directory path specified, file name is generated automatically.
FileOrDir string
// Prefix for file name.
// This option only affect if FileOrDir points to a directory.
Prefix string
// Additional file extensions.
// This option only affect if FileOrDir points to a directory.
//
// Supported extensions list:
// ".zst"
// ".gz"
// "" (without compression)
Suffix string
// Flag for open a file.
Flag int
// Default file mode.
Mode os.FileMode
}
var DefaultOpenOption = OpenOption{
FileOrDir: "-",
Prefix: filepath.Base(os.Args[0]),
Suffix: ".zst",
Flag: os.O_WRONLY | os.O_APPEND | os.O_CREATE,
Mode: 0666,
}
type TearDown func() error
func Setup(option OpenOption) (TearDown, error) {
w, err := Open(option)
if err != nil {
return nil, err
}
old := log.Writer()
log.SetOutput(w)
return func() error {
log.SetOutput(old)
return w.Close()
}, nil
}
func Open(option OpenOption) (io.WriteCloser, error) {
w, ok := openFast(option)
if ok {
return w, nil
}
return openSlow(option)
}
func openFast(option OpenOption) (io.WriteCloser, bool) {
p := option.FileOrDir
if p == "" || p == "/dev/null" {
// Drop all logs.
return Discard, true
}
if p == "-" {
// Send to stderr.
// To prevent closing stderr unexpectedly, wraps os.Stderr by nopCloserFile.
return &nopCloserFile{os.Stderr}, true
}
return nil, false
}
func openSlow(opt OpenOption) (io.WriteCloser, error) {
filePath := opt.FileOrDir
stat, err := os.Stat(opt.FileOrDir)
if err != nil && !os.IsNotExist(err) {
// Unexpected error occurred.
return nil, err
}
if err == nil && stat.IsDir() {
dir := opt.FileOrDir
filename := fmt.Sprintf(
"%s.%s-%d.log%s",
opt.Prefix,
time.Now().Format(time.RFC3339Nano), os.Getpid(),
opt.Suffix,
)
filePath = path.Join(dir, filename)
}
// Ignore os.ErrNotExist.
return openSuitableLogger(filePath, opt)
}
func openSuitableLogger(filePath string, opt OpenOption) (w io.WriteCloser, err error) {
w, err = os.OpenFile(filePath, opt.Flag, opt.Mode)
if err != nil {
return nil, err
}
// Select compression algorithm
if strings.HasSuffix(filePath, ".gz") {
w = NewCompressedWriter(w, &GzipAlgorithm{})
} else if strings.HasSuffix(filePath, ".zst") {
w = NewCompressedWriter(w, &ZstdAlgorithm{})
}
// Add write buffer to improve compression efficiency.
w = NewBuffer(os.Getpagesize(), time.Second, w)
w = NewTickWriter(w, time.Second)
return
}