Ensure that the millRun goroutine terminates when Close called.

linux_test.go

@@ -4,9 +4,9 @@ package lumberjack
 
 import (
 	"os"
+	"sync"
 	"syscall"
 	"testing"
-	"time"
 )
 
 func TestMaintainMode(t *testing.T) {
@@ -97,11 +97,13 @@ func TestCompressMaintainMode(t *testing.T) {
 	isNil(err, t)
 	f.Close()
 
+	notify := make(chan struct{})
 	l := &Logger{
-		Compress: true,
-		Filename:   filename,
-		MaxBackups: 1,
-		MaxSize:    100, // megabytes
+		Compress:         true,
+		Filename:         filename,
+		MaxBackups:       1,
+		MaxSize:          100, // megabytes
+		notifyCompressed: notify,
 	}
 	defer l.Close()
 	b := []byte("boo!")
@@ -114,16 +116,14 @@ func TestCompressMaintainMode(t *testing.T) {
 	err = l.Rotate()
 	isNil(err, t)
 
-	// we need to wait a little bit since the files get compressed on a different
-	// goroutine.
-	<-time.After(10 * time.Millisecond)
+	waitForNotify(notify, t)
 
 	// a compressed version of the log file should now exist with the correct
 	// mode.
 	filename2 := backupFile(dir)
 	info, err := os.Stat(filename)
 	isNil(err, t)
-	info2, err := os.Stat(filename2+compressSuffix)
+	info2, err := os.Stat(filename2 + compressSuffix)
 	isNil(err, t)
 	equals(mode, info.Mode(), t)
 	equals(mode, info2.Mode(), t)
@@ -147,11 +147,13 @@ func TestCompressMaintainOwner(t *testing.T) {
 	isNil(err, t)
 	f.Close()
 
+	notify := make(chan struct{})
 	l := &Logger{
-		Compress:   true,
-		Filename:   filename,
-		MaxBackups: 1,
-		MaxSize:    100, // megabytes
+		Compress:         true,
+		Filename:         filename,
+		MaxBackups:       1,
+		MaxSize:          100, // megabytes
+		notifyCompressed: notify,
 	}
 	defer l.Close()
 	b := []byte("boo!")
@@ -164,15 +166,14 @@ func TestCompressMaintainOwner(t *testing.T) {
 	err = l.Rotate()
 	isNil(err, t)
 
-	// we need to wait a little bit since the files get compressed on a different
-	// goroutine.
-	<-time.After(10 * time.Millisecond)
+	waitForNotify(notify, t)
 
 	// a compressed version of the log file should now exist with the correct
 	// owner.
 	filename2 := backupFile(dir)
-	equals(555, fakeFS.files[filename2+compressSuffix].uid, t)
-	equals(666, fakeFS.files[filename2+compressSuffix].gid, t)
+	uid, gid := fakeFS.fileOwners(filename2 + compressSuffix)
+	equals(555, uid, t)
+	equals(666, gid, t)
 }
 
 type fakeFile struct {
@@ -182,18 +183,30 @@ type fakeFile struct {
 
 type fakeFS struct {
 	files map[string]fakeFile
+	mu    sync.Mutex
 }
 
 func newFakeFS() *fakeFS {
 	return &fakeFS{files: make(map[string]fakeFile)}
 }
 
+func (fs *fakeFS) fileOwners(name string) (int, int) {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
+	result := fs.files[name]
+	return result.uid, result.gid
+}
+
 func (fs *fakeFS) Chown(name string, uid, gid int) error {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
 	fs.files[name] = fakeFile{uid: uid, gid: gid}
 	return nil
 }
 
 func (fs *fakeFS) Stat(name string) (os.FileInfo, error) {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
 	info, err := os.Stat(name)
 	if err != nil {
 		return nil, err

lumberjack.go

@@ -111,8 +111,18 @@ type Logger struct {
 	file *os.File
 	mu   sync.Mutex
 
-	millCh    chan bool
-	startMill sync.Once
+	wg     *sync.WaitGroup
+	millCh chan struct{}
+
+	// notifyCompressed is only set and used for tests. It is signalled when
+	// millRunOnce compresses some files. If no files are compressed,
+	// notifyCompressed is not signalled.
+	notifyCompressed chan struct{}
+
+	// notifyRemoved is only set and used for tests. It is signalled when the
+	// millRunOnce method removes some old log files. If no files are removed,
+	// notifyRemoved is not signalled.
+	notifyRemoved chan struct{}
 }
 
 var (
@@ -165,7 +175,16 @@ func (l *Logger) Write(p []byte) (n int, err error) {
 func (l *Logger) Close() error {
 	l.mu.Lock()
 	defer l.mu.Unlock()
-	return l.close()
+	if err := l.close(); err != nil {
+		return err
+	}
+	if l.millCh != nil {
+		close(l.millCh)
+		l.wg.Wait()
+		l.millCh = nil
+		l.wg = nil
+	}
+	return nil
 }
 
 // close closes the file if it is open.
@@ -356,27 +375,37 @@ func (l *Logger) millRunOnce() error {
 		}
 	}
 
+	filesRemoved := false
 	for _, f := range remove {
 		errRemove := os.Remove(filepath.Join(l.dir(), f.Name()))
 		if err == nil && errRemove != nil {
 			err = errRemove
 		}
+		filesRemoved = true
 	}
+	if filesRemoved && l.notifyRemoved != nil {
+		l.notifyRemoved <- struct{}{}
+	}
+
+	filesCompressed := false
 	for _, f := range compress {
 		fn := filepath.Join(l.dir(), f.Name())
 		errCompress := compressLogFile(fn, fn+compressSuffix)
 		if err == nil && errCompress != nil {
 			err = errCompress
 		}
+		filesCompressed = true
+	}
+	if filesCompressed && l.notifyCompressed != nil {
+		l.notifyCompressed <- struct{}{}
 	}
-
 	return err
 }
 
 // millRun runs in a goroutine to manage post-rotation compression and removal
 // of old log files.
-func (l *Logger) millRun() {
-	for _ = range l.millCh {
+func (l *Logger) millRun(ch <-chan struct{}) {
+	for range ch {
 		// what am I going to do, log this?
 		_ = l.millRunOnce()
 	}
@@ -385,12 +414,18 @@ func (l *Logger) millRun() {
 // mill performs post-rotation compression and removal of stale log files,
 // starting the mill goroutine if necessary.
 func (l *Logger) mill() {
-	l.startMill.Do(func() {
-		l.millCh = make(chan bool, 1)
-		go l.millRun()
-	})
+	// It is safe to check the millCh here as we are inside the mutex lock.
+	if l.millCh == nil {
+		l.millCh = make(chan struct{}, 1)
+		l.wg = &sync.WaitGroup{}
+		l.wg.Add(1)
+		go func() {
+			l.millRun(l.millCh)
+			l.wg.Done()
+		}()
+	}
 	select {
-	case l.millCh <- true:
+	case l.millCh <- struct{}{}:
 	default:
 	}
 }