mlog: refactor to handle certain use cases better, and make KV merging more consistent

Logger is now initialized with its WriteCloser (which can never change
after initialization) and then has its properties modified via method
calls. These method calls return new copies of the Logger, all writing
to the same original WriteCloser still. This allows for doing things
like having a Logger within some entity that has its own KV, for
example.

Also I revamped the merging code a bit, and fixed a precedence
inconsitency in ErrWithKV and CtxWithKV.

mlog/errctx.go

@@ -12,15 +12,18 @@ type kvKey int
 
 // ErrWithKV embeds the merging of a set of KVs into an error, returning a new
 // error instance. If the error already has a KV embedded in it then the
-// returned error will have the merging of them all.
+// returned error will have the merging of them all, with the given KVs taking
+// precedence.
 func ErrWithKV(err error, kvs ...KVer) merry.Error {
 	if err == nil {
 		return nil
 	}
 	merr := merry.WrapSkipping(err, 1)
-	kv := merge(kvs...)
+	var kv KV
 	if exKV := merry.Value(merr, kvKey(0)); exKV != nil {
-		kv = merge(exKV.(KV), kv)
+		kv = mergeInto(exKV.(KV), kvs...)
+	} else {
+		kv = merge(kvs...)
 	}
 	return merr.WithValue(kvKey(0), kv)
 }
@@ -47,12 +50,14 @@ func ErrKV(err error) KVer {
 
 // CtxWithKV embeds a KV into a Context, returning a new Context instance. If
 // the Context already has a KV embedded in it then the returned error will have
-// the merging of the two.
+// the merging of the two, with the given KVs taking precedence.
 func CtxWithKV(ctx context.Context, kvs ...KVer) context.Context {
-	kv := merge(kvs...)
 	existingKV := ctx.Value(kvKey(0))
+	var kv KV
 	if existingKV != nil {
-		kv = merge(existingKV.(KV), kv)
+		kv = mergeInto(existingKV.(KV), kvs...)
+	} else {
+		kv = merge(kvs...)
 	}
 	return context.WithValue(ctx, kvKey(0), kv)
 }

mlog/mlog.go

@@ -15,6 +15,7 @@
 package mlog
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"os"
@@ -124,17 +125,37 @@ func (kv KV) Set(k string, v interface{}) KV {
 	return nkv
 }
 
-// separate from Merge because it's convenient to not return a KVer if that KVer
+// separate from MergeInto because it's convenient to not return a KVer if that
-// is going to immediately have KV called on it (and thereby create a copy for
+// KVer is going to immediately have KV called on it (and thereby create a copy
-// no reason).
+// for no reason).
-func merge(kvs ...KVer) KV {
+//
-	kvm := kvs[0].KV()
+// this may take in any amount of nil values, but should never return nil
-	for _, kv := range kvs[1:] {
+func mergeInto(kv KVer, kvs ...KVer) KV {
+	if kv == nil {
+		kv = KV(nil) // will return empty map when KV is called on it
+	}
+	kvm := kv.KV()
+	for _, kv := range kvs {
+		if kv == nil {
+			continue
+		}
 		for k, v := range kv.KV() {
 			kvm[k] = v
 		}
 	}
-	return KV(kvm)
+	return kvm
+}
+
+// separate from Merge because it's convenient to not return a KVer if that KVer
+// is going to immediately have KV called on it (and thereby create a copy for
+// no reason).
+//
+// this may take in any amount of nil values, but should never return nil
+func merge(kvs ...KVer) KV {
+	if len(kvs) == 0 {
+		return KV{}
+	}
+	return mergeInto(kvs[0], kvs[1:]...)
 }
 
 // Merge takes in multiple KVers and returns a single KVer which is the union of
@@ -144,6 +165,11 @@ func Merge(kvs ...KVer) KVer {
 	return merge(kvs...)
 }
 
+// MergeInto is a convenience function which acts similarly to Merge.
+func MergeInto(kv KVer, kvs ...KVer) KVer {
+	return mergeInto(kv, kvs...)
+}
+
 // Message describes a message to be logged, after having already resolved the
 // KVer
 type Message struct {
@@ -153,8 +179,8 @@ type Message struct {
 }
 
 // WriteFn describes a function which formats a single log message and writes it
-// to the given io.Writer. If the io.Writer returns an error this should return
+// to the given io.Writer. If the io.Writer returns an error WriteFn should
-// that error.
+// return that error.
 type WriteFn func(w io.Writer, msg Message) error
 
 func stringSlice(kv KV) [][2]string {
@@ -190,102 +216,122 @@ func DefaultWriteFn(w io.Writer, msg Message) error {
 	return err
 }
 
+type msg struct {
+	buf *bytes.Buffer
+	msg Message
+}
+
 // Logger wraps a WriteFn and an io.WriteCloser such that logging calls on the
 // Logger will use them (in a thread-safe manner) to write out log messages.
-//
-// Logger will auto-initialize itself on the first method call.
 type Logger struct {
-	WriteFn        // Defaults to DefaultWriteFn
+	wc       io.WriteCloser
-	io.WriteCloser // Defaults to os.Stderr
+	wfn      WriteFn
-
-	levelL      sync.RWMutex
 	maxLevel uint
-	maxLevelSet bool
+	kv       KV
 
-	init             sync.Once
+	msgBufPool       *sync.Pool
-	msgCh            chan Message
+	msgCh            chan msg
 	testMsgWrittenCh chan struct{} // only initialized/used in tests
 }
 
-func (l *Logger) getMaxLevel() uint {
+// NewLogger initializes and returns a new instance of Logger which will write
-	l.levelL.RLock()
+// to the given WriteCloser.
-	defer l.levelL.RUnlock()
+func NewLogger(wc io.WriteCloser) *Logger {
-	if !l.maxLevelSet {
+	l := &Logger{
-		return InfoLevel.Uint()
+		wc:  wc,
+		wfn: DefaultWriteFn,
+		msgBufPool: &sync.Pool{
+			New: func() interface{} {
+				return new(bytes.Buffer)
+			},
+		},
+		msgCh:    make(chan msg, 1024),
+		maxLevel: InfoLevel.Uint(),
 	}
-	return l.maxLevel
+	go l.spin()
+	return l
 }
 
-// SetMaxLevelUint sets the maximum (up-to-and-including) level priority which
+func (l *Logger) cp() *Logger {
-// the Logger will output a log for. Any Levels with an Uint value higher than
+	l2 := *l
-// this will not be logged. This may be called at any time.
+	return &l2
-func (l *Logger) SetMaxLevelUint(i uint) {
-	l.levelL.Lock()
-	l.maxLevel = i
-	l.maxLevelSet = true
-	l.levelL.Unlock()
-}
-
-// SetMaxLevel sets the maximum (up-to-and-including) Level which the Logger
-// will output a log for. Any Levels whose Uint value is higher than this one's
-// will not be logged. This may be called at any time.
-func (l *Logger) SetMaxLevel(lvl Level) {
-	l.SetMaxLevelUint(lvl.Uint())
 }
 
 func (l *Logger) spin() {
 	for msg := range l.msgCh {
-		if err := l.WriteFn(l.WriteCloser, msg); err != nil {
+		if _, err := io.Copy(l.wc, msg.buf); err != nil {
-			go l.Error("could not write to Logger.WriteCloser", ErrKV(err))
+			go l.Error("error writing to Logger's WriteCloser", ErrKV(err))
 		}
+		l.msgBufPool.Put(msg.buf)
 		if l.testMsgWrittenCh != nil {
 			l.testMsgWrittenCh <- struct{}{}
 		}
-		if msg.Level.Uint() == 0 {
+		if msg.msg.Level.Uint() == 0 {
-			l.WriteCloser.Close()
+			l.wc.Close()
 			os.Exit(1)
 		}
 	}
-	l.WriteCloser.Close()
+	l.wc.Close()
 }
 
-func (l *Logger) initDo() {
+// WithMaxLevelUint returns a copy of the Logger with its max logging level set
-	l.init.Do(func() {
+// to the given uint. The Logger will not log any messages with a higher
-		if l.WriteFn == nil {
+// Level.Uint value.
-			l.WriteFn = DefaultWriteFn
+func (l *Logger) WithMaxLevelUint(i uint) *Logger {
+	l = l.cp()
+	l.maxLevel = i
+	return l
 }
-		if l.WriteCloser == nil {
+
-			l.WriteCloser = os.Stderr
+// WithMaxLevel returns a copy of the Logger with its max Level set to the given
+// one. The Logger will not log any messages with a higher Level.Uint value.
+func (l *Logger) WithMaxLevel(lvl Level) *Logger {
+	return l.WithMaxLevelUint(lvl.Uint())
 }
-		// maxLevel's default is implicitly handled by getMaxLevel
+
-		l.msgCh = make(chan Message, 1000)
+// WithWriteFn returns a copy of the Logger which will use the given WriteFn
-		go l.spin()
+// to format and write Messages to the Logger's WriteCloser. This does not
-	})
+// affect the WriteFn of the original Logger, and both can be used at the same
+// time.
+func (l *Logger) WithWriteFn(wfn WriteFn) *Logger {
+	l = l.cp()
+	l.wfn = wfn
+	return l
+}
+
+// WithKV returns a copy of Logger which will implicitly use the KVers for all
+// log messages.
+func (l *Logger) WithKV(kvs ...KVer) *Logger {
+	l = l.cp()
+	l.kv = mergeInto(l.kv, kvs...)
+	return l
 }
 
 // Stop stops and cleans up any running go-routines and resources held by the
 // Logger, allowing it to be garbage-collected. The Logger should not be used
 // after Stop is called
 func (l *Logger) Stop() {
-	l.initDo()
 	close(l.msgCh)
 }
 
 // Log can be used to manually log a message of some custom defined Level. kvs
 // will be Merge'd automatically. If the Level is a fatal (Uint() == 0) then
 // calling this will never return, and the process will have os.Exit(1) called.
-func (l *Logger) Log(lvl Level, msg string, kvs ...KVer) {
+func (l *Logger) Log(lvl Level, msgStr string, kvs ...KVer) {
-	if maxLevel := l.getMaxLevel(); maxLevel < lvl.Uint() {
+	if l.maxLevel < lvl.Uint() {
 		return
 	}
 
-	var kv KV
+	m := Message{Level: lvl, Msg: msgStr, KV: mergeInto(l.kv, kvs...)}
-	if len(kvs) > 0 {
+	buf := l.msgBufPool.Get().(*bytes.Buffer)
-		kv = merge(kvs...)
+	buf.Reset()
+	if err := l.wfn(buf, m); err != nil {
+		// TODO welp, hopefully this doesn't infinite loop
+		l.Log(ErrorLevel, "Logger could not write to WriteCloser", ErrKV(err))
+		return
 	}
 
-	l.initDo()
+	l.msgCh <- msg{buf: buf, msg: m}
-	l.msgCh <- Message{Level: lvl, Msg: msg, KV: kv}
 
 	// if a Fatal is logged then we're merely waiting here for spin to call
 	// os.Exit, and this go-routine shouldn't be allowed to continue
@@ -326,7 +372,7 @@ func (l *Logger) Fatal(msg string, kvs ...KVer) {
 // till the first time it is called any of DefaultLogger's fields may be
 // modified before using one of the Debug, Info, Warn, etc... global functions.
 var (
-	DefaultLogger = new(Logger)
+	DefaultLogger = NewLogger(os.Stderr)
 	Debug         = DefaultLogger.Debug
 	Info          = DefaultLogger.Info
 	Warn          = DefaultLogger.Warn

mlog/mlog_test.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"io/ioutil"
 	"regexp"
+	"strings"
 	. "testing"
 	"time"
 
@@ -34,26 +35,18 @@ func TestKV(t *T) {
 	kv2 = kv.Set("bar", "wat")
 	assert.Equal(t, KV{"foo": "a"}, kv)
 	assert.Equal(t, KV{"foo": "a", "bar": "wat"}, kv2)
-
-	assert.Equal(t, KV{"foo": "a", "bar": "b"}, Merge(
-		KV{"foo": "aaaaa"},
-		KV{"foo": "a", "bar": "bbbbb"},
-		KV{"bar": "b"},
-	))
 }
 
 func TestLLog(t *T) {
 	buf := new(bytes.Buffer)
-	l := &Logger{
+	l := NewLogger(struct {
-		WriteCloser: struct {
 		io.Writer
 		io.Closer
 	}{
 		Writer: buf,
 		Closer: ioutil.NopCloser(nil),
-		},
+	})
-		testMsgWrittenCh: make(chan struct{}, 10),
+	l.testMsgWrittenCh = make(chan struct{}, 10)
-	}
 
 	assertOut := func(expected string) {
 		select {
@@ -75,7 +68,8 @@ func TestLLog(t *T) {
 	assertOut("~ WARN -- baz\n")
 	assertOut("~ ERROR -- buz\n")
 
-	l.SetMaxLevel(WarnLevel)
+	{
+		l := l.WithMaxLevel(WarnLevel)
 		l.Log(DebugLevel, "foo")
 		l.Log(InfoLevel, "bar")
 		l.Log(WarnLevel, "baz")
@@ -84,6 +78,20 @@ func TestLLog(t *T) {
 		assertOut("~ ERROR -- buz -- a=\"b\"\n")
 	}
 
+	{
+		l2 := l.WithWriteFn(func(w io.Writer, msg Message) error {
+			msg.Msg = strings.ToUpper(msg.Msg)
+			return DefaultWriteFn(w, msg)
+		})
+		l2.Log(InfoLevel, "bar")
+		l2.Log(WarnLevel, "baz")
+		l.Log(ErrorLevel, "buz")
+		assertOut("~ INFO -- BAR\n")
+		assertOut("~ WARN -- BAZ\n")
+		assertOut("~ ERROR -- buz\n")
+	}
+}
+
 func TestDefaultWriteFn(t *T) {
 	assertFormat := func(postfix string, msg Message) {
 		expectedRegex := regexp.MustCompile(`^~ ` + postfix + `\n$`)
@@ -106,3 +114,31 @@ func TestDefaultWriteFn(t *T) {
 	msg.KV = KV{"foo": "a", "bar": "b"}.KV()
 	assertFormat("INFO -- this is a test -- bar=\"b\" foo=\"a\"", msg)
 }
+
+func TestMerge(t *T) {
+	assertMerge := func(exp KV, kvs ...KVer) {
+		got := merge(kvs...)
+		assert.Equal(t, exp, got)
+	}
+
+	assertMerge(KV{})
+	assertMerge(KV{}, nil)
+	assertMerge(KV{}, nil, nil)
+
+	assertMerge(KV{"a": "a"}, KV{"a": "a"})
+	assertMerge(KV{"a": "a"}, nil, KV{"a": "a"})
+	assertMerge(KV{"a": "a"}, KV{"a": "a"}, nil)
+
+	assertMerge(
+		KV{"a": "a", "b": "b"},
+		KV{"a": "a"}, KV{"b": "b"},
+	)
+	assertMerge(
+		KV{"a": "a", "b": "b"},
+		KV{"a": "a"}, KV{"b": "b"},
+	)
+	assertMerge(
+		KV{"a": "b"},
+		KV{"a": "a"}, KV{"a": "b"},
+	)
+}