talos

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
package logging
6

7
import (
8
	"io"
9
	"log"
10
	"strings"
11

12
	"github.com/siderolabs/gen/xslices"
13
	"go.uber.org/zap"
14
	"go.uber.org/zap/zapcore"
15
)
16

17
// LogWriter is a wrapper around zap.Logger that implements io.Writer interface.
18
type LogWriter struct {
19
	dest  *zap.Logger
20
	level zapcore.Level
21
}
22

23
// NewWriter creates new log zap log writer.
24
func NewWriter(l *zap.Logger, level zapcore.Level) io.Writer {
25
	return &LogWriter{
26
		dest:  l,
27
		level: level,
28
	}
29
}
30

31
// Write implements io.Writer interface.
32
func (lw *LogWriter) Write(line []byte) (int, error) {
33
	checked := lw.dest.Check(lw.level, strings.TrimSpace(string(line)))
34
	if checked == nil {
35
		return 0, nil
36
	}
37

38
	checked.Write()
39

40
	return len(line), nil
41
}
42

43
// logWrapper wraps around standard logger.
44
type logWrapper struct {
45
	log *log.Logger
46
}
47

48
// Write implements io.Writer interface.
49
func (lw *logWrapper) Write(line []byte) (int, error) {
50
	if lw.log == nil {
51
		log.Print(string(line))
52
	} else {
53
		lw.log.Print(string(line))
54
	}
55

56
	return len(line), nil
57
}
58

59
// StdWriter creates a sync writer that writes all logs to the std logger.
60
var StdWriter = &logWrapper{nil}
61

62
// LogDestination defines logging destination Config.
63
type LogDestination struct {
64
	// Level log level.
65
	level             zapcore.LevelEnabler
66
	writer            io.Writer
67
	config            zapcore.EncoderConfig
68
	suppressThreshold int64
69
}
70

71
// LogDestinationOption defines a log destination encoder config setter.
72
type LogDestinationOption func(dest *LogDestination)
73

74
// WithoutTimestamp disables timestamp.
75
func WithoutTimestamp() LogDestinationOption {
76
	return func(dest *LogDestination) {
77
		dest.config.EncodeTime = nil
78
	}
79
}
80

81
// WithoutLogLevels disable log level.
82
func WithoutLogLevels() LogDestinationOption {
83
	return func(dest *LogDestination) {
84
		dest.config.EncodeLevel = nil
85
	}
86
}
87

88
// WithColoredLevels enables log level colored output.
89
func WithColoredLevels() LogDestinationOption {
90
	return func(dest *LogDestination) {
91
		dest.config.EncodeLevel = zapcore.CapitalColorLevelEncoder
92
	}
93
}
94

95
// WithControllerErrorSuppressor creates a new console log controller error suppressor.
96
func WithControllerErrorSuppressor(threshold int64) LogDestinationOption {
97
	return func(dest *LogDestination) {
98
		dest.suppressThreshold = threshold
99
	}
100
}
101

102
// NewLogDestination creates new log destination.
103
func NewLogDestination(writer io.Writer, logLevel zapcore.LevelEnabler, options ...LogDestinationOption) *LogDestination {
104
	config := zap.NewDevelopmentEncoderConfig()
105
	config.ConsoleSeparator = " "
106
	config.StacktraceKey = "error"
107

108
	dest := &LogDestination{
109
		level:  logLevel,
110
		config: config,
111
		writer: writer,
112
	}
113

114
	for _, option := range options {
115
		option(dest)
116
	}
117

118
	return dest
119
}
120

121
// Wrap is a simple helper to wrap io.Writer with default arguments.
122
func Wrap(writer io.Writer) *zap.Logger {
123
	return ZapLogger(
124
		NewLogDestination(writer, zapcore.DebugLevel),
125
	)
126
}
127

128
// ZapLogger creates new default Zap Logger.
129
func ZapLogger(dests ...*LogDestination) *zap.Logger {
130
	if len(dests) == 0 {
131
		panic("at least one writer must be defined")
132
	}
133

134
	cores := xslices.Map(dests, func(dest *LogDestination) zapcore.Core {
135
		core := zapcore.NewCore(
136
			zapcore.NewConsoleEncoder(dest.config),
137
			zapcore.AddSync(dest.writer),
138
			dest.level,
139
		)
140

141
		if dest.suppressThreshold > 0 {
142
			core = NewControllerErrorSuppressor(core, dest.suppressThreshold)
143
		}
144

145
		return core
146
	})
147

148
	return zap.New(zapcore.NewTee(cores...))
149
}
150

151
// Component helper for creating zap.Field.
152
func Component(name string) zapcore.Field {
153
	return zap.String("component", name)
154
}
155