open-nomad/client/driver/logrotator/logs.go

package logrotator

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
)

const (
	bufSize = 32 * 1024 // Max number of bytes read from a buffer
)

// LogRotator ingests data and writes out to a rotated set of files
type LogRotator struct {
	MaxFiles int    // maximum number of rotated files retained by the log rotator
	FileSize int64  // maximum file size of a rotated file
	path     string // path where the rotated files are created
	fileName string // base file name of the rotated files

	logFileIdx       int // index to the current file
	oldestLogFileIdx int // index to the oldest log file

	logger  *log.Logger
	purgeCh chan struct{}
}

// NewLogRotator configures and returns a new LogRotator
func NewLogRotator(path string, fileName string, maxFiles int, fileSize int64, logger *log.Logger) (*LogRotator, error) {
	files, err := ioutil.ReadDir(path)
	if err != nil {
		return nil, err
	}

	// Finding out the log file with the largest index
	logFileIdx := 0
	prefix := fmt.Sprintf("%s.", fileName)
	for _, f := range files {
		if strings.HasPrefix(f.Name(), prefix) {
			fileIdx := strings.TrimPrefix(f.Name(), prefix)
			n, err := strconv.Atoi(fileIdx)
			if err != nil {
				continue
			}
			if n > logFileIdx {
				logFileIdx = n
			}
		}
	}

	lr := &LogRotator{
		MaxFiles:   maxFiles,
		FileSize:   fileSize,
		path:       path,
		fileName:   fileName,
		logFileIdx: logFileIdx,
		logger:     logger,
		purgeCh:    make(chan struct{}, 1),
	}
	go lr.PurgeOldFiles()

	return lr, nil
}

// Start reads from a Reader and writes them to files and rotates them when the
// size of the file becomes equal to the max size configured
func (l *LogRotator) Start(r io.Reader) error {
	buf := make([]byte, bufSize)
	for {
		logFileName := filepath.Join(l.path, fmt.Sprintf("%s.%d", l.fileName, l.logFileIdx))
		var fileSize int64
		if f, err := os.Stat(logFileName); err == nil {
			// Skipping the current file if it happens to be a directory
			if f.IsDir() {
				l.logFileIdx += 1
				continue
			}
			fileSize = f.Size()
			// Calculating the remaining capacity of the log file
		}
		f, err := os.OpenFile(logFileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			return err
		}
		l.logger.Printf("[DEBUG] client.logrotator: opened a new file: %s", logFileName)

		// Closing the current log file if it doesn't have any more capacity
		if fileSize >= l.FileSize {
			l.logFileIdx += 1
			f.Close()
			continue
		}

		// Reading from the reader and writing into the current log file as long
		// as it has capacity or the reader closes
		totalWritten := 0
		for {
			if l.FileSize-(fileSize+int64(totalWritten)) < 1 {
				f.Close()
				break
			}
			var nr int
			var err error
			remainingSize := l.FileSize - (int64(totalWritten) + fileSize)
			if remainingSize < bufSize {
				nr, err = r.Read(buf[0:remainingSize])
			} else {
				nr, err = r.Read(buf)
			}
			if err != nil {
				f.Close()
				return err
			}
			nw, err := f.Write(buf[:nr])
			if err != nil {
				f.Close()
				return err
			}
			if nr != nw {
				f.Close()
				return fmt.Errorf("failed to write data read from the reader into file, R: %d W: %d", nr, nw)
			}
			totalWritten += nr
		}
		l.logFileIdx = l.logFileIdx + 1
		// Purge old files if we have more files than MaxFiles
		if l.logFileIdx-l.oldestLogFileIdx >= l.MaxFiles {
			select {
			case l.purgeCh <- struct{}{}:
			default:
			}
		}
	}
}

// PurgeOldFiles removes older files and keeps only the last N files rotated for
// a file
func (l *LogRotator) PurgeOldFiles() {
	for {
		select {
		case <-l.purgeCh:
			var fIndexes []int
			files, err := ioutil.ReadDir(l.path)
			if err != nil {
				return
			}
			// Inserting all the rotated files in a slice
			for _, f := range files {
				if strings.HasPrefix(f.Name(), l.fileName) {
					fileIdx := strings.TrimPrefix(f.Name(), fmt.Sprintf("%s.", l.fileName))
					n, err := strconv.Atoi(fileIdx)
					if err != nil {
						continue
					}
					fIndexes = append(fIndexes, n)
				}
			}

			// Sorting the file indexes so that we can purge the older files and keep
			// only the number of files as configured by the user
			sort.Sort(sort.IntSlice(fIndexes))
			var toDelete []int
			toDelete = fIndexes[0 : len(fIndexes)-l.MaxFiles]
			for _, fIndex := range toDelete {
				fname := filepath.Join(l.path, fmt.Sprintf("%s.%d", l.fileName, fIndex))
				os.RemoveAll(fname)
			}
			l.oldestLogFileIdx = fIndexes[0]
		}
	}
}