feat: JobLoggerFactory (#1496)

Remove overriding io.Stdout in TestMaskValues to prevent deadlock in GitHub Actions

pkg/runner/logger.go

@@ -57,31 +57,48 @@ func WithMasks(ctx context.Context, masks *[]string) context.Context {
 	return context.WithValue(ctx, masksContextKeyVal, masks)
 }
 
+type JobLoggerFactory interface {
+	WithJobLogger() *logrus.Logger
+}
+
+type jobLoggerFactoryContextKey string
+
+var jobLoggerFactoryContextKeyVal = (jobLoggerFactoryContextKey)("jobloggerkey")
+
+func WithJobLoggerFactory(ctx context.Context, factory JobLoggerFactory) context.Context {
+	return context.WithValue(ctx, jobLoggerFactoryContextKeyVal, factory)
+}
+
 // WithJobLogger attaches a new logger to context that is aware of steps
 func WithJobLogger(ctx context.Context, jobID string, jobName string, config *Config, masks *[]string, matrix map[string]interface{}) context.Context {
-	mux.Lock()
-	defer mux.Unlock()
-
-	var formatter logrus.Formatter
-	if config.JSONLogger {
-		formatter = &jobLogJSONFormatter{
-			formatter: &logrus.JSONFormatter{},
-			masker:    valueMasker(config.InsecureSecrets, config.Secrets),
-		}
-	} else {
-		formatter = &jobLogFormatter{
-			color:  colors[nextColor%len(colors)],
-			masker: valueMasker(config.InsecureSecrets, config.Secrets),
-		}
-	}
-
-	nextColor++
 	ctx = WithMasks(ctx, masks)
 
-	logger := logrus.New()
+	var logger *logrus.Logger
-	logger.SetFormatter(formatter)
+	if jobLoggerFactory, ok := ctx.Value(jobLoggerFactoryContextKeyVal).(JobLoggerFactory); ok && jobLoggerFactory != nil {
+		logger = jobLoggerFactory.WithJobLogger()
+	} else {
+		var formatter logrus.Formatter
+		if config.JSONLogger {
+			formatter = &logrus.JSONFormatter{}
+		} else {
+			mux.Lock()
+			defer mux.Unlock()
+			nextColor++
+			formatter = &jobLogFormatter{
+				color: colors[nextColor%len(colors)],
+			}
+		}
+
+		logger = logrus.New()
 		logger.SetOutput(os.Stdout)
 		logger.SetLevel(logrus.GetLevel())
+		logger.SetFormatter(formatter)
+	}
+
+	logger.SetFormatter(&maskedFormatter{
+		Formatter: logger.Formatter,
+		masker:    valueMasker(config.InsecureSecrets, config.Secrets),
+	})
 	rtn := logger.WithFields(logrus.Fields{
 		"job":    jobName,
 		"jobID":  jobID,
@@ -149,16 +166,22 @@ func valueMasker(insecureSecrets bool, secrets map[string]string) entryProcessor
 	}
 }
 
+type maskedFormatter struct {
+	logrus.Formatter
+	masker entryProcessor
+}
+
+func (f *maskedFormatter) Format(entry *logrus.Entry) ([]byte, error) {
+	return f.Formatter.Format(f.masker(entry))
+}
+
 type jobLogFormatter struct {
 	color int
-	masker entryProcessor
 }
 
 func (f *jobLogFormatter) Format(entry *logrus.Entry) ([]byte, error) {
 	b := &bytes.Buffer{}
 
-	entry = f.masker(entry)
-
 	if f.isColored(entry) {
 		f.printColored(b, entry)
 	} else {
@@ -225,12 +248,3 @@ func checkIfTerminal(w io.Writer) bool {
 		return false
 	}
 }
-
-type jobLogJSONFormatter struct {
-	masker    entryProcessor
-	formatter *logrus.JSONFormatter
-}
-
-func (f *jobLogJSONFormatter) Format(entry *logrus.Entry) ([]byte, error) {
-	return f.formatter.Format(f.masker(entry))
-}

pkg/runner/runner_test.go

@@ -1,8 +1,10 @@
 package runner
 
 import (
+	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -343,6 +345,17 @@ func TestRunDifferentArchitecture(t *testing.T) {
 	tjfi.runTest(context.Background(), t, &Config{ContainerArchitecture: "linux/arm64"})
 }
 
+type maskJobLoggerFactory struct {
+	Output bytes.Buffer
+}
+
+func (f *maskJobLoggerFactory) WithJobLogger() *log.Logger {
+	logger := log.New()
+	logger.SetOutput(io.MultiWriter(&f.Output, os.Stdout))
+	logger.SetLevel(log.DebugLevel)
+	return logger
+}
+
 func TestMaskValues(t *testing.T) {
 	assertNoSecret := func(text string, secret string) {
 		index := strings.Index(text, "composite secret")
@@ -366,9 +379,9 @@ func TestMaskValues(t *testing.T) {
 		platforms:    platforms,
 	}
 
-	output := captureOutput(t, func() {
+	logger := &maskJobLoggerFactory{}
-		tjfi.runTest(context.Background(), t, &Config{})
+	tjfi.runTest(WithJobLoggerFactory(common.WithLogger(context.Background(), logger.WithJobLogger()), logger), t, &Config{})
-	})
+	output := logger.Output.String()
 
 	assertNoSecret(output, "secret value")
 	assertNoSecret(output, "YWJjCg==")