From 4dc1bbf981ed4d3fea8d629406c4a1a679550bb4 Mon Sep 17 00:00:00 2001
From: Patrick Ohly <patrick.ohly@intel.com>
Date: Thu, 10 Mar 2022 17:54:48 +0100
Subject: [PATCH] contextual logging: hide the additional interface

If we only ever expect this to be used via the wrapper function, then we don't
need an interface at all and can simply check for our own type.
---
 contextual.go      | 52 +++++++++++++++++++---------------------------
 contextual_test.go | 16 ++++++--------
 2 files changed, 27 insertions(+), 41 deletions(-)

diff --git a/contextual.go b/contextual.go
index 1ed9b968..dc6cfaae 100644
--- a/contextual.go
+++ b/contextual.go
@@ -65,41 +65,37 @@ var (
 // To remove a backing logr implemention, use ClearLogger. Setting an
 // empty logger with SetLogger(logr.Logger{}) does not work.
 //
-// If the logger instance implements the ContextualLogger interface in addition
-// to logr.Logger, then the instance will also be called directly
-// by code that retrieves a logger with FromContext, TODO, or Background.
-// In that case the logger instance should do its own verbosity checks
-// in Enabled because klog will not do that for it.
-//
-// AddSupportsContext can be used for loggers which don't know about this
-// klog specific convention.
+// Ideally, loggers should implement the verbosity check in Enabled themselves.
+// Then they can get called directly by application code after retrieving the
+// logger with FromContext, Background, or TODO. SetLogger needs to know
+// whether the logger is ready for that. Wrapping it with WithContextLogging
+// provides that information.
 //
 // Modifying the logger is not thread-safe and should be done while no other
 // goroutines invoke log calls, usually during program initialization.
 func SetLogger(logger logr.Logger) {
-	globalLogger = &logger
 	contextualLogger = false
-	if contextual, ok := logger.GetSink().(ContextualLogSink); ok && contextual.SupportsContext() {
+	if contextual, ok := logger.GetSink().(contextualLogSinkWrapper); ok {
 		contextualLogger = true
+		// Strip the extra layer, we don't need it anymore.
+		logger = logr.New(contextual.LogSink)
 	}
-}
 
-// ContextualLogSink is an extension of the logr.LogSink interface that
-// indicates to SetLogger whether the logger may get called directly by
-// application code. Only the LogSink seen by SetLogger needs to implement
-// this, clones of it as they occur during logger.V, logger.WithName or
-// logger.WithValues don't need it.
-type ContextualLogSink interface {
-	logr.LogSink
-
-	// SupportsContext returns true if the logger may get called directly
-	// by application code.
-	SupportsContext() bool
+	globalLogger = &logger
 }
 
-// AddSupportsContext wraps a logger that doesn't know about ContextualLogSink
-// so that SetLogger knows that it can also be called directly.
-func AddSupportsContext(logger logr.Logger) logr.Logger {
+// WithContextLogging wraps a logger so that SetLogger knows that the logger
+// can also be called directly.
+//
+// Experimental
+//
+// Notice: This function is EXPERIMENTAL and may be changed or removed in a
+// later release.
+func WithContextLogging(logger logr.Logger) logr.Logger {
+	// We can only wrap the LogSink interface but not
+	// the Logger struct. Besides the LogSink, Logger also
+	// contains a level. We cannot clone that one,
+	// so this is slightly broken:
 	return logr.New(contextualLogSinkWrapper{logger.GetSink()})
 }
 
@@ -107,12 +103,6 @@ type contextualLogSinkWrapper struct {
 	logr.LogSink
 }
 
-func (c contextualLogSinkWrapper) SupportsContext() bool {
-	return true
-}
-
-var _ ContextualLogSink = contextualLogSinkWrapper{}
-
 // ClearLogger removes a backing Logger implementation if one was set earlier
 // with SetLogger.
 //
diff --git a/contextual_test.go b/contextual_test.go
index 32128691..04cf881b 100644
--- a/contextual_test.go
+++ b/contextual_test.go
@@ -24,17 +24,13 @@ import (
 )
 
 func ExampleSetLogger() {
-	logger := logr.Discard()
-	klog.SetLogger(logger)
-	equal := logger == klog.Background()
-	fmt.Printf("original logr.Discard logger can be used directly: %v\n", equal)
+	klog.SetLogger(logr.Discard())
+	fmt.Printf("after setting logr.Discard as logger: %T\n", klog.Background().GetSink())
 
-	logger = klog.AddSupportsContext(logger)
-	klog.SetLogger(logger)
-	equal = logger == klog.Background()
-	fmt.Printf("wrapped logger can be used directly: %v\n", equal)
+	klog.SetLogger(klog.WithContextLogging(logr.Discard()))
+	fmt.Printf("after wrapping it: %T\n", klog.Background().GetSink())
 
 	// Output:
-	// original logr.Discard logger can be used directly: false
-	// wrapped logger can be used directly: true
+	// after setting logr.Discard as logger: *klog.klogger
+	// after wrapping it: logr.discardLogSink
 }