From 3a1705db158adb256f7af24955d5e44d649e9162 Mon Sep 17 00:00:00 2001
From: Nick Miyake <nmiyake@palantir.com>
Date: Fri, 15 Mar 2019 12:45:38 -0700
Subject: [PATCH] Stop mill goroutine when logger is closed

Fixes #56
---
 lumberjack.go      | 14 ++++++++++++--
 lumberjack_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/lumberjack.go b/lumberjack.go
index 46d97c5..da966f3 100644
--- a/lumberjack.go
+++ b/lumberjack.go
@@ -112,6 +112,7 @@ type Logger struct {
 	mu   sync.Mutex
 
 	millCh    chan bool
+	wg        *sync.WaitGroup
 	startMill sync.Once
 }
 
@@ -161,11 +162,17 @@ func (l *Logger) Write(p []byte) (n int, err error) {
 	return n, err
 }
 
-// Close implements io.Closer, and closes the current logfile.
+// Close implements io.Closer, and closes the current logfile.  Also stops the logger's mill goroutine if it is running.
 func (l *Logger) Close() error {
 	l.mu.Lock()
 	defer l.mu.Unlock()
-	return l.close()
+	err := l.close()
+	if l.millCh != nil {
+		close(l.millCh)
+		l.wg.Wait()
+		l.millCh = nil
+	}
+	return err
 }
 
 // close closes the file if it is open.
@@ -376,6 +383,7 @@ func (l *Logger) millRunOnce() error {
 // millRun runs in a goroutine to manage post-rotation compression and removal
 // of old log files.
 func (l *Logger) millRun() {
+	defer l.wg.Done()
 	for _ = range l.millCh {
 		// what am I going to do, log this?
 		_ = l.millRunOnce()
@@ -386,6 +394,8 @@ func (l *Logger) millRun() {
 // starting the mill goroutine if necessary.
 func (l *Logger) mill() {
 	l.startMill.Do(func() {
+		l.wg = new(sync.WaitGroup)
+		l.wg.Add(1)
 		l.millCh = make(chan bool, 1)
 		go l.millRun()
 	})
diff --git a/lumberjack_test.go b/lumberjack_test.go
index 484ee9d..720edb0 100644
--- a/lumberjack_test.go
+++ b/lumberjack_test.go
@@ -8,6 +8,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"runtime/pprof"
 	"testing"
 	"time"
 
@@ -127,6 +128,30 @@ func TestDefaultFilename(t *testing.T) {
 	existsWithContent(filename, b, t)
 }
 
+func TestGoRoutinesNotLeaked(t *testing.T) {
+	dir := makeTempDir("TestGoRoutinesNotLeaked", t)
+	defer os.RemoveAll(dir)
+
+	numGoRoutinesBefore := pprof.Lookup("goroutine").Count()
+	filename := logFile(dir)
+	for i := 0; i < 25; i++ {
+		func() {
+			l := &Logger{
+				Filename: filename,
+				MaxSize:  10,
+			}
+			defer l.Close()
+			b := []byte("boo!")
+			_, err := l.Write(b)
+			isNil(err, t)
+		}()
+	}
+	numGoRoutinesAfter := pprof.Lookup("goroutine").Count()
+
+	// all loggers have been closed, so number of goroutines should not have increased
+	equals(numGoRoutinesBefore, numGoRoutinesAfter, t)
+}
+
 func TestAutoRotate(t *testing.T) {
 	currentTime = fakeTime
 	megabyte = 1