From 01178fe2b21f0dce6d3d428baed32c1a7c4fe691 Mon Sep 17 00:00:00 2001
From: Tiago Queiroz <tiago.queiroz@elastic.co>
Date: Thu, 11 Apr 2024 13:10:18 +0200
Subject: [PATCH] Correctly handle log selectors when using TypedCore

---
 logp/core.go           | 12 +++++++-
 logp/selective_test.go | 63 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/logp/core.go b/logp/core.go
index c6240a15..5ad47db5 100644
--- a/logp/core.go
+++ b/logp/core.go
@@ -153,13 +153,21 @@ func ConfigureWithOutputs(defaultLoggerCfg Config, outputs ...zapcore.Core) erro
 //   - `value` is the value compared against the `logKey` entry
 //   - `outputs` is a list of cores that will be added together with the core
 //     generated by `defaultLoggerCfg` as the default output for the loggger.
+//
+// If `defaultLoggerCfg.toObserver` is true, then `typedLoggerCfg` is ignored
+// and a single sink is used so all logs can be observed.
 func ConfigureWithTypedOutput(defaultLoggerCfg, typedLoggerCfg Config, key, value string, outputs ...zapcore.Core) error {
 	sink, level, observedLogs, selectors, err := createSink(defaultLoggerCfg, outputs...)
 	if err != nil {
 		return err
 	}
 
-	typedCore, err := createLogOutput(typedLoggerCfg, level)
+	var typedCore zapcore.Core
+	if defaultLoggerCfg.toObserver {
+		typedCore = sink
+	} else {
+		typedCore, err = createLogOutput(typedLoggerCfg, level)
+	}
 	if err != nil {
 		return fmt.Errorf("could not create typed logger output: %w", err)
 	}
@@ -171,6 +179,8 @@ func ConfigureWithTypedOutput(defaultLoggerCfg, typedLoggerCfg Config, key, valu
 		value:       value,
 	}
 
+	sink = selectiveWrapper(sink, selectors)
+
 	root := zap.New(sink, makeOptions(defaultLoggerCfg)...)
 	storeLogger(&coreLogger{
 		selectors:    selectors,
diff --git a/logp/selective_test.go b/logp/selective_test.go
index e03d5bd5..9533e5c1 100644
--- a/logp/selective_test.go
+++ b/logp/selective_test.go
@@ -54,3 +54,66 @@ func TestLoggerSelectors(t *testing.T) {
 	logs = ObserverLogs().TakeAll()
 	assert.Len(t, logs, 1)
 }
+
+func TestTypedCoreSelectors(t *testing.T) {
+	logSelector := "enabled-log-selector"
+	expectedMsg := "this should be logged"
+
+	defaultCfg := DefaultConfig(DefaultEnvironment)
+	eventsCfg := DefaultEventConfig(DefaultEnvironment)
+
+	defaultCfg.Level = DebugLevel
+	defaultCfg.toObserver = true
+	defaultCfg.ToStderr = false
+	defaultCfg.ToFiles = false
+	defaultCfg.Selectors = []string{logSelector}
+
+	eventsCfg.Level = defaultCfg.Level
+	eventsCfg.toObserver = defaultCfg.toObserver
+	eventsCfg.ToStderr = defaultCfg.ToStderr
+	eventsCfg.ToFiles = defaultCfg.ToFiles
+	eventsCfg.Selectors = defaultCfg.Selectors
+
+	ConfigureWithTypedOutput(defaultCfg, eventsCfg, "log.type", "event")
+
+	enabledSelector := NewLogger(logSelector)
+	disabledSelector := NewLogger("foo-selector")
+
+	enabledSelector.Debugw(expectedMsg)
+	enabledSelector.Debugw(expectedMsg, "log.type", "event")
+	disabledSelector.Debug("this should not be logged")
+
+	logEntries := ObserverLogs().TakeAll()
+	if len(logEntries) != 2 {
+		t.Errorf("expecting 2 log entries, got %d", len(logEntries))
+		t.Log("Log entries:")
+		for _, e := range logEntries {
+			t.Log("Message:", e.Message, "Fields:", e.Context)
+		}
+		t.FailNow()
+	}
+
+	for i, logEntry := range logEntries {
+		msg := logEntry.Message
+		if msg != expectedMsg {
+			t.Fatalf("[%d] expecting log message '%s', got '%s'", i, expectedMsg, msg)
+		}
+
+		// The second entry should also contain `log.type: event`
+		if i == 1 {
+			fields := logEntry.Context
+			if len(fields) != 1 {
+				t.Errorf("expecting one field, got %d", len(fields))
+			}
+
+			k := fields[0].Key
+			v := fields[0].String
+			if k != "log.type" {
+				t.Errorf("expecting key 'log.type', got '%s'", k)
+			}
+			if v != "event" {
+				t.Errorf("expecting value 'event', got '%s'", v)
+			}
+		}
+	}
+}