Add NewGoKitWithFields constructor

Signed-off-by: Yuri Nikolic <[email protected]>

log/gokit.go

@@ -17,32 +17,68 @@ const (
 	LogfmtFormat = "logfmt"
 )
 
+type Fields []interface{}
+
+func NewFields(fields ...interface{}) Fields {
+	return fields
+}
+
+func (f Fields) empty() bool {
+	return len(f) == 0
+}
+
 // NewGoKit creates a new GoKit logger with the given level, format and writer,
 // enriched with the standard "ts" and "caller" fields.
 // If the given writer is nil, os.Stderr is used.
 // If the given format is nil, logfmt is used.
 func NewGoKit(lvl Level, format string, writer io.Writer) log.Logger {
-	if writer == nil {
-		writer = log.NewSyncWriter(os.Stderr)
-	}
-	if format == "json" {
-		return log.NewJSONLogger(writer)
-	}
-	logger := addStandardFields(log.NewLogfmtLogger(writer))
+	logger := newGoKit(format, writer, nil)
+	logger = addStandardFields(logger)
 	return level.NewFilter(logger, lvl.Option)
 }
 
-// NewGoKitWithFields creates a new GoKit logger with the given level and writer,
+// NewGoKitLogFmt creates a new GoKit logger with the given level and writer,
 // enriched with the standard "ts" and "caller" fields.
 // The "logfmt" format is used.
 // If the given writer is nil, os.Stderr is used.
-func NewGoKitWithFields(l Level, writer io.Writer) log.Logger {
+func NewGoKitLogFmt(l Level, writer io.Writer) log.Logger {
 	return NewGoKit(l, "logfmt", writer)
 }
 
+// NewGoKitWithFields creates a new GoKit logger configured with the given level, format, writer and
+// rate limit-related configuration, enriched with the given fields.
+// If the given format is nil, logfmt is used.
+// If the given writer is nil, os.Stderr is used.
+// If the given rate limit configuration is nil, no rate limited logger is created.
+// If the given fields are empty, no fields are added.
+func NewGoKitWithFields(lvl Level, format string, writer io.Writer, rateLimitedCfg *RateLimitedLoggerCfg, fields Fields) log.Logger {
+	logger := newGoKit(format, writer, rateLimitedCfg)
+	if !fields.empty() {
+		logger = log.With(logger, fields...)
+	}
+	return level.NewFilter(logger, lvl.Option)
+}
+
+func newGoKit(format string, writer io.Writer, rateLimitedCfg *RateLimitedLoggerCfg) log.Logger {
+	var logger log.Logger
+	if writer == nil {
+		writer = log.NewSyncWriter(os.Stderr)
+	}
+	if format == "json" {
+		logger = log.NewJSONLogger(writer)
+	} else {
+		logger = log.NewLogfmtLogger(writer)
+	}
+	if rateLimitedCfg != nil {
+		logger = NewRateLimitedLogger(logger, *rateLimitedCfg)
+	}
+	return logger
+}
+
 // stand-alone for test purposes
 func addStandardFields(logger log.Logger) log.Logger {
-	return log.With(logger, "ts", log.DefaultTimestampUTC, "caller", log.Caller(5))
+	fields := NewFields("ts", log.DefaultTimestampUTC, "caller", log.Caller(5))
+	return log.With(logger, fields...)
 }
 
 type Sprintf struct {

log/gokit_test.go

@@ -12,12 +12,13 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
 )
 
 func BenchmarkLazySprintf(b *testing.B) {
 	g := log.NewNopLogger()
-	logger := addStandardFields(g)
+	logger := level.NewFilter(addStandardFields(g), level.AllowInfo())
 	// Simulate the parameters used in middleware/logging.go
 	var (
 		method     = "method"
@@ -55,10 +56,10 @@ func TestLazySprintf(t *testing.T) {
 		var lvl Level
 		require.NoError(t, lvl.Set(test.lvl))
 		buf.Reset()
-		logger := NewGoKitWithFields(lvl, buf)
+		logger := NewGoKitLogFmt(lvl, buf)
 		now := time.Now()
 		expectedMessage := fmt.Sprintf(test.format, test.id, now)
-		lazySprintf := newLazySprintfWithCount("debug %d has been logged %v", []interface{}{test.id, now})
+		lazySprintf := newLazySprintfWithCount("debug %d has been logged %v", test.id, now)
 		level.Debug(logger).Log("msg", lazySprintf)
 		if test.lvl == "debug" {
 			require.True(t, bytes.Contains(buf.Bytes(), []byte(expectedMessage)))
@@ -70,12 +71,41 @@ func TestLazySprintf(t *testing.T) {
 	}
 }
 
+func TestNewGoKitWithFields(t *testing.T) {
+	buf := bytes.NewBuffer(nil)
+	var lvl Level
+	require.NoError(t, lvl.Set("info"))
+	rateLimitedCfg := &RateLimitedLoggerCfg{
+		LogsPerSecond:      1,
+		LogsPerSecondBurst: 2,
+		Registry:           prometheus.DefaultRegisterer,
+	}
+	now := log.DefaultTimestampUTC()
+	fields := NewFields("ts", now, "caller", log.Caller(5))
+	logger := NewGoKitWithFields(lvl, LogfmtFormat, buf, rateLimitedCfg, fields)
+	for i := 0; i < 1000; i++ {
+		level.Info(logger).Log("msg", LazySprintf("info %d", i))
+		level.Debug(logger).Log("msg", LazySprintf("debug %d", i))
+	}
+
+	format := "ts=%s caller=gokit_test.go:%d level=%s msg=\"info %d\""
+	for i := 0; i < 1000; i++ {
+		if i < 2 {
+			require.True(t, bytes.Contains(buf.Bytes(), []byte(fmt.Sprintf(format, now, 87, "info", i))))
+		} else {
+			require.False(t, bytes.Contains(buf.Bytes(), []byte(fmt.Sprintf(format, now, 87, "info", i))))
+		}
+
+		require.False(t, bytes.Contains(buf.Bytes(), []byte(fmt.Sprintf(format, now, 88, "debug", i))))
+	}
+}
+
 type lazySprintfWithCount struct {
 	next  Sprintf
 	count int
 }
 
-func newLazySprintfWithCount(format string, args []interface{}) *lazySprintfWithCount {
+func newLazySprintfWithCount(format string, args ...interface{}) *lazySprintfWithCount {
 	return &lazySprintfWithCount{
 		Sprintf{
 			format: format,

log/ratelimit.go

@@ -8,6 +8,12 @@ import (
 	"golang.org/x/time/rate"
 )
 
+type RateLimitedLoggerCfg struct {
+	LogsPerSecond      float64
+	LogsPerSecondBurst int
+	Registry           prometheus.Registerer
+}
+
 type RateLimitedLogger struct {
 	next    log.Logger
 	limiter *rate.Limiter
@@ -20,15 +26,15 @@ type RateLimitedLogger struct {
 
 // NewRateLimitedLogger returns a log.Logger that is limited to the given number of logs per second,
 // with the given burst size.
-func NewRateLimitedLogger(logger log.Logger, logsPerSecond rate.Limit, burstSize int, reg prometheus.Registerer) log.Logger {
-	discardedLogLinesCounter := promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
+func NewRateLimitedLogger(logger log.Logger, cfg RateLimitedLoggerCfg) log.Logger {
+	discardedLogLinesCounter := promauto.With(cfg.Registry).NewCounterVec(prometheus.CounterOpts{
 		Name: "logger_rate_limit_discarded_log_lines_total",
 		Help: "Total number of discarded log lines per level.",
 	}, []string{"level"})
 
 	return &RateLimitedLogger{
 		next:                          logger,
-		limiter:                       rate.NewLimiter(logsPerSecond, burstSize),
+		limiter:                       rate.NewLimiter(rate.Limit(cfg.LogsPerSecond), cfg.LogsPerSecondBurst),
 		discardedInfoLogLinesCounter:  discardedLogLinesCounter.WithLabelValues(level.InfoValue().String()),
 		discardedDebugLogLinesCounter: discardedLogLinesCounter.WithLabelValues(level.DebugValue().String()),
 		discardedWarnLogLinesCounter:  discardedLogLinesCounter.WithLabelValues(level.WarnValue().String()),

log/ratelimit_test.go

@@ -18,7 +18,12 @@ func TestRateLimitedLoggerLogs(t *testing.T) {
 	buf := bytes.NewBuffer(nil)
 	c := newCounterLogger(buf)
 	reg := prometheus.NewPedanticRegistry()
-	r := NewRateLimitedLogger(c, 1, 1, reg)
+	cfg := RateLimitedLoggerCfg{
+		LogsPerSecond:      1,
+		LogsPerSecondBurst: 1,
+		Registry:           reg,
+	}
+	r := NewRateLimitedLogger(c, cfg)
 
 	level.Error(r).Log("msg", "error will be logged")
 	assert.Equal(t, 1, c.count)
@@ -40,7 +45,12 @@ func TestRateLimitedLoggerLimits(t *testing.T) {
 	buf := bytes.NewBuffer(nil)
 	c := newCounterLogger(buf)
 	reg := prometheus.NewPedanticRegistry()
-	r := NewRateLimitedLogger(c, 2, 2, reg)
+	cfg := RateLimitedLoggerCfg{
+		LogsPerSecond:      2,
+		LogsPerSecondBurst: 2,
+		Registry:           reg,
+	}
+	r := NewRateLimitedLogger(c, cfg)
 
 	level.Error(r).Log("msg", "error 1 will be logged")
 	assert.Equal(t, 1, c.count)
@@ -99,7 +109,12 @@ func TestRateLimitedLoggerWithFields(t *testing.T) {
 	buf := bytes.NewBuffer(nil)
 	c := newCounterLogger(buf)
 	reg := prometheus.NewPedanticRegistry()
-	logger := NewRateLimitedLogger(c, 0.0001, 1, reg)
+	cfg := RateLimitedLoggerCfg{
+		LogsPerSecond:      1,
+		LogsPerSecondBurst: 1,
+		Registry:           reg,
+	}
+	logger := NewRateLimitedLogger(c, cfg)
 	loggerWithFields := log.With(logger, "key", "value")
 
 	level.Error(loggerWithFields).Log("msg", "error will be logged")

middleware/logging_test.go

@@ -35,7 +35,7 @@ func TestBadWriteLogging(t *testing.T) {
 		buf := bytes.NewBuffer(nil)
 		var level log.Level
 		require.NoError(t, level.Set("debug"))
-		logger := log.NewGoKitWithFields(level, buf)
+		logger := log.NewGoKitLogFmt(level, buf)
 
 		loggingMiddleware := Log{
 			Log: logger,
@@ -80,7 +80,7 @@ func TestDisabledSuccessfulRequestsLogging(t *testing.T) {
 		require.NoError(t, level.Set("debug"))
 
 		loggingMiddleware := Log{
-			Log:                      log.NewGoKitWithFields(level, buf),
+			Log:                      log.NewGoKitLogFmt(level, buf),
 			DisableRequestSuccessLog: tc.disableLog,
 		}
 
@@ -125,7 +125,7 @@ func TestLoggingRequestsAtInfoLevel(t *testing.T) {
 		require.NoError(t, level.Set("debug"))
 
 		loggingMiddleware := Log{
-			Log:                   log.NewGoKitWithFields(level, buf),
+			Log:                   log.NewGoKitLogFmt(level, buf),
 			LogRequestAtInfoLevel: true,
 		}
 		handler := func(w http.ResponseWriter, r *http.Request) {
@@ -179,7 +179,7 @@ func TestLoggingRequestWithExcludedHeaders(t *testing.T) {
 			var level log.Level
 			require.NoError(t, level.Set("debug"))
 
-			loggingMiddleware := NewLogMiddleware(log.NewGoKitWithFields(level, buf), true, false, nil, tc.excludeHeaderList)
+			loggingMiddleware := NewLogMiddleware(log.NewGoKitLogFmt(level, buf), true, false, nil, tc.excludeHeaderList)
 
 			handler := func(w http.ResponseWriter, r *http.Request) {
 				_, _ = io.WriteString(w, "<html><body>Hello world!</body></html>")

server/server_test.go

@@ -723,7 +723,7 @@ type FakeLogger struct {
 
 func newFakeLogger(level log.Level) *FakeLogger {
 	buf := bytes.NewBuffer(nil)
-	log := log.NewGoKitWithFields(level, buf)
+	log := log.NewGoKitLogFmt(level, buf)
 	return &FakeLogger{
 		logger: log,
 		buf:    buf,

server/server_tracing_test.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/grafana/dskit/httpgrpc"
 	httpgrpcServer "github.com/grafana/dskit/httpgrpc/server"
+	"github.com/grafana/dskit/log"
 	"github.com/grafana/dskit/middleware"
 )
 
@@ -201,6 +202,9 @@ func TestHTTPGRPCTracing(t *testing.T) {
 			cfg.GRPCServerMaxSendMsgSize = 4 * 1024 * 1024
 			cfg.Router = middleware.InitHTTPGRPCMiddleware(mux.NewRouter())
 			cfg.MetricsNamespace = "testing_httpgrpc_tracing_" + middleware.MakeLabelValue(testName)
+			var lvl log.Level
+			require.NoError(t, lvl.Set("info"))
+			cfg.LogLevel = lvl
 
 			server, err := New(cfg)
 			require.NoError(t, err)