feat: add ability to use custom zap core

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)

setup_test.go

@@ -36,7 +36,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 +202,41 @@ func TestSubsystemLevels(t *testing.T) {
 		t.Errorf("got %q, wanted it to contain info2", buf.String())
 	}
 }
+
+func TestCustomCore(t *testing.T) {
+	r1, w1, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("failed to open pipe: %v", err)
+	}
+	r2, w2, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("failed to open pipe: %v", err)
+	}
+
+	// logging should work with the custom core
+	SetPrimaryCore(newCore(PlaintextOutput, w1, LevelDebug))
+	log := getLogger("test")
+	log.Error("scooby")
+
+	// 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 _, 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())
+	}
+}