1// Package logutils augments the standard log package with levels.
2package logutils
3
4import (
5	"bytes"
6	"io"
7	"sync"
8)
9
10type LogLevel string
11
12// LevelFilter is an io.Writer that can be used with a logger that
13// will filter out log messages that aren't at least a certain level.
14//
15// Once the filter is in use somewhere, it is not safe to modify
16// the structure.
17type LevelFilter struct {
18	// Levels is the list of log levels, in increasing order of
19	// severity. Example might be: {"DEBUG", "WARN", "ERROR"}.
20	Levels []LogLevel
21
22	// MinLevel is the minimum level allowed through
23	MinLevel LogLevel
24
25	// The underlying io.Writer where log messages that pass the filter
26	// will be set.
27	Writer io.Writer
28
29	badLevels map[LogLevel]struct{}
30	once      sync.Once
31}
32
33// Check will check a given line if it would be included in the level
34// filter.
35func (f *LevelFilter) Check(line []byte) bool {
36	f.once.Do(f.init)
37
38	// Check for a log level
39	var level LogLevel
40	x := bytes.IndexByte(line, '[')
41	if x >= 0 {
42		y := bytes.IndexByte(line[x:], ']')
43		if y >= 0 {
44			level = LogLevel(line[x+1 : x+y])
45		}
46	}
47
48	_, ok := f.badLevels[level]
49	return !ok
50}
51
52func (f *LevelFilter) Write(p []byte) (n int, err error) {
53	// Note in general that io.Writer can receive any byte sequence
54	// to write, but the "log" package always guarantees that we only
55	// get a single line. We use that as a slight optimization within
56	// this method, assuming we're dealing with a single, complete line
57	// of log data.
58
59	if !f.Check(p) {
60		return len(p), nil
61	}
62
63	return f.Writer.Write(p)
64}
65
66// SetMinLevel is used to update the minimum log level
67func (f *LevelFilter) SetMinLevel(min LogLevel) {
68	f.MinLevel = min
69	f.init()
70}
71
72func (f *LevelFilter) init() {
73	badLevels := make(map[LogLevel]struct{})
74	for _, level := range f.Levels {
75		if level == f.MinLevel {
76			break
77		}
78		badLevels[level] = struct{}{}
79	}
80	f.badLevels = badLevels
81}
82