internal/log: add log level support

This change introduces minimum log level support.
GO_DISCOVERY_LOG_LEVEL config can be used to set the log level.

When set, no logs below GO_DISCOVERY_LOG_LEVEL will be printed.
Possible log levels are debug, info, error, fatal.

Log severity sequence is debug < info < error < fatal.
In case of invalid/empty log level, it print all logs.

Fixes golang/go#40339

Change-Id: I66116e6778eb6386bf1e8df6947078a78fece691
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/245797
Reviewed-by: Julie Qiu <julie@golang.org>

cmd/frontend/main.go

@@ -57,6 +57,9 @@ func main() {
 			log.Fatalf(ctx, "profiler.Start: %v", err)
 		}
 	}
+
+	log.SetLevel(cfg.LogLevel)
+
 	var (
 		dsg        func(context.Context) internal.DataSource
 		expg       func(context.Context) internal.ExperimentSource

cmd/prober/main.go

@@ -222,6 +222,7 @@ func main() {
 	}
 	cfg.Dump(os.Stderr)
 
+	log.SetLevel(cfg.LogLevel)
 	if _, err := log.UseStackdriver(ctx, cfg, "prober-log"); err != nil {
 		log.Fatal(ctx, err)
 	}

cmd/teeproxy/main.go

@@ -37,6 +37,8 @@ func main() {
 		log.Fatal(ctx, err)
 	}
 	cfg.Dump(os.Stderr)
+
+	log.SetLevel(cfg.LogLevel)
 	if cfg.OnAppEngine() {
 		_, err := log.UseStackdriver(ctx, cfg, "teeproxy-log")
 		if err != nil {

cmd/worker/main.go

@@ -56,6 +56,8 @@ func main() {
 	}
 	cfg.Dump(os.Stderr)
 
+	log.SetLevel(cfg.LogLevel)
+
 	if cfg.UseProfiler {
 		if err := profiler.Start(profiler.Config{}); err != nil {
 			log.Fatalf(ctx, "profiler.Start: %v", err)

internal/config/config.go

@@ -123,6 +123,11 @@ type Config struct {
 	// TeeproxyTargetHosts is a list of hosts that teeproxy will forward
 	// requests to.
 	TeeproxyForwardedHosts []string
+
+	// Minimum log level below which no logs will be printed.
+	// Possible values are [debug, info, error, fatal].
+	// In case of invalid/empty value, all logs will be printed.
+	LogLevel string
 }
 
 // AppVersionLabel returns the version label for the current instance.  This is
@@ -272,6 +277,7 @@ func Init(ctx context.Context) (_ *Config, err error) {
 		UseProfiler:            os.Getenv("GO_DISCOVERY_USE_PROFILER") == "TRUE",
 		TeeproxyAuthValue:      os.Getenv("GO_DISCOVERY_TEEPROXY_AUTH_VALUE"),
 		TeeproxyForwardedHosts: parseCommaList(os.Getenv("GO_DISCOVERY_TEEPROXY_FORWARDED_HOSTS")),
+		LogLevel:               os.Getenv("GO_DISCOVERY_LOG_LEVEL"),
 	}
 	cfg.AppMonitoredResource = &mrpb.MonitoredResource{
 		Type: "gae_app",

internal/log/log.go

@@ -25,6 +25,10 @@ var (
 	logger interface {
 		log(context.Context, logging.Severity, interface{})
 	} = stdlibLogger{}
+
+	// currentLevel holds current log level.
+	// No logs will be printed below currentLevel.
+	currentLevel = logging.Default
 )
 
 type (
@@ -35,6 +39,19 @@ type (
 	labelsKey struct{}
 )
 
+// Set the log level
+func SetLevel(v string) {
+	mu.Lock()
+	defer mu.Unlock()
+	currentLevel = toLevel(v)
+}
+
+func getLevel() logging.Severity {
+	mu.Lock()
+	defer mu.Unlock()
+	return currentLevel
+}
+
 // NewContextWithTraceID creates a new context from ctx that adds the trace ID.
 func NewContextWithTraceID(ctx context.Context, traceID string) context.Context {
 	return context.WithValue(ctx, traceIDKey{}, traceID)
@@ -156,9 +173,9 @@ func Debugf(ctx context.Context, format string, args ...interface{}) {
 	logf(ctx, logging.Debug, format, args)
 }
 
-// Fatalf is equivalent to Errorf followed by exiting the program.
+// Fatalf logs formatted string at the Critical level followed by exiting the program.
 func Fatalf(ctx context.Context, format string, args ...interface{}) {
-	Errorf(ctx, format, args...)
+	logf(ctx, logging.Critical, format, args)
 	die()
 }
 
@@ -175,13 +192,16 @@ func Error(ctx context.Context, arg interface{}) { doLog(ctx, logging.Error, arg
 // Debug logs arg, which can be a string or a struct, at the Debug level.
 func Debug(ctx context.Context, arg interface{}) { doLog(ctx, logging.Debug, arg) }
 
-// Fatal is equivalent to Error followed by exiting the program.
+// Fatal logs arg, which can be a string or a struct, at the Critical level followed by exiting the program.
 func Fatal(ctx context.Context, arg interface{}) {
-	Error(ctx, arg)
+	doLog(ctx, logging.Critical, arg)
 	die()
 }
 
 func doLog(ctx context.Context, s logging.Severity, payload interface{}) {
+	if getLevel() > s {
+		return
+	}
 	mu.Lock()
 	l := logger
 	mu.Unlock()
@@ -196,3 +216,28 @@ func die() {
 	mu.Unlock()
 	os.Exit(1)
 }
+
+// toLevel returns the logging.Severity for a given string.
+// Possible input values are "", "debug", "info", "error", "fatal".
+// In case of invalid string input, it maps to DefaultLevel.
+func toLevel(v string) logging.Severity {
+	v = strings.ToLower(v)
+
+	switch v {
+	case "":
+		// default log level will print everything.
+		return logging.Default
+	case "debug":
+		return logging.Debug
+	case "info":
+		return logging.Info
+	case "error":
+		return logging.Error
+	case "fatal":
+		return logging.Critical
+	}
+
+	// Default log level in case of invalid input.
+	log.Printf("Error: %s is invalid LogLevel. Possible values are [debug, info, error, fatal]", v)
+	return logging.Default
+}

internal/log/log_test.go

@@ -0,0 +1,121 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package log
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"testing"
+
+	"cloud.google.com/go/logging"
+)
+
+const (
+	debugMsg = "debugMsg"
+	infoMsg  = "infoMsg"
+	errorMsg = "errorMsg"
+)
+
+// Do not run in parallel. It overrides currentLevel.
+func TestSetLogLevel(t *testing.T) {
+	oldLevel := getLevel()
+	defer func() { currentLevel = oldLevel }()
+
+	tests := []struct {
+		name      string
+		newLevel  string
+		wantLevel logging.Severity
+	}{
+		{name: "default level", newLevel: "", wantLevel: logging.Default},
+		{name: "invalid level", newLevel: "xyz", wantLevel: logging.Default},
+		{name: "debug level", newLevel: "debug", wantLevel: logging.Debug},
+		{name: "info level", newLevel: "info", wantLevel: logging.Info},
+		{name: "error level", newLevel: "error", wantLevel: logging.Error},
+		{name: "fatal level", newLevel: "fatal", wantLevel: logging.Critical},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			SetLevel(test.newLevel)
+			gotLevel := getLevel()
+			if test.wantLevel != gotLevel {
+				t.Errorf("Error: want=%s, got=%s", test.wantLevel, gotLevel)
+			}
+		})
+	}
+}
+
+// Do not run in parallel. It overrides logger with mockLogger.
+func TestLogLevel(t *testing.T) {
+	oldLogger := logger
+	defer func() { logger = oldLogger }()
+	logger = &mockLogger{}
+
+	// logs below info(like debug) won't print
+	SetLevel("info")
+
+	tests := []struct {
+		name     string
+		logFunc  func(context.Context, interface{})
+		logMsg   string
+		expected bool
+	}{
+		{name: "debug", logFunc: Debug, logMsg: debugMsg, expected: false},
+		{name: "info", logFunc: Info, logMsg: infoMsg, expected: true},
+		{name: "error", logFunc: Error, logMsg: errorMsg, expected: true},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+
+			test.logFunc(context.Background(), test.logMsg)
+			logs := logger.(*mockLogger).logs
+			got := strings.Contains(logs, test.logMsg)
+
+			if got != test.expected {
+				t.Errorf("expected : %v, got %v", test.expected, got)
+			}
+		})
+	}
+}
+
+// Do not run in parallel. It overrides logger with mockLogger.
+func TestDefaultLogLevel(t *testing.T) {
+	oldLogger := logger
+	defer func() { logger = oldLogger }()
+	logger = &mockLogger{}
+
+	SetLevel("") // default behaviour; print everything
+
+	tests := []struct {
+		name    string
+		logFunc func(context.Context, interface{})
+		logMsg  string
+	}{
+		{name: "debug", logFunc: Debug, logMsg: debugMsg},
+		{name: "info", logFunc: Info, logMsg: infoMsg},
+		{name: "error", logFunc: Error, logMsg: errorMsg},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+
+			test.logFunc(context.Background(), test.logMsg)
+			logs := logger.(*mockLogger).logs
+
+			if !strings.Contains(logs, test.logMsg) {
+				t.Errorf("%v not logged.", test.logMsg)
+			}
+		})
+	}
+}
+
+type mockLogger struct {
+	logs string
+}
+
+func (l *mockLogger) log(ctx context.Context, s logging.Severity, payload interface{}) {
+	l.logs += fmt.Sprintf("%s: %+v", s, payload)
+}