From 60a3410f2dc71e9564d30a095d6f63d4b6060b61 Mon Sep 17 00:00:00 2001
From: Michal Dobaczewski <mdobak@gmail.com>
Date: Tue, 1 Jun 2021 06:48:38 +0200
Subject: [PATCH 1/2] feat: add ability to use custom zap core

---
 setup.go      | 26 +++++++++++++++++++-------
 setup_test.go | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/setup.go b/setup.go
index 06064952e7e8..75fed3f5a99f 100644
--- a/setup.go
+++ b/setup.go
@@ -134,13 +134,7 @@ func SetupLogging(cfg Config) {
 		newPrimaryCore = newPrimaryCore.With([]zap.Field{zap.String(k, v)})
 	}
 
-	if primaryCore != nil {
-		loggerCore.ReplaceCore(primaryCore, newPrimaryCore)
-	} else {
-		loggerCore.AddCore(newPrimaryCore)
-	}
-	primaryCore = newPrimaryCore
-
+	setPrimaryCore(newPrimaryCore)
 	setAllLoggers(defaultLevel)
 
 	for name, level := range cfg.SubsystemLevels {
@@ -152,6 +146,24 @@ func SetupLogging(cfg Config) {
 	}
 }
 
+// SetPrimaryCore changes the primary logging core. If the SetupLogging was
+// called then the previously configured core will be replaced.
+func SetPrimaryCore(core zapcore.Core) {
+	loggerMutex.Lock()
+	defer loggerMutex.Unlock()
+
+	setPrimaryCore(core)
+}
+
+func setPrimaryCore(core zapcore.Core) {
+	if primaryCore != nil {
+		loggerCore.ReplaceCore(primaryCore, core)
+	} else {
+		loggerCore.AddCore(core)
+	}
+	primaryCore = core
+}
+
 // SetDebugLogging calls SetAllLoggers with logging.DEBUG
 func SetDebugLogging() {
 	SetAllLoggers(LevelDebug)
diff --git a/setup_test.go b/setup_test.go
index 691c6eee8ffa..337f2f762b61 100644
--- a/setup_test.go
+++ b/setup_test.go
@@ -7,6 +7,8 @@ import (
 	"os"
 	"strings"
 	"testing"
+
+	"go.uber.org/zap"
 )
 
 func TestGetLoggerDefault(t *testing.T) {
@@ -36,7 +38,6 @@ func TestGetLoggerDefault(t *testing.T) {
 	if !strings.Contains(buf.String(), "scooby") {
 		t.Errorf("got %q, wanted it to contain log output", buf.String())
 	}
-
 }
 
 func TestLogToFileAndStderr(t *testing.T) {
@@ -203,3 +204,37 @@ func TestSubsystemLevels(t *testing.T) {
 		t.Errorf("got %q, wanted it to contain info2", buf.String())
 	}
 }
+
+func TestCustomCore(t *testing.T) {
+	r, w, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("failed to open pipe: %v", err)
+	}
+
+	stderr := os.Stderr
+	os.Stderr = w
+	defer func() {
+		os.Stderr = stderr
+	}()
+
+	ws, _, err := zap.Open("stdout", "stderr")
+	if err != nil {
+		t.Fatalf("unable to open logging output: %v", err)
+	}
+
+	// logging should work with the custom core
+	SetPrimaryCore(newCore(PlaintextOutput, ws, LevelDebug))
+	log := getLogger("test")
+
+	log.Error("scooby")
+	w.Close()
+
+	buf := &bytes.Buffer{}
+	if _, err := io.Copy(buf, r); err != nil && err != io.ErrClosedPipe {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if !strings.Contains(buf.String(), "scooby") {
+		t.Errorf("got %q, wanted it to contain log output", buf.String())
+	}
+}

From 109a95fa73e4aa68b771bd97103323bad1b96627 Mon Sep 17 00:00:00 2001
From: Michal Dobaczewski <mdobak@gmail.com>
Date: Fri, 4 Jun 2021 01:37:36 +0200
Subject: [PATCH 2/2] test: test whether the core is replaced in already
 created loggers

---
 setup_test.go | 42 ++++++++++++++++++++++--------------------
 1 file changed, 22 insertions(+), 20 deletions(-)

diff --git a/setup_test.go b/setup_test.go
index 337f2f762b61..d3b1c28f3456 100644
--- a/setup_test.go
+++ b/setup_test.go
@@ -7,8 +7,6 @@ import (
 	"os"
 	"strings"
 	"testing"
-
-	"go.uber.org/zap"
 )
 
 func TestGetLoggerDefault(t *testing.T) {
@@ -206,35 +204,39 @@ func TestSubsystemLevels(t *testing.T) {
 }
 
 func TestCustomCore(t *testing.T) {
-	r, w, err := os.Pipe()
+	r1, w1, err := os.Pipe()
 	if err != nil {
 		t.Fatalf("failed to open pipe: %v", err)
 	}
-
-	stderr := os.Stderr
-	os.Stderr = w
-	defer func() {
-		os.Stderr = stderr
-	}()
-
-	ws, _, err := zap.Open("stdout", "stderr")
+	r2, w2, err := os.Pipe()
 	if err != nil {
-		t.Fatalf("unable to open logging output: %v", err)
+		t.Fatalf("failed to open pipe: %v", err)
 	}
 
 	// logging should work with the custom core
-	SetPrimaryCore(newCore(PlaintextOutput, ws, LevelDebug))
+	SetPrimaryCore(newCore(PlaintextOutput, w1, LevelDebug))
 	log := getLogger("test")
-
 	log.Error("scooby")
-	w.Close()
 
-	buf := &bytes.Buffer{}
-	if _, err := io.Copy(buf, r); err != nil && err != io.ErrClosedPipe {
+	// SetPrimaryCore should replace the core in previously created loggers
+	SetPrimaryCore(newCore(PlaintextOutput, w2, LevelDebug))
+	log.Error("doo")
+
+	w1.Close()
+	w2.Close()
+
+	buf1 := &bytes.Buffer{}
+	buf2 := &bytes.Buffer{}
+	if _, err := io.Copy(buf1, r1); err != nil && err != io.ErrClosedPipe {
 		t.Fatalf("unexpected error: %v", err)
 	}
-
-	if !strings.Contains(buf.String(), "scooby") {
-		t.Errorf("got %q, wanted it to contain log output", buf.String())
+	if _, err := io.Copy(buf2, r2); err != nil && err != io.ErrClosedPipe {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !strings.Contains(buf1.String(), "scooby") {
+		t.Errorf("got %q, wanted it to contain log output", buf1.String())
+	}
+	if !strings.Contains(buf2.String(), "doo") {
+		t.Errorf("got %q, wanted it to contain log output", buf2.String())
 	}
 }