Improve caller logging (#2)

* Log caller after message * Add caller logging options * Fix test * Add tests * Change name Co-authored-by: Joseph DeChicchis <[email protected]>
clear-street · Oct 13, 2020 · 4f2977c · 4f2977c
1 parent 0867ffc
commit 4f2977c
Show file tree

Hide file tree

Showing 2 changed files with 150 additions and 15 deletions.
diff --git a/encoder.go b/encoder.go
@@ -42,16 +42,54 @@ func putEncoder(enc *logfmtEncoder) {
 	_logfmtPool.Put(enc)
 }
 
+// Custom logger encoding configuration.
+type config struct {
+	alternativeCallerEncoder zapcore.CallerEncoder
+	callerLogLevel           zapcore.Level
+}
+
 type logfmtEncoder struct {
 	*zapcore.EncoderConfig
 	buf        *buffer.Buffer
 	namespaces []string
+	config     *config
 }
 
-func NewEncoder(cfg zapcore.EncoderConfig) zapcore.Encoder {
+// NewEncoder return a new encoder.
+func NewEncoder(cfg zapcore.EncoderConfig, opts ...Option) zapcore.Encoder {
+	c := &config{
+		alternativeCallerEncoder: nil,
+		callerLogLevel:           zapcore.ErrorLevel,
+	}
+
+	for _, opt := range opts {
+		opt(c)
+	}
+
 	return &logfmtEncoder{
 		EncoderConfig: &cfg,
 		buf:           bufferpool.Get(),
+		config:        c,
+	}
+}
+
+// Option is a function which can be used to modify the default custom config.
+type Option func(*config)
+
+// WithCallerLevel sets the minimum log level for which the caller is logged.
+// e.g. If it is set to WarnLevel the caller will be logged for only WarnLevel and above logs.
+func WithCallerLevel(level zapcore.Level) Option {
+	return func(c *config) {
+		c.callerLogLevel = level
+	}
+}
+
+// WithAlternativeCallerEncoder sets the caller encoder for logs which are less than the caller level.
+// e.g. If set to the ShortCallerEncoder, the caller is still logged for logs which are less than the
+// caller level using the abbreviated format.
+func WithAlternativeCallerEncoder(encoder zapcore.CallerEncoder) Option {
+	return func(c *config) {
+		c.alternativeCallerEncoder = encoder
 	}
 }
 
@@ -268,6 +306,7 @@ func (enc *logfmtEncoder) clone() *logfmtEncoder {
 	clone.EncoderConfig = enc.EncoderConfig
 	clone.buf = bufferpool.Get()
 	clone.namespaces = enc.namespaces
+	clone.config = enc.config
 	return clone
 }
 
@@ -290,16 +329,6 @@ func (enc *logfmtEncoder) EncodeEntry(ent zapcore.Entry, fields []zapcore.Field)
 		final.addKey(enc.NameKey)
 		final.AppendString(ent.LoggerName)
 	}
-	if ent.Caller.Defined && final.CallerKey != "" {
-		final.addKey(final.CallerKey)
-		cur := final.buf.Len()
-		final.EncodeCaller(ent.Caller, final)
-		if cur == final.buf.Len() {
-			// User-supplied EncodeCaller was a no-op. Fall back to strings to
-			// keep output valid.
-			final.AppendString(ent.Caller.String())
-		}
-	}
 	if final.MessageKey != "" {
 		final.addKey(enc.MessageKey)
 		final.AppendString(ent.Message)
@@ -309,6 +338,27 @@ func (enc *logfmtEncoder) EncodeEntry(ent zapcore.Entry, fields []zapcore.Field)
 		final.buf.Write(enc.buf.Bytes())
 	}
 	addFields(final, fields)
+	if ent.Caller.Defined && final.CallerKey != "" {
+		if ent.Level >= final.config.callerLogLevel {
+			final.addKey(final.CallerKey)
+			cur := final.buf.Len()
+			final.EncodeCaller(ent.Caller, final)
+			if cur == final.buf.Len() {
+				// User-supplied EncodeCaller was a no-op. Fall back to strings to
+				// keep output valid.
+				final.AppendString(ent.Caller.String())
+			}
+		} else if final.config.alternativeCallerEncoder != nil {
+			final.addKey(final.CallerKey)
+			cur := final.buf.Len()
+			final.config.alternativeCallerEncoder(ent.Caller, final)
+			if cur == final.buf.Len() {
+				// User-supplied EncodeCaller was a no-op. Fall back to strings to
+				// keep output valid.
+				final.AppendString(ent.Caller.TrimmedPath())
+			}
+		}
+	}
 	if ent.Stack != "" && final.StacktraceKey != "" {
 		final.AddString(final.StacktraceKey, ent.Stack)
 	}

diff --git a/encoder_test.go b/encoder_test.go
@@ -112,11 +112,96 @@ func assertOutput(t testing.TB, desc string, expected string, f func(zapcore.Enc
 }
 
 func TestEncodeCaller(t *testing.T) {
-	enc := &logfmtEncoder{buf: bufferpool.Get(), EncoderConfig: &zapcore.EncoderConfig{
+	enc := NewEncoder(zapcore.EncoderConfig{
 		EncodeTime:     zapcore.EpochTimeEncoder,
 		EncodeDuration: zapcore.SecondsDurationEncoder,
-		EncodeCaller:   zapcore.ShortCallerEncoder,
-	}}
+		EncodeCaller:   zapcore.FullCallerEncoder,
+	},
+		WithCallerLevel(zapcore.DebugLevel),
+	).(*logfmtEncoder)
+
+	var buf *buffer.Buffer
+	var err error
+	encodeEntry := func() {
+		buf, err = enc.EncodeEntry(
+			zapcore.Entry{
+				Level:      zapcore.DebugLevel,
+				Time:       time.Time{},
+				LoggerName: "test",
+				Message:    "caller test",
+				Caller: zapcore.EntryCaller{
+					Defined: true,
+					File:    "h2g2.go",
+					Line:    42,
+				},
+			},
+			[]zapcore.Field{
+				zap.String("k", "v"),
+			},
+		)
+	}
+
+	encodeEntry()
+	assert.Nil(t, err)
+	assert.Equal(t, "k=v\n", buf.String())
+
+	enc.truncate()
+	enc.EncoderConfig.CallerKey = "caller"
+	encodeEntry()
+	assert.Nil(t, err)
+	assert.Equal(t, "k=v caller=h2g2.go:42\n", buf.String())
+}
+
+func TestEncodeCallerNoAlternative(t *testing.T) {
+	enc := NewEncoder(zapcore.EncoderConfig{
+		EncodeTime:     zapcore.EpochTimeEncoder,
+		EncodeDuration: zapcore.SecondsDurationEncoder,
+		EncodeCaller:   zapcore.FullCallerEncoder,
+	},
+		WithCallerLevel(zapcore.ErrorLevel),
+	).(*logfmtEncoder)
+
+	var buf *buffer.Buffer
+	var err error
+	encodeEntry := func() {
+		buf, err = enc.EncodeEntry(
+			zapcore.Entry{
+				Level:      zapcore.DebugLevel,
+				Time:       time.Time{},
+				LoggerName: "test",
+				Message:    "caller test",
+				Caller: zapcore.EntryCaller{
+					Defined: true,
+					File:    "h2g2.go",
+					Line:    42,
+				},
+			},
+			[]zapcore.Field{
+				zap.String("k", "v"),
+			},
+		)
+	}
+
+	encodeEntry()
+	assert.Nil(t, err)
+	assert.Equal(t, "k=v\n", buf.String())
+
+	enc.truncate()
+	enc.EncoderConfig.CallerKey = "caller"
+	encodeEntry()
+	assert.Nil(t, err)
+	assert.Equal(t, "k=v\n", buf.String())
+}
+
+func TestEncodeCallerWithAlternative(t *testing.T) {
+	enc := NewEncoder(zapcore.EncoderConfig{
+		EncodeTime:     zapcore.EpochTimeEncoder,
+		EncodeDuration: zapcore.SecondsDurationEncoder,
+		EncodeCaller:   zapcore.FullCallerEncoder,
+	},
+		WithAlternativeCallerEncoder(zapcore.ShortCallerEncoder),
+		WithCallerLevel(zapcore.ErrorLevel),
+	).(*logfmtEncoder)
 
 	var buf *buffer.Buffer
 	var err error
@@ -147,7 +232,7 @@ func TestEncodeCaller(t *testing.T) {
 	enc.EncoderConfig.CallerKey = "caller"
 	encodeEntry()
 	assert.Nil(t, err)
-	assert.Equal(t, "caller=h2g2.go:42 k=v\n", buf.String())
+	assert.Equal(t, "k=v caller=h2g2.go:42\n", buf.String())
 }
 
 func TestEncodeStacktrace(t *testing.T) {