Repository: uber-go/zap
File: checked_message.go (Go)

Code before:
// 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"

	"github.com/uber-go/atomic"
)

var cmPool sync.Pool

// 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

	// singly linked list built by Chain
	next *CheckedMessage // carried by each part of Chain-ed list
	tail *CheckedMessage // non-nil only in the head of a Chain-ed list
	// NOTE: suffixes are NOT valid lists, since they lack a tail pointer; tail
	// pointer exists solely as an optimization for calling head.Chain.
}

// 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 {
	for {
		val := cmPool.Get()
		if val == nil {
			break
		}
		m := val.(*CheckedMessage)
		if n := m.used.Load(); n > 1 {
			// we've already logged (and survived!) a DFatal log in the second
			// CheckedMessage.Write, so just skip it and move on
			continue
		}
		*m = CheckedMessage{
			logger: logger,
			lvl:    lvl,
			msg:    msg,
		}
		return m
	}

	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.
//
// 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.
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))
		}
		return
	}

	switch m.lvl {
	case DebugLevel:
		m.logger.Debug(m.msg, fields...)
	case InfoLevel:
		m.logger.Info(m.msg, fields...)
	case WarnLevel:
		m.logger.Warn(m.msg, fields...)
	case ErrorLevel:
		m.logger.Error(m.msg, fields...)
	case PanicLevel:
		m.logger.Panic(m.msg, fields...)
	case FatalLevel:
		m.logger.Fatal(m.msg, fields...)
	default:
		m.logger.Log(m.lvl, m.msg, fields...)
	}

	m.next.Write(fields...)
	m.lvl, m.next, m.tail = invalidLevel, nil, nil
	cmPool.Put(m)
}

// Chain combines two or more CheckedMessages. If the receiver message is not
// OK(), the passed message is returned. Otherwise if the passed message is
// OK(), then it is retained such that its Write() will be called after the
// receiver's Write(), and any previously Chain()'ed messages already so
// retained.
func (m *CheckedMessage) Chain(ms ...*CheckedMessage) *CheckedMessage {
	for _, m2 := range ms {
		if !m2.OK() {
			continue
		}
		if m.OK() {
			m.push(m2)
		} else {
			m = m2
		}
	}
	return m
}

func (m *CheckedMessage) push(next *CheckedMessage) {
	if m.tail != nil {
		m.tail.next = next
	} else if m.next != nil {
		m.logger.DFatal("invalid CheckedMessage linked list; did we lose our head?", String("original", m.msg))
	} else {
		m.next = next
	}
	if next.tail != nil {
		m.tail, next.tail = next.tail, nil
	} else {
		m.tail = next
	}
}

// OK checks whether it's safe to call Write.
func (m *CheckedMessage) OK() bool {
	return m != nil
}


Review comment (suggestion):
I'm not sure that `used` is providing us the ironclad guarantees it used to - it's entirely possible to use a checked message once, keep a reference to it, and concurrently pull it out of the pool in another location. The second reference passes this test, since we've only used the message once, and then we hit line 61 and both callers see the new message. Now the original call site (which is misbehaving and re-using the checked message) is polluted with the message and level of the new call site. Both call sites have references to the same checked message, with a fresh atomic, so it's a race to see who gets to use it.

I don't know of a good pattern to avoid this sort of problem with pooling - it's inherently fiddly and error-prone, which is why we avoided it in the current implementation. I'm inclined to remove `used` altogether and make the documentation for `Write` and `Check` much, much more emphatic. (If you can think of a way to preserve a guarantee of safety, I'm all ears.)

Code after:
// 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"

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
	safeToWrite bool
	lvl         Level
	msg         string

	// singly linked list built by Chain
	next *CheckedMessage // carried by each part of Chain-ed list
	tail *CheckedMessage // non-nil only in the head of a Chain-ed list
	// NOTE: suffixes are NOT valid lists, since they lack a tail pointer; tail
	// pointer exists solely as an optimization for calling head.Chain.
}

// 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 {
	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 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 !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:
		m.logger.Debug(m.msg, fields...)
	case InfoLevel:
		m.logger.Info(m.msg, fields...)
	case WarnLevel:
		m.logger.Warn(m.msg, fields...)
	case ErrorLevel:
		m.logger.Error(m.msg, fields...)
	case PanicLevel:
		m.logger.Panic(m.msg, fields...)
	case FatalLevel:
		m.logger.Fatal(m.msg, fields...)
	default:
		m.logger.Log(m.lvl, m.msg, fields...)
	}

	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
// OK(), the passed message is returned. Otherwise if the passed message is
// OK(), then it is retained such that its Write() will be called after the
// receiver's Write(), and any previously Chain()'ed messages already so
// retained.
func (m *CheckedMessage) Chain(ms ...*CheckedMessage) *CheckedMessage {
	for _, m2 := range ms {
		if !m2.OK() {
			continue
		}
		if m.OK() {
			m.push(m2)
		} else {
			m = m2
		}
	}
	return m
}

func (m *CheckedMessage) push(next *CheckedMessage) {
	if m.tail != nil {
		m.tail.next = next
	} else if m.next != nil {
		m.logger.DFatal("invalid CheckedMessage linked list; did we lose our head?", String("original", m.msg))
	} else {
		m.next = next
	}
	if next.tail != nil {
		m.tail, next.tail = next.tail, nil
	} else {
		m.tail = next
	}
}

// OK checks whether it's safe to call Write.
func (m *CheckedMessage) OK() bool {
	return m != nil
}