Add a rate limit on logger

CHANGELOG.md

@@ -53,6 +53,7 @@
 * [FEATURE] Add support for waiting on the rate limiter using the new `WaitN` method. #279
 * [FEATURE] Add `log.BufferedLogger` type. #338
 * [FEATURE] Add `flagext.ParseFlagsAndArguments()` and `flagext.ParseFlagsWithoutArguments()` utilities. #341
+* [FEATURE] Add `log.RateLimitedLogger` for limiting the rate of logging. The `rate_limit_logger_discarded_log_lines_total` metrics traces the total number of discarded log lines per level. #352
 * [ENHANCEMENT] Add configuration to customize backoff for the gRPC clients.
 * [ENHANCEMENT] Use `SecretReader` interface to fetch secrets when configuring TLS. #274
 * [ENHANCEMENT] Add middleware package. #38

log/ratelimit.go

@@ -0,0 +1,116 @@
+package log
+
+import (
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	"golang.org/x/time/rate"
+)
+
+const (
+	infoLevel  = "info"
+	debugLevel = "debug"
+	warnLevel  = "warning"
+	errorLevel = "error"
+)
+
+type RateLimitedLogger struct {
+	next    Interface
+	limiter *rate.Limiter
+
+	discardedLogLinesCounter *prometheus.CounterVec
+}
+
+// NewRateLimitedLogger returns a logger.Interface that is limited to the given number of logs per second,
+// with the given burst size.
+func NewRateLimitedLogger(logger Interface, logsPerSecond rate.Limit, burstSize int, reg prometheus.Registerer) Interface {
+	discardedLogLinesCounter := promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
+		Name: "rate_limit_logger_discarded_log_lines_total",
+		Help: "Total number of discarded log lines per level.",
+	}, []string{"level"})
+
+	return &RateLimitedLogger{
+		next:                     logger,
+		limiter:                  rate.NewLimiter(logsPerSecond, burstSize),
+		discardedLogLinesCounter: discardedLogLinesCounter,
+	}
+}
+
+func (l *RateLimitedLogger) Debugf(format string, args ...interface{}) {
+	if l.limiter.Allow() {
+		l.next.Debugf(format, args...)
+	} else {
+		l.discardedLogLinesCounter.WithLabelValues(debugLevel).Inc()
+	}
+}
+
+func (l *RateLimitedLogger) Debugln(args ...interface{}) {
+	if l.limiter.Allow() {
+		l.next.Debugln(args...)
+	} else {
+		l.discardedLogLinesCounter.WithLabelValues(debugLevel).Inc()
+	}
+}
+
+func (l *RateLimitedLogger) Infof(format string, args ...interface{}) {
+	if l.limiter.Allow() {
+		l.next.Infof(format, args...)
+	} else {
+		l.discardedLogLinesCounter.WithLabelValues(infoLevel).Inc()
+	}
+}
+
+func (l *RateLimitedLogger) Infoln(args ...interface{}) {
+	if l.limiter.Allow() {
+		l.next.Infoln(args...)
+	} else {
+		l.discardedLogLinesCounter.WithLabelValues(infoLevel).Inc()
+	}
+}
+
+func (l *RateLimitedLogger) Errorf(format string, args ...interface{}) {
+	if l.limiter.Allow() {
+		l.next.Errorf(format, args...)
+	} else {
+		l.discardedLogLinesCounter.WithLabelValues(errorLevel).Inc()
+	}
+}
+
+func (l *RateLimitedLogger) Errorln(args ...interface{}) {
+	if l.limiter.Allow() {
+		l.next.Errorln(args...)
+	} else {
+		l.discardedLogLinesCounter.WithLabelValues(errorLevel).Inc()
+	}
+}
+
+func (l *RateLimitedLogger) Warnf(format string, args ...interface{}) {
+	if l.limiter.Allow() {
+		l.next.Warnf(format, args...)
+	} else {
+		l.discardedLogLinesCounter.WithLabelValues(warnLevel).Inc()
+	}
+}
+
+func (l *RateLimitedLogger) Warnln(args ...interface{}) {
+	if l.limiter.Allow() {
+		l.next.Warnln(args...)
+	} else {
+		l.discardedLogLinesCounter.WithLabelValues(warnLevel).Inc()
+	}
+}
+
+func (l *RateLimitedLogger) WithField(key string, value interface{}) Interface {
+	return &RateLimitedLogger{
+		next:                     l.next.WithField(key, value),
+		limiter:                  l.limiter,
+		discardedLogLinesCounter: l.discardedLogLinesCounter,
+	}
+}
+
+func (l *RateLimitedLogger) WithFields(f Fields) Interface {
+	return &RateLimitedLogger{
+		next:                     l.next.WithFields(f),
+		limiter:                  l.limiter,
+		discardedLogLinesCounter: l.discardedLogLinesCounter,
+	}
+}

log/ratelimit_test.go

@@ -0,0 +1,176 @@
+package log
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRateLimitedLoggerLogs(t *testing.T) {
+	buf := bytes.NewBuffer(nil)
+	c := newCounterLogger(buf)
+	reg := prometheus.NewPedanticRegistry()
+	r := NewRateLimitedLogger(c, 1, 1, reg)
+
+	r.Errorln("Error will be logged")
+	assert.Equal(t, 1, c.count)
+
+	logContains := []string{"error", "Error will be logged"}
+	c.verify(t, logContains)
+}
+
+func TestRateLimitedLoggerLimits(t *testing.T) {
+	buf := bytes.NewBuffer(nil)
+	c := newCounterLogger(buf)
+	reg := prometheus.NewPedanticRegistry()
+	r := NewRateLimitedLogger(c, 2, 2, reg)
+
+	r.Errorln("error 1 will be logged")
+	assert.Equal(t, 1, c.count)
+	c.verify(t, []string{"error", "error 1 will be logged"})
+
+	r.Infoln("info 1 will be logged")
+	assert.Equal(t, 2, c.count)
+	c.verify(t, []string{"info", "info 1 will be logged"})
+
+	r.Debugln("debug 1 will be discarded")
+	assert.Equal(t, 2, c.count)
+
+	r.Warnln("warning 1 will be discarded")
+	assert.Equal(t, 2, c.count)
+
+	require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(`
+		# HELP rate_limit_logger_discarded_log_lines_total Total number of discarded log lines per level.
+		# TYPE rate_limit_logger_discarded_log_lines_total counter
+        rate_limit_logger_discarded_log_lines_total{level="debug"} 1
+        rate_limit_logger_discarded_log_lines_total{level="warning"} 1
+	`)))
+
+	// we wait 1 second, so the next group of lines can be logged
+	time.Sleep(time.Second)
+	r.Debugln("debug 2 will be logged")
+	assert.Equal(t, 3, c.count)
+	c.verify(t, []string{"debug", "debug 2 will be logged"})
+
+	r.Infoln("info 2 will be logged")
+	assert.Equal(t, 4, c.count)
+	c.verify(t, []string{"info", "info 2 will be logged"})
+
+	r.Errorln("error 2 will be discarded")
+	assert.Equal(t, 4, c.count)
+
+	r.Warnln("warning 2 will be discarded")
+	assert.Equal(t, 4, c.count)
+
+	require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(`
+		# HELP rate_limit_logger_discarded_log_lines_total Total number of discarded log lines per level.
+		# TYPE rate_limit_logger_discarded_log_lines_total counter
+        rate_limit_logger_discarded_log_lines_total{level="debug"} 1
+        rate_limit_logger_discarded_log_lines_total{level="error"} 1
+        rate_limit_logger_discarded_log_lines_total{level="warning"} 2
+	`)))
+}
+
+func TestRateLimitedLoggerWithFields(t *testing.T) {
+	buf := bytes.NewBuffer(nil)
+	c := newCounterLogger(buf)
+	reg := prometheus.NewPedanticRegistry()
+	logger := NewRateLimitedLogger(c, 0.0001, 1, reg)
+	loggerWithFields := logger.WithField("key", "value")
+
+	loggerWithFields.Errorln("Error will be logged")
+	assert.Equal(t, 1, c.count)
+	c.verify(t, []string{"key", "value", "error", "Error will be logged"})
+
+	logger.Infoln("Info will not be logged")
+	loggerWithFields.Debugln("Debug will not be logged")
+	loggerWithFields.Warnln("Warning will not be logged")
+	assert.Equal(t, 1, c.count)
+
+	require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(`
+		# HELP rate_limit_logger_discarded_log_lines_total Total number of discarded log lines per level.
+		# TYPE rate_limit_logger_discarded_log_lines_total counter
+        rate_limit_logger_discarded_log_lines_total{level="info"} 1
+        rate_limit_logger_discarded_log_lines_total{level="debug"} 1
+        rate_limit_logger_discarded_log_lines_total{level="warning"} 1
+	`)))
+}
+
+type counterLogger struct {
+	logger Interface
+	buf    *bytes.Buffer
+	count  int
+}
+
+func (c *counterLogger) Debugf(format string, args ...interface{}) {
+	c.logger.Debugf(format, args...)
+	c.count++
+}
+
+func (c *counterLogger) Debugln(args ...interface{}) {
+	c.logger.Debugln(args...)
+	c.count++
+}
+
+func (c *counterLogger) Infof(format string, args ...interface{}) {
+	c.logger.Infof(format, args...)
+	c.count++
+}
+
+func (c *counterLogger) Infoln(args ...interface{}) {
+	c.logger.Infoln(args...)
+	c.count++
+}
+
+func (c *counterLogger) Warnf(format string, args ...interface{}) {
+	c.logger.Warnf(format, args...)
+	c.count++
+}
+
+func (c *counterLogger) Warnln(args ...interface{}) {
+	c.logger.Warnln(args...)
+	c.count++
+}
+
+func (c *counterLogger) Errorf(format string, args ...interface{}) {
+	c.logger.Errorf(format, args...)
+	c.count++
+}
+
+func (c *counterLogger) Errorln(args ...interface{}) {
+	c.logger.Errorln(args...)
+	c.count++
+}
+
+func (c *counterLogger) WithField(key string, value interface{}) Interface {
+	c.logger = c.logger.WithField(key, value)
+	return c
+}
+
+func (c *counterLogger) WithFields(fields Fields) Interface {
+	c.logger = c.logger.WithFields(fields)
+	return c
+}
+
+func (c *counterLogger) verify(t *testing.T, logContains []string) {
+	for _, content := range logContains {
+		require.True(t, bytes.Contains(c.buf.Bytes(), []byte(content)))
+	}
+}
+
+func newCounterLogger(buf *bytes.Buffer) *counterLogger {
+	logrusLogger := logrus.New()
+	logrusLogger.Out = buf
+	logrusLogger.Level = logrus.DebugLevel
+	return &counterLogger{
+		logger: Logrus(logrusLogger),
+		buf:    buf,
+	}
+}