Pool CheckedMessage (free after Write) (uber-go#191)

* Pool CheckedMessages
* Replace atomic used with weaker safeToWrite bool

checked_message.go

@@ -20,16 +20,24 @@
 
 package zap
 
-import "github.com/uber-go/atomic"
+import "sync"
+
+var _cmPool = sync.Pool{
+	New: func() interface{} {
+		return &CheckedMessage{
+			lvl: invalidLevel,
+		}
+	},
+}
 
 // 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   atomic.Uint32
-	lvl    Level
-	msg    string
+	logger      Logger
+	safeToWrite bool
+	lvl         Level
+	msg         string
 
 	// singly linked list built by Chain
 	next *CheckedMessage // carried by each part of Chain-ed list
@@ -41,33 +49,39 @@ type CheckedMessage struct {
 // 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,
-	}
+	m := _cmPool.Get().(*CheckedMessage)
+	m.safeToWrite, m.logger, m.lvl, m.msg = true, logger, lvl, msg
+	return m
 }
 
-// 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.
+// Write logs the pre-checked message with the supplied fields. It will call
+// the underlying level method (Debug, Info, Warn, Error, Panic, and Fatal) for
+// the defined levels; the Log method is only called for unknown logging
+// levels.
 //
-// Write will call the underlying level method (Debug, Info, Warn, Error,
-// Panic, and Fatal) for the defined levels; the Log method is only called for
-// unknown logging levels.
+// It MUST be called at most once, since Write will return the *CheckedMessage
+// to an internal pool for potentially immediate re-use; re-using a
+// *CheckedMessage after calling Write() will result in data races or other
+// undefined behavior. An attempt is made to detect and DFatal log any re-use,
+// but such detection is not guaranteed due to race conditions.
 func (m *CheckedMessage) Write(fields ...Field) {
 	if m == nil {
 		return
 	}
 
-	if n := m.used.Inc(); 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))
+	if !m.safeToWrite {
+		// we're living in racy times, so copy what we can out of the pointer
+		// that we have, and at least tell the user something
+		if logger := m.logger; logger != nil {
+			lvl, msg := m.lvl, m.msg
+			logger.DFatal(
+				"Must not call zap.(*CheckedMessage).Write() more than once",
+				Nest("prior", Stringer("level", lvl), String("msg", msg)),
+			)
 		}
 		return
 	}
+	m.safeToWrite = false
 
 	switch m.lvl {
 	case DebugLevel:
@@ -87,6 +101,8 @@ func (m *CheckedMessage) Write(fields ...Field) {
 	}
 
 	m.next.Write(fields...)
+	m.next, m.tail = nil, nil
+	_cmPool.Put(m)
 }
 
 // Chain combines two or more CheckedMessages. If the receiver message is not

checked_message_test.go

@@ -47,17 +47,18 @@ func TestJSONLoggerCheck(t *testing.T) {
 	})
 }
 
-func TestCheckedMessageIsSingleUse(t *testing.T) {
-	expected := []string{
-		`{"level":"info","msg":"single-use"}`,
-		`{"level":"error","msg":"Shouldn't re-use a CheckedMessage.","original":"single-use"}`,
-	}
-	withJSONLogger(t, nil, func(logger Logger, buf *testBuffer) {
-		cm := logger.Check(InfoLevel, "single-use")
-		cm.Write() // ok
-		cm.Write() // first re-use logs error
-		cm.Write() // second re-use is silently ignored
-		assert.Equal(t, expected, buf.Lines(), "Expected re-using a CheckedMessage to log an error.")
+func TestCheckedMessageUnsafeWrite(t *testing.T) {
+	withJSONLogger(t, opts(InfoLevel), func(logger Logger, buf *testBuffer) {
+		stub := stubExit()
+		defer stub.Unstub()
+		cm := logger.Check(InfoLevel, "bob lob law blog")
+		cm.Write()
+		cm.Write()
+		assert.Equal(t, []string{
+			`{"level":"info","msg":"bob lob law blog"}`,
+			`{"level":"error","msg":"Must not call zap.(*CheckedMessage).Write() more than once","prior":{"level":"info","msg":"bob lob law blog"}}`,
+		}, buf.Lines(), "Expected one lob log, and a fatal")
+		stub.AssertNoExit(t)
 	})
 }
 

level.go

@@ -50,9 +50,11 @@ type LevelEnabler interface {
 }
 
 const (
+	invalidLevel Level = iota - 2
+
 	// DebugLevel logs are typically voluminous, and are usually disabled in
 	// production.
-	DebugLevel Level = iota - 1
+	DebugLevel
 	// InfoLevel is the default logging priority.
 	InfoLevel
 	// WarnLevel logs are more important than Info, but don't need individual