Reduce allocations for disabled/sampled logs (uber-go#45)

* Reduce allocations for disabled/sampled logs

Even when a log level is disabled or heavily sampled, calling
`logger.Error(msg, someField, someOtherField)` will make at least one
allocation to hold the fields. If we're using the more expensive field
types, we may allocate even more, only to discard all that work.

This PR introduces the `logger.Check` API, which allows particularly
performance-sensitive applications to avoid these allocations by using
a slightly wordier call pattern.

Using `Check` makes it virtually free to pass fields to a sampled
logger:

```
BenchmarkZapDisabledLevelsAddingFields-4         3000000        439 ns/op             704 B/op          2 allocs/op
BenchmarkZapDisabledLevelsCheckAddingFields-4   200000000         5.74 ns/op            0 B/op          0 allocs/op

BenchmarkZapSampleWithoutFields-4       20000000                60.9 ns/op             0 B/op          0 allocs/op
BenchmarkZapSampleCheckWithoutFields-4  20000000                61.4 ns/op             0 B/op          0 allocs/op

BenchmarkZapSampleAddingFields-4         2000000               575 ns/op             704 B/op          2 allocs/op
BenchmarkZapSampleCheckAddingFields-4   20000000                65.2 ns/op             0 B/op          0 allocs/op
```

* Improve CheckedMessage ergonomics

* Don't log on CheckedMessage re-use

benchmarks/zap_bench_test.go

@@ -66,6 +66,14 @@ func fakeFields() []zap.Field {
 	}
 }
 
+func fakeMessages(n int) []string {
+	messages := make([]string, n)
+	for i := range messages {
+		messages[i] = fmt.Sprintf("Test logging, but use a somewhat realistic message length. (#%v)", i)
+	}
+	return messages
+}
+
 func BenchmarkZapDisabledLevelsWithoutFields(b *testing.B) {
 	logger := zap.NewJSON(zap.Error, zap.Output(zap.Discard))
 	b.ResetTimer()
@@ -96,6 +104,19 @@ func BenchmarkZapDisabledLevelsAddingFields(b *testing.B) {
 		}
 	})
 }
+
+func BenchmarkZapDisabledLevelsCheckAddingFields(b *testing.B) {
+	logger := zap.NewJSON(zap.Error, zap.Output(zap.Discard))
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			if m := logger.Check(zap.Info, "Should be discarded."); m.OK() {
+				m.Write(fakeFields()...)
+			}
+		}
+	})
+}
+
 func BenchmarkZapAddingFields(b *testing.B) {
 	logger := zap.NewJSON(zap.All, zap.Output(zap.Discard))
 	b.ResetTimer()
@@ -128,12 +149,9 @@ func BenchmarkZapWithoutFields(b *testing.B) {
 }
 
 func BenchmarkZapSampleWithoutFields(b *testing.B) {
-	messages := make([]string, 1000)
-	for i := range messages {
-		messages[i] = fmt.Sprintf("Sample the logs, but use a somewhat realistic message length. (#%v)", i)
-	}
+	messages := fakeMessages(1000)
 	base := zap.NewJSON(zap.All, zap.Output(zap.Discard))
-	logger := zwrap.Sample(base, time.Second, 10, 100)
+	logger := zwrap.Sample(base, time.Second, 10, 10000)
 	b.ResetTimer()
 	b.RunParallel(func(pb *testing.PB) {
 		i := 0
@@ -143,3 +161,49 @@ func BenchmarkZapSampleWithoutFields(b *testing.B) {
 		}
 	})
 }
+
+func BenchmarkZapSampleAddingFields(b *testing.B) {
+	messages := fakeMessages(1000)
+	base := zap.NewJSON(zap.All, zap.Output(zap.Discard))
+	logger := zwrap.Sample(base, time.Second, 10, 10000)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		i := 0
+		for pb.Next() {
+			i++
+			logger.Info(messages[i%1000], fakeFields()...)
+		}
+	})
+}
+
+func BenchmarkZapSampleCheckWithoutFields(b *testing.B) {
+	messages := fakeMessages(1000)
+	base := zap.NewJSON(zap.All, zap.Output(zap.Discard))
+	logger := zwrap.Sample(base, time.Second, 10, 10000)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		i := 0
+		for pb.Next() {
+			i++
+			if cm := logger.Check(zap.Info, messages[i%1000]); cm.OK() {
+				cm.Write()
+			}
+		}
+	})
+}
+
+func BenchmarkZapSampleCheckAddingFields(b *testing.B) {
+	messages := fakeMessages(1000)
+	base := zap.NewJSON(zap.All, zap.Output(zap.Discard))
+	logger := zwrap.Sample(base, time.Second, 10, 10000)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		i := 0
+		for pb.Next() {
+			i++
+			if m := logger.Check(zap.Info, messages[i%1000]); m.OK() {
+				m.Write(fakeFields()...)
+			}
+		}
+	})
+}

checked_message.go

@@ -0,0 +1,63 @@
+// Copyright (c) 2016 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package zap
+
+import "sync/atomic"
+
+// A CheckedMessage is the result of a call to Logger.Check, which allows
+// especially performance-sensitive applications to avoid allocations for disabled
+// or heavily sampled log levels.
+type CheckedMessage struct {
+	logger Logger
+	used   uint32
+	lvl    Level
+	msg    string
+}
+
+// NewCheckedMessage constructs a CheckedMessage. It's only intended for use by
+// wrapper libraries, and shouldn't be necessary in application code.
+func NewCheckedMessage(logger Logger, lvl Level, msg string) *CheckedMessage {
+	return &CheckedMessage{
+		logger: logger,
+		lvl:    lvl,
+		msg:    msg,
+	}
+}
+
+// Write logs the pre-checked message with the supplied fields. It should only
+// be used once; if a CheckedMessage is re-used, it also logs an error message
+// with the underlying logger's DFatal method.
+func (m *CheckedMessage) Write(fields ...Field) {
+	if n := atomic.AddUint32(&m.used, 1); n > 1 {
+		if n == 2 {
+			// Log an error on the first re-use. After that, skip the I/O and
+			// allocations and just return.
+			m.logger.DFatal("Shouldn't re-use a CheckedMessage.", String("original", m.msg))
+		}
+		return
+	}
+	m.logger.Log(m.lvl, m.msg, fields...)
+}
+
+// OK checks whether it's safe to call Write.
+func (m *CheckedMessage) OK() bool {
+	return m != nil
+}

checked_message_test.go

@@ -0,0 +1,58 @@
+// Copyright (c) 2016 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package zap
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestJSONLoggerCheck(t *testing.T) {
+	withJSONLogger(t, opts(Info), func(jl *jsonLogger, output func() []string) {
+		assert.False(t, jl.Check(Debug, "Debug.").OK(), "Expected CheckedMessage to be not OK at disabled levels.")
+
+		cm := jl.Check(Info, "Info.")
+		require.True(t, cm.OK(), "Expected CheckedMessage to be OK at enabled levels.")
+		cm.Write(Int("magic", 42))
+		assert.Equal(
+			t,
+			`{"msg":"Info.","level":"info","ts":0,"fields":{"magic":42}}`,
+			output()[0],
+			"Unexpected output after writing a CheckedMessage.",
+		)
+	})
+}
+
+func TestCheckedMessageIsSingleUse(t *testing.T) {
+	expected := []string{
+		`{"msg":"Single-use.","level":"info","ts":0,"fields":{}}`,
+		`{"msg":"Shouldn't re-use a CheckedMessage.","level":"error","ts":0,"fields":{"original":"Single-use."}}`,
+	}
+	withJSONLogger(t, nil, func(jl *jsonLogger, output func() []string) {
+		cm := jl.Check(Info, "Single-use.")
+		cm.Write() // ok
+		cm.Write() // first re-use logs error
+		cm.Write() // second re-use is silently ignored
+		assert.Equal(t, expected, output(), "Expected re-using a CheckedMessage to log an error.")
+	})
+}

example_test.go

@@ -97,3 +97,26 @@ func ExampleNewJSON_options() {
 	// {"msg":"This is a debug log.","level":"debug","ts":0,"fields":{"count":1}}
 	// {"msg":"This is an info log.","level":"info","ts":0,"fields":{"count":1}}
 }
+
+func ExampleCheckedMessage() {
+	logger := zap.NewJSON()
+	// Stub the current time in tests.
+	logger.StubTime()
+
+	// By default, the debug logging level is disabled. However, calls to
+	// logger.Debug will still allocate a slice to hold any passed fields.
+	// Particularly performance-sensitive applications can avoid paying this
+	// penalty by using checked messages.
+	if cm := logger.Check(zap.Debug, "This is a debug log."); cm.OK() {
+		// Debug-level logging is disabled, so we won't get here.
+		cm.Write(zap.Int("foo", 42), zap.Stack())
+	}
+
+	if cm := logger.Check(zap.Info, "This is an info log."); cm.OK() {
+		// Since info-level logging is enabled, we expect to write out this message.
+		cm.Write()
+	}
+
+	// Output:
+	// {"msg":"This is an info log.","level":"info","ts":0,"fields":{}}
+}

logger.go

@@ -51,6 +51,13 @@ type Logger interface {
 	// option.
 	StubTime()
 
+	// Check returns a CheckedMessage if logging a message at the specified level
+	// is enabled. It's a completely optional optimization; in high-performance
+	// applications, Check can help avoid allocating a slice to hold fields.
+	//
+	// See CheckedMessage for an example.
+	Check(Level, string) *CheckedMessage
+
 	// Log a message at the given level. Messages include any context that's
 	// accumulated on the logger, as well as any fields added at the log site.
 	Log(Level, string, ...Field)
@@ -133,6 +140,13 @@ func (jl *jsonLogger) StubTime() {
 	jl.alwaysEpoch = true
 }
 
+func (jl *jsonLogger) Check(lvl Level, msg string) *CheckedMessage {
+	if !jl.Enabled(lvl) {
+		return nil
+	}
+	return NewCheckedMessage(jl, lvl, msg)
+}
+
 func (jl *jsonLogger) Log(lvl Level, msg string, fields ...Field) {
 	switch lvl {
 	case Panic:

logger_test.go

@@ -28,8 +28,9 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
 	"github.com/uber-go/zap/spywrite"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func opts(opts ...Option) []Option {
@@ -70,7 +71,10 @@ func withJSONLogger(t testing.TB, opts []Option, f func(*jsonLogger, func() []st
 	jl := NewJSON(allOpts...)
 	jl.StubTime()
 
-	f(jl.(*jsonLogger), func() []string { return strings.Split(sink.String(), "\n") })
+	f(jl.(*jsonLogger), func() []string {
+		output := strings.Split(sink.String(), "\n")
+		return output[:len(output)-1]
+	})
 	assert.Empty(t, errSink.String(), "Expected error sink to be empty")
 }
 
@@ -245,7 +249,7 @@ func TestJSONLoggerNoOpsDisabledLevels(t *testing.T) {
 	withJSONLogger(t, nil, func(jl *jsonLogger, output func() []string) {
 		jl.SetLevel(Warn)
 		jl.Info("silence!")
-		assert.Equal(t, []string{""}, output(), "Expected logging at a disabled level to produce no output.")
+		assert.Equal(t, []string{}, output(), "Expected logging at a disabled level to produce no output.")
 	})
 }
 

spy/logger.go

@@ -122,6 +122,14 @@ func (l *Logger) With(fields ...zap.Field) zap.Logger {
 	}
 }
 
+// Check returns a CheckedMessage if logging a particular message would succeed.
+func (l *Logger) Check(lvl zap.Level, msg string) *zap.CheckedMessage {
+	if !l.Enabled(lvl) {
+		return nil
+	}
+	return zap.NewCheckedMessage(l, lvl, msg)
+}
+
 // Log writes a message at the specified level.
 func (l *Logger) Log(lvl zap.Level, msg string, fields ...zap.Field) {
 	l.sink.WriteLog(lvl, msg, l.allFields(fields))

zwrap/sample.go

@@ -97,6 +97,13 @@ func (s *sampler) With(fields ...zap.Field) zap.Logger {
 	}
 }
 
+func (s *sampler) Check(lvl zap.Level, msg string) *zap.CheckedMessage {
+	if !s.check(lvl, msg) {
+		return nil
+	}
+	return zap.NewCheckedMessage(s.Logger, lvl, msg)
+}
+
 func (s *sampler) Log(lvl zap.Level, msg string, fields ...zap.Field) {
 	if s.check(lvl, msg) {
 		s.Logger.Log(lvl, msg, fields...)

zwrap/sample_test.go

@@ -165,6 +165,22 @@ func TestSamplerTicks(t *testing.T) {
 	assert.Equal(t, expected, sink.Logs(), "Expected sleeping for a tick to reset sampler.")
 }
 
+func TestSamplerCheck(t *testing.T) {
+	sampler, sink := fakeSampler(time.Millisecond, 1, 10, false)
+	sampler.SetLevel(zap.Info)
+
+	assert.Nil(t, sampler.Check(zap.Debug, "foo"), "Expected a nil CheckedMessage at disabled log levels.")
+
+	for i := 1; i < 12; i++ {
+		if cm := sampler.Check(zap.Info, "sample"); cm.OK() {
+			cm.Write(zap.Int("iter", i))
+		}
+	}
+
+	expected := buildExpectation(zap.Info, 1, 11)
+	assert.Equal(t, expected, sink.Logs(), "Unexpected output when sampling with Check.")
+}
+
 func TestSamplerRaces(t *testing.T) {
 	sampler, _ := fakeSampler(time.Minute, 1, 1000, false)