lib/mlog: add method Close to MultiLogger

The Close method flush and close all log forwarders.
Any write to a closed MultiLogger will be ignored.

This changes require adding sync.Mutex to mark if the instance has been
closed or not; which affect createMultiLogger and defaultMLog to return
a pointer to prevent copy on Mutex.

4 files changed, 177 insertions, 63 deletions

diff --git a/lib/mlog/example_test.go b/lib/mlog/example_test.go
index 1a98638b..ed3d3b56 100644
--- a/lib/mlog/example_test.go
+++ b/lib/mlog/example_test.go
@@ -47,9 +47,11 @@ func ExampleMultiLogger() {
 		//mlog.RegisterErrorWriter(NewNamedWriter("slack", slackc))
 	}
 
-	mlog.Outf("writing to standard output and buffer\n")
-	mlog.Errf("writing to standard error and slack\n")
-	mlog.Flush()
+	mlog.Outf("writing to standard output and buffer")
+	mlog.Errf("writing to standard error and slack")
+	mlog.Close()
+	// Try writing to closed mlog.
+	mlog.Outf("writing to standard output and buffer after close")
 	fmt.Println("Output on buffer:", buf.String())
 
 	// Unordered output:
diff --git a/lib/mlog/mlog.go b/lib/mlog/mlog.go
index 6aedb963..a54ca602 100644
--- a/lib/mlog/mlog.go
+++ b/lib/mlog/mlog.go
@@ -37,7 +37,7 @@ const (
 	defTimeFormat = "2006-01-02 15:04:05 MST"
 )
 
-var defaultMLog MultiLogger = createMultiLogger(defTimeFormat, "",
+var defaultMLog *MultiLogger = createMultiLogger(defTimeFormat, "",
 	[]NamedWriter{
 		NewNamedWriter("stdout", os.Stdout),
 	},
@@ -126,5 +126,5 @@ func UnregisterOutputWriter(name string) {
 // A call to Write() on returned io.Writer will forward it to all registered
 // error writers.
 func ErrorWriter() io.Writer {
-	return &defaultMLog
+	return defaultMLog
 }
diff --git a/lib/mlog/mlog_test.go b/lib/mlog/mlog_test.go
new file mode 100644
index 00000000..bc3b6b37
--- /dev/null
+++ b/lib/mlog/mlog_test.go
@@ -0,0 +1,60 @@
+// Copyright 2022, Shulhan <ms@kilabit.info>. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package mlog
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+)
+
+// Test writing to mlog after closed.
+func TestMultiLogger_Close(t *testing.T) {
+	var (
+		bufOut = bytes.Buffer{}
+		bufErr = bytes.Buffer{}
+		outs   = []NamedWriter{
+			NewNamedWriter("bufOut", &bufOut),
+		}
+		errs = []NamedWriter{
+			NewNamedWriter("bufErr", &bufErr),
+		}
+		mlog = NewMultiLogger("", "test", outs, errs)
+
+		outq = make(chan struct{})
+		errq = make(chan struct{})
+	)
+
+	go func() {
+		var x int
+		for x = 0; x < 10; x++ {
+			mlog.Outf("out: %d", x)
+			if x == 2 {
+				outq <- struct{}{}
+				<-outq
+			}
+		}
+	}()
+	go func() {
+		var x int
+		for x = 0; x < 10; x++ {
+			mlog.Errf("err: %d", x)
+			if x == 2 {
+				errq <- struct{}{}
+				<-errq
+			}
+		}
+	}()
+
+	<-outq
+	<-errq
+	mlog.Close()
+	outq <- struct{}{}
+	errq <- struct{}{}
+	mlog.Flush()
+
+	fmt.Println(bufOut.String())
+	fmt.Println(bufErr.String())
+}
diff --git a/lib/mlog/multi_logger.go b/lib/mlog/multi_logger.go
index b616b1a1..9a593431 100644
--- a/lib/mlog/multi_logger.go
+++ b/lib/mlog/multi_logger.go
@@ -29,7 +29,6 @@ type MultiLogger struct {
 	qout chan []byte
 
 	qerrFlush chan bool
-	qflush    chan bool
 	qoutFlush chan bool
 
 	errs map[string]NamedWriter
@@ -38,23 +37,23 @@ type MultiLogger struct {
 	timeFormat string
 
 	prefix []byte
+
+	sync.Mutex
+	isClosed bool
 }
 
 // NewMultiLogger create and initialize new MultiLogger.
 func NewMultiLogger(timeFormat, prefix string, outs, errs []NamedWriter) *MultiLogger {
-	var (
-		mlog = createMultiLogger(timeFormat, prefix, outs, errs)
-	)
-	return &mlog
+	return createMultiLogger(timeFormat, prefix, outs, errs)
 }
 
-func createMultiLogger(timeFormat, prefix string, outs, errs []NamedWriter) (mlog MultiLogger) {
+func createMultiLogger(timeFormat, prefix string, outs, errs []NamedWriter) (mlog *MultiLogger) {
 	var (
 		w    NamedWriter
 		name string
 	)
 
-	mlog = MultiLogger{
+	mlog = &MultiLogger{
 		bufPool: &sync.Pool{
 			New: func() interface{} {
 				return new(bytes.Buffer)
@@ -68,7 +67,6 @@ func createMultiLogger(timeFormat, prefix string, outs, errs []NamedWriter) (mlo
 		qerr:       make(chan []byte, 512),
 		qerrFlush:  make(chan bool, 1),
 		qoutFlush:  make(chan bool, 1),
-		qflush:     make(chan bool, 1),
 	}
 	for _, w = range outs {
 		name = w.Name()
@@ -90,14 +88,48 @@ func createMultiLogger(timeFormat, prefix string, outs, errs []NamedWriter) (mlo
 	return mlog
 }
 
+func flush(qlog chan []byte, writers map[string]NamedWriter) {
+	var (
+		name string
+		err  error
+		nw   NamedWriter
+		b    []byte
+		x    int
+	)
+
+	for x = 0; x < len(qlog); x++ {
+		b = <-qlog
+		if len(b) == 0 {
+			b = append(b, '\n')
+		} else if b[len(b)-1] != '\n' {
+			b = append(b, '\n')
+		}
+		for name, nw = range writers {
+			_, err = nw.Write(b)
+			if err != nil {
+				log.Printf("MultiLogger: %s: %s", name, err)
+			}
+		}
+	}
+}
+
+// Close flush and close all log forwarders.
+// Any write to a closed MultiLogger will be ignored.
+func (mlog *MultiLogger) Close() {
+	mlog.Lock()
+	mlog.isClosed = true
+	close(mlog.qerr)
+	close(mlog.qout)
+	<-mlog.qerrFlush
+	<-mlog.qoutFlush
+	mlog.Unlock()
+}
+
 // Errf write the formatted string and its optional values to all error
 // writers.
 //
 // If the generated string does not end with new line, it will be added.
 func (mlog *MultiLogger) Errf(format string, v ...interface{}) {
-	if len(mlog.errs) == 0 {
-		return
-	}
 	mlog.writeTo(mlog.qerr, format, v...)
 }
 
@@ -110,19 +142,23 @@ func (mlog *MultiLogger) Fatalf(format string, v ...interface{}) {
 
 // Flush all writes and wait until it finished.
 func (mlog *MultiLogger) Flush() {
+	mlog.Lock()
+	if mlog.isClosed {
+		mlog.Unlock()
+		return
+	}
+	mlog.Unlock()
+
 	mlog.qerrFlush <- true
 	mlog.qoutFlush <- true
-	<-mlog.qflush
-	<-mlog.qflush
+	<-mlog.qerrFlush
+	<-mlog.qoutFlush
 }
 
 // Outf write the formatted string and its values to all output writers.
 //
 // If the generated string does not end with new line, it will be added.
 func (mlog *MultiLogger) Outf(format string, v ...interface{}) {
-	if len(mlog.outs) == 0 {
-		return
-	}
 	mlog.writeTo(mlog.qout, format, v...)
 }
 
@@ -147,7 +183,7 @@ func (mlog *MultiLogger) Panicf(format string, v ...interface{}) {
 //		os.Exit(1)
 //	}
 func (mlog *MultiLogger) PrintStack() {
-	mlog.Errf("%s\n", debug.Stack())
+	mlog.Errf("%s", debug.Stack())
 	mlog.Flush()
 }
 
@@ -200,88 +236,99 @@ func (mlog *MultiLogger) Write(b []byte) (n int, err error) {
 
 func (mlog *MultiLogger) processErrorQueue() {
 	var (
-		b   []byte
-		w   NamedWriter
-		err error
-		x   int
+		name string
+		b    []byte
+		w    NamedWriter
+		err  error
+		ok   bool
 	)
 	for {
 		select {
-		case b = <-mlog.qerr:
+		case b, ok = <-mlog.qerr:
+			if !ok {
+				// A closed channel is already empty, no need
+				// to flush it.
+				for name = range mlog.errs {
+					delete(mlog.errs, name)
+				}
+				mlog.qerrFlush <- true
+				return
+			}
+
 			if len(b) == 0 {
 				b = append(b, '\n')
 			} else if b[len(b)-1] != '\n' {
 				b = append(b, '\n')
 			}
-			for _, w = range mlog.errs {
+			for name, w = range mlog.errs {
 				_, err = w.Write(b)
 				if err != nil {
-					log.Printf("MultiLogger: %s: %s", w.Name(), err)
+					log.Printf("MultiLogger: %s: %s", name, err)
 				}
 			}
+
 		case <-mlog.qerrFlush:
-			for x = 0; x < len(mlog.qerr); x++ {
-				b = <-mlog.qerr
-				if len(b) == 0 {
-					b = append(b, '\n')
-				} else if b[len(b)-1] != '\n' {
-					b = append(b, '\n')
-				}
-				for _, w = range mlog.errs {
-					_, err = w.Write(b)
-					if err != nil {
-						log.Printf("MultiLogger: %s: %s", w.Name(), err)
-					}
-				}
-			}
-			mlog.qflush <- true
+			flush(mlog.qerr, mlog.errs)
+			mlog.qerrFlush <- true
 		}
 	}
 }
 
 func (mlog *MultiLogger) processOutputQueue() {
 	var (
-		b   []byte
-		w   NamedWriter
-		err error
-		x   int
+		name string
+		b    []byte
+		w    NamedWriter
+		err  error
+		ok   bool
 	)
 
 	for {
 		select {
-		case b = <-mlog.qout:
+		case b, ok = <-mlog.qout:
+			if !ok {
+				// A closed channel is already empty, no need
+				// to flush it.
+				for name = range mlog.outs {
+					delete(mlog.outs, name)
+				}
+				mlog.qoutFlush <- true
+				return
+			}
+
 			if len(b) == 0 {
 				b = append(b, '\n')
 			} else if b[len(b)-1] != '\n' {
 				b = append(b, '\n')
 			}
-			for _, w = range mlog.outs {
+			for name, w = range mlog.outs {
 				_, err = w.Write(b)
 				if err != nil {
-					log.Printf("MultiLogger: %s: %s", w.Name(), err)
+					log.Printf("MultiLogger: %s: %s", name, err)
 				}
 			}
-		case <-mlog.qoutFlush:
-			for x = 0; x < len(mlog.qout); x++ {
-				b = <-mlog.qout
-				for _, w = range mlog.outs {
-					_, err = w.Write(b)
-					if err != nil {
-						log.Printf("MultiLogger: %s: %s", w.Name(), err)
-					}
 
-				}
-			}
-			mlog.qflush <- true
+		case <-mlog.qoutFlush:
+			flush(mlog.qout, mlog.outs)
+			mlog.qoutFlush <- true
 		}
 	}
 }
 
 func (mlog *MultiLogger) writeTo(q chan []byte, format string, v ...interface{}) {
+	mlog.Lock()
+	if mlog.isClosed {
+		mlog.Unlock()
+		return
+	}
+	mlog.Unlock()
+
 	var (
 		buf    = mlog.bufPool.Get().(*bytes.Buffer)
 		bufFmt = mlog.bufPool.Get().(*bytes.Buffer)
 		args   = make([]interface{}, 0, len(v)+2)
+
+		b []byte
 	)
 	buf.Reset()
 	bufFmt.Reset()
@@ -298,7 +345,12 @@ func (mlog *MultiLogger) writeTo(q chan []byte, format string, v ...interface{})
 	args = append(args, v...)
 	fmt.Fprintf(buf, bufFmt.String(), args...)
 
-	q <- libbytes.Copy(buf.Bytes())
+	b = libbytes.Copy(buf.Bytes())
+	select {
+	case q <- b:
+	default:
+		// Queue is full or closed.
+	}
 
 	mlog.bufPool.Put(bufFmt)
 	mlog.bufPool.Put(buf)