mlog: make stopping logic a little better

This commit is contained in:
Brian Picciano 2018-10-28 14:10:11 -04:00
parent 23045168cf
commit d9480abb59

View File

@ -252,7 +252,9 @@ type Logger struct {
msgBufPool *sync.Pool msgBufPool *sync.Pool
msgCh chan msg msgCh chan msg
testMsgWrittenCh chan struct{} // only initialized/used in tests testMsgWrittenCh chan struct{} // only initialized/used in tests
wg sync.WaitGroup
stopCh chan struct{}
wg *sync.WaitGroup
} }
// NewLogger initializes and returns a new instance of Logger which will write // NewLogger initializes and returns a new instance of Logger which will write
@ -268,6 +270,8 @@ func NewLogger(wc io.WriteCloser) *Logger {
}, },
msgCh: make(chan msg, 1024), msgCh: make(chan msg, 1024),
maxLevel: InfoLevel.Uint(), maxLevel: InfoLevel.Uint(),
stopCh: make(chan struct{}),
wg: new(sync.WaitGroup),
} }
l.wg.Add(1) l.wg.Add(1)
go func() { go func() {
@ -282,21 +286,42 @@ func (l *Logger) cp() *Logger {
return &l2 return &l2
} }
func (l *Logger) spin() { func (l *Logger) drain() {
for msg := range l.msgCh { for {
if _, err := io.Copy(l.wc, msg.buf); err != nil { select {
go l.Error("error writing to Logger's WriteCloser", ErrKV(err)) case m := <-l.msgCh:
} l.writeMsg(m)
l.msgBufPool.Put(msg.buf) default:
if l.testMsgWrittenCh != nil { return
l.testMsgWrittenCh <- struct{}{} }
} }
if msg.msg.Level.Uint() == 0 { }
l.wc.Close()
os.Exit(1) func (l *Logger) writeMsg(m msg) {
if _, err := m.buf.WriteTo(l.wc); err != nil {
go l.Error("error writing to Logger's WriteCloser", ErrKV(err))
}
l.msgBufPool.Put(m.buf)
if l.testMsgWrittenCh != nil {
l.testMsgWrittenCh <- struct{}{}
}
if m.msg.Level.Uint() == 0 {
l.wc.Close()
os.Exit(1)
}
}
func (l *Logger) spin() {
defer l.wc.Close()
for {
select {
case m := <-l.msgCh:
l.writeMsg(m)
case <-l.stopCh:
l.drain()
return
} }
} }
l.wc.Close()
} }
// WithMaxLevelUint returns a copy of the Logger with its max logging level set // WithMaxLevelUint returns a copy of the Logger with its max logging level set
@ -338,7 +363,7 @@ func (l *Logger) WithKV(kvs ...KVer) *Logger {
// //
// The Logger should not be used after Stop is called // The Logger should not be used after Stop is called
func (l *Logger) Stop() { func (l *Logger) Stop() {
close(l.msgCh) close(l.stopCh)
l.wg.Wait() l.wg.Wait()
} }
@ -359,7 +384,10 @@ func (l *Logger) Log(lvl Level, msgStr string, kvs ...KVer) {
return return
} }
l.msgCh <- msg{buf: buf, msg: m} select {
case l.msgCh <- msg{buf: buf, msg: m}:
case <-l.stopCh:
}
// if a Fatal is logged then we're merely waiting here for spin to call // 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 // os.Exit, and this go-routine shouldn't be allowed to continue