level.go 1.9 KB

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