From e38adbbdc64f8abd7933ca83946e12ff1fca4a88 Mon Sep 17 00:00:00 2001
From: Adieu <adieu@adieu.me>
Date: Thu, 23 Jan 2020 16:42:50 +0800
Subject: [PATCH] properly close the writer and wait for all writes to finish

---
 writer.go | 43 +++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/writer.go b/writer.go
index 72e8e3a..a449795 100644
--- a/writer.go
+++ b/writer.go
@@ -47,14 +47,52 @@ func (entry *Entry) WriterLevel(level Level) *io.PipeWriter {
 	default:
 		printFunc = entry.Print
 	}
+	done := make(chan struct{})
 
-	go entry.writerScanner(reader, printFunc)
+	go entry.writerScanner(reader, printFunc, done)
 	runtime.SetFinalizer(writer, writerFinalizer)
 
 	return writer
 }
 
-func (entry *Entry) writerScanner(reader *io.PipeReader, printFunc func(args ...interface{})) {
+func (entry *Entry) WriterLevelWithClose(level Level) (*io.PipeWriter, func()) {
+	reader, writer := io.Pipe()
+
+	var printFunc func(args ...interface{})
+
+	switch level {
+	case TraceLevel:
+		printFunc = entry.Trace
+	case DebugLevel:
+		printFunc = entry.Debug
+	case InfoLevel:
+		printFunc = entry.Info
+	case WarnLevel:
+		printFunc = entry.Warn
+	case ErrorLevel:
+		printFunc = entry.Error
+	case FatalLevel:
+		printFunc = entry.Fatal
+	case PanicLevel:
+		printFunc = entry.Panic
+	default:
+		printFunc = entry.Print
+	}
+
+	done := make(chan struct{})
+
+	go entry.writerScanner(reader, printFunc, done)
+	runtime.SetFinalizer(writer, writerFinalizer)
+
+	closeFunc := func() {
+		writer.Close()
+		<-done
+	}
+
+	return writer, closeFunc
+}
+
+func (entry *Entry) writerScanner(reader *io.PipeReader, printFunc func(args ...interface{}), done chan struct{}) {
 	scanner := bufio.NewScanner(reader)
 	for scanner.Scan() {
 		printFunc(scanner.Text())
@@ -63,6 +101,7 @@ func (entry *Entry) writerScanner(reader *io.PipeReader, printFunc func(args ...
 		entry.Errorf("Error while reading from Writer: %s", err)
 	}
 	reader.Close()
+	close(done)
 }
 
 func writerFinalizer(writer *io.PipeWriter) {