trufflehog/pkg/log/level.go
Miccah f3367d7910
[THOG-643] Implement independent log level controls (#733)
* [THOG-643] Implement independent log level controls

There are two log level controls to mentally distinguish. Log levels
associated with a sink (e.g. stdout and streamed), and log levels
associated with a logger (e.g. a GitHub source).

The level is determined to be the minimum of the two. If a sink is at
level 0, then it will only output level 0 logs regardless of the
logger's level. This is best demonstrated by TestSinkWithName.

* Rename WithName to WithNamedLevel

* Check flush errors

* Replace IncreaseLevelCore with custom LevelCore

Adding a leveler that was less verbose would cause the initialization
fail, and therefore not be added to the core. This check is only at
the time of initialization.

An alternative approach to creating our own core is to set the child log
level equal to the parent, so initialization is guaranteed (with the
added benefit of intuitive behavior).

* Use controller if it exists, otherwise inherit parent's log level

* Cleanup some tests
2022-08-26 15:27:09 -05:00

114 lines
3.6 KiB
Go

package log
import (
"sort"
"sync"
"github.com/go-logr/logr"
"github.com/go-logr/zapr"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
)
// TODO: Use a struct to make testing easier.
var (
// Global, default log level control.
globalLogLevel levelSetter = zap.NewAtomicLevel()
// Map of name -> level control for independently setting log levels. A new
// control is registered via WithNamedLevel. This map is never cleaned up
// and new entries will overwrite previous values. Currently, this is
// acceptable behavior because WithNamedLevel is used sparingly.
globalControls map[string]levelSetter = make(map[string]levelSetter, 16)
// globalControls is protected (both read and write) by a mutex to make it
// thread safe. Access is low frequency, so performance is not a concern.
globalControlsLock sync.Mutex
)
type levelSetter interface {
zapcore.LevelEnabler
SetLevel(zapcore.Level)
Level() zapcore.Level
}
// SetLevel sets the log level for loggers created with the default level
// controller.
func SetLevel(level int8) {
SetLevelForControl(globalLogLevel, level)
}
// SetLevelForControl sets the log level for a given control.
func SetLevelForControl(control levelSetter, level int8) {
// Zap's levels get more verbose as the number gets smaller, as explained
// by zapr here: https://github.com/go-logr/zapr#increasing-verbosity
// For example setting the level to -2 below, means log.V(2) will be enabled.
control.SetLevel(zapcore.Level(-level))
}
// SetLevelFor sets the log level for a given named control.
func SetLevelFor(name string, level int8) {
globalControlsLock.Lock()
defer globalControlsLock.Unlock()
if control, ok := globalControls[name]; ok {
SetLevelForControl(control, level)
return
}
// Create a new control so registering a control with the same name will
// inherit the existing level.
globalControls[name] = newAtomicLevelAt(level)
}
// AddLeveler adds a log level control to a logr.Logger.
func AddLeveler(l logr.Logger, control levelSetter) (logr.Logger, error) {
zapLogger, err := getZapLogger(l)
if err != nil {
return l, err
}
zapLogger = zapLogger.WithOptions(zap.WrapCore(func(core zapcore.Core) zapcore.Core {
return NewLevelCore(core, control)
}))
return zapr.NewLogger(zapLogger), nil
}
// WithNamedLevel creates a child logger with a new name and independent log
// level control (see SetLevelFor). NOTE: if name already exists, the existing
// controller will be used, otherwise a new controller is created with level
// matching the parent's log level.
func WithNamedLevel(logger logr.Logger, name string) logr.Logger {
logger = logger.WithName(name)
globalControlsLock.Lock()
defer globalControlsLock.Unlock()
var leveler levelSetter
if currentControl, ok := globalControls[name]; ok {
leveler = currentControl
} else {
leveler = newAtomicLevelAt(findLevel(logger))
globalControls[name] = leveler
}
newLogger, err := AddLeveler(logger, leveler)
if err != nil {
return logger
}
return newLogger
}
// newAtomicLevelAt is a helper function to create a zap.AtomicLevel
// initialized with a level. We cannot use zap.NewAtomicLevelAt here because of
// a quirk with logr levels (see SetLevelForControl).
func newAtomicLevelAt(level int8) zap.AtomicLevel {
control := zap.NewAtomicLevel()
SetLevelForControl(control, level)
return control
}
// findLevel probes a logr.Logger to figure out what level it is at via binary
// search. We only search [0, 128), so worst case is ~7 checks.
func findLevel(logger logr.Logger) int8 {
sink := logger.GetSink()
return int8(sort.Search(128, func(i int) bool {
return !sink.Enabled(i)
}) - 1)
}