Ensure that the millRun goroutine terminates when Close called.

linux_test.go

@@ -1,12 +1,13 @@
+//go:build linux
 // +build linux
 
 package lumberjack
 
 import (
 	"os"
+	"sync"
 	"syscall"
 	"testing"
-	"time"
 )
 
 func TestMaintainMode(t *testing.T) {
@@ -98,10 +99,11 @@ func TestCompressMaintainMode(t *testing.T) {
 	f.Close()
 
 	l := &Logger{
-		Compress:   true,
-		Filename:   filename,
-		MaxBackups: 1,
-		MaxSize:    100, // megabytes
+		Compress:         true,
+		Filename:         filename,
+		MaxBackups:       1,
+		MaxSize:          100, // megabytes
+		notifyCompressed: make(chan struct{}),
 	}
 	defer l.Close()
 	b := []byte("boo!")
@@ -114,9 +116,7 @@ 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(l.notifyCompressed, t)
 
 	// a compressed version of the log file should now exist with the correct
 	// mode.
@@ -148,10 +148,11 @@ func TestCompressMaintainOwner(t *testing.T) {
 	f.Close()
 
 	l := &Logger{
-		Compress:   true,
-		Filename:   filename,
-		MaxBackups: 1,
-		MaxSize:    100, // megabytes
+		Compress:         true,
+		Filename:         filename,
+		MaxBackups:       1,
+		MaxSize:          100, // megabytes
+		notifyCompressed: make(chan struct{}),
 	}
 	defer l.Close()
 	b := []byte("boo!")
@@ -164,15 +165,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(l.notifyCompressed, 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 +182,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

@@ -3,7 +3,7 @@
 // Note that this is v2.0 of lumberjack, and should be imported using gopkg.in
 // thusly:
 //
-//   import "gopkg.in/natefinch/lumberjack.v2"
+//	import "gopkg.in/natefinch/lumberjack.v2"
 //
 // The package name remains simply lumberjack, and the code resides at
 // https://github.com/natefinch/lumberjack under the v2.0 branch.
@@ -66,7 +66,7 @@ var _ io.WriteCloser = (*Logger)(nil)
 // `/var/log/foo/server.log`, a backup created at 6:30pm on Nov 11 2016 would
 // use the filename `/var/log/foo/server-2016-11-04T18-30-00.000.log`
 //
-// Cleaning Up Old Log Files
+// # Cleaning Up Old Log Files
 //
 // Whenever a new logfile gets created, old log files may be deleted.  The most
 // recent files according to the encoded timestamp will be retained, up to a
@@ -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,20 @@ 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()
+
+	// Always close the mill channel, even if the close fails. This way we
+	// guarantee that the mill goroutine will exit.
+	err := l.close()
+
+	if l.millCh != nil {
+		close(l.millCh)
+		l.wg.Wait()
+		l.millCh = nil
+		l.wg = nil
+	}
+
+	// Return the result of the file close.
+	return err
 }
 
 // close closes the file if it is open.
@@ -356,20 +379,30 @@ 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
 }
 
@@ -385,12 +418,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.wg.Done()
+		}()
+	}
 	select {
-	case l.millCh <- true:
+	case l.millCh <- struct{}{}:
 	default:
 	}
 }