crawlab/core/task/log/file_driver.go

package log

import (
	"bufio"
	"bytes"
	"errors"
	"github.com/apex/log"
	"github.com/crawlab-team/crawlab/core/utils"
	"github.com/crawlab-team/crawlab/trace"
	"github.com/spf13/viper"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

type FileLogDriver struct {
	// settings
	logFileName string
	rootPath    string

	// internals
	mu sync.Mutex
}

func (d *FileLogDriver) Init() {
	go d.cleanup()
}

func (d *FileLogDriver) Close() (err error) {
	return nil
}

func (d *FileLogDriver) WriteLine(id string, line string) (err error) {
	d.initDir(id)

	d.mu.Lock()
	defer d.mu.Unlock()
	filePath := d.getLogFilePath(id, d.logFileName)

	f, err := os.OpenFile(filePath, os.O_WRONLY|os.O_APPEND|os.O_CREATE, os.FileMode(0760))
	if err != nil {
		return trace.TraceError(err)
	}
	defer func(f *os.File) {
		err := f.Close()
		if err != nil {
			log.Errorf("close file error: %s", err.Error())
		}
	}(f)

	_, err = f.WriteString(line + "\n")
	if err != nil {
		return trace.TraceError(err)
	}

	return nil
}

func (d *FileLogDriver) WriteLines(id string, lines []string) (err error) {
	linesString := strings.Join(lines, "\n")
	if err := d.WriteLine(id, linesString); err != nil {
		return err
	}
	return nil
}

func (d *FileLogDriver) Find(id string, pattern string, skip int, limit int) (lines []string, err error) {
	if pattern != "" {
		return lines, errors.New("not implemented")
	}
	if !utils.Exists(d.getLogFilePath(id, d.logFileName)) {
		return nil, nil
	}

	f, err := os.Open(d.getLogFilePath(id, d.logFileName))
	if err != nil {
		return nil, trace.TraceError(err)
	}
	defer f.Close()

	sc := bufio.NewReaderSize(f, 1024*1024*10)

	i := -1
	for {
		line, err := sc.ReadString(byte('\n'))
		if err != nil {
			break
		}
		line = strings.TrimSuffix(line, "\n")

		i++

		if i < skip {
			continue
		}

		if i >= skip+limit {
			break
		}

		lines = append(lines, line)
	}

	return lines, nil
}

func (d *FileLogDriver) Count(id string, pattern string) (n int, err error) {
	if pattern != "" {
		return n, errors.New("not implemented")
	}
	if !utils.Exists(d.getLogFilePath(id, d.logFileName)) {
		return 0, nil
	}

	f, err := os.Open(d.getLogFilePath(id, d.logFileName))
	if err != nil {
		return n, trace.TraceError(err)
	}
	return d.lineCounter(f)
}

func (d *FileLogDriver) Flush() (err error) {
	return nil
}

func (d *FileLogDriver) getBasePath(id string) (filePath string) {
	return filepath.Join(utils.GetTaskLogPath(), id)
}

func (d *FileLogDriver) getMetadataPath(id string) (filePath string) {
	return filepath.Join(d.getBasePath(id), MetadataName)
}

func (d *FileLogDriver) getLogFilePath(id, fileName string) (filePath string) {
	return filepath.Join(d.getBasePath(id), fileName)
}

func (d *FileLogDriver) getLogFiles(id string) (files []os.FileInfo) {
	files, err := utils.ListDir(d.getBasePath(id))
	if err != nil {
		log.Errorf("failed to list log files: %s", err.Error())
		return nil
	}
	return
}

func (d *FileLogDriver) initDir(id string) {
	if !utils.Exists(d.getBasePath(id)) {
		if err := os.MkdirAll(d.getBasePath(id), os.FileMode(0770)); err != nil {
			trace.PrintError(err)
		}
	}
}

func (d *FileLogDriver) lineCounter(r io.Reader) (n int, err error) {
	buf := make([]byte, 32*1024)
	count := 0
	lineSep := []byte{'\n'}

	for {
		c, err := r.Read(buf)
		count += bytes.Count(buf[:c], lineSep)

		switch {
		case err == io.EOF:
			return count, nil

		case err != nil:
			return count, err
		}
	}
}

func (d *FileLogDriver) getTtl() time.Duration {
	ttl := viper.GetString("log.ttl")
	if ttl == "" {
		return DefaultLogTtl
	}

	if strings.HasSuffix(ttl, "s") {
		ttl = strings.TrimSuffix(ttl, "s")
		n, err := strconv.Atoi(ttl)
		if err != nil {
			return DefaultLogTtl
		}
		return time.Duration(n) * time.Second
	} else if strings.HasSuffix(ttl, "m") {
		ttl = strings.TrimSuffix(ttl, "m")
		n, err := strconv.Atoi(ttl)
		if err != nil {
			return DefaultLogTtl
		}
		return time.Duration(n) * time.Minute
	} else if strings.HasSuffix(ttl, "h") {
		ttl = strings.TrimSuffix(ttl, "h")
		n, err := strconv.Atoi(ttl)
		if err != nil {
			return DefaultLogTtl
		}
		return time.Duration(n) * time.Hour

	} else if strings.HasSuffix(ttl, "d") {
		ttl = strings.TrimSuffix(ttl, "d")
		n, err := strconv.Atoi(ttl)
		if err != nil {
			return DefaultLogTtl
		}
		return time.Duration(n) * 24 * time.Hour
	} else {
		return DefaultLogTtl
	}
}

func (d *FileLogDriver) cleanup() {
	// check if log path is set
	if utils.GetTaskLogPath() == "" {
		log.Errorf("log path is not set")
		return
	}

	// check if log path exists
	if !utils.Exists(utils.GetTaskLogPath()) {
		// create log directory if not exists
		if err := os.MkdirAll(utils.GetTaskLogPath(), os.FileMode(0770)); err != nil {
			log.Errorf("failed to create log directory: %s", utils.GetTaskLogPath())
			return
		}
	}

	ticker := time.NewTicker(10 * time.Minute)

	for {
		select {
		case <-ticker.C:
			dirs, err := utils.ListDir(utils.GetTaskLogPath())
			if err != nil {
				log.Errorf("failed to list log directory: %s", utils.GetTaskLogPath())
				continue
			}
			for _, dir := range dirs {
				if time.Now().After(dir.ModTime().Add(d.getTtl())) {
					if err := os.RemoveAll(d.getBasePath(dir.Name())); err != nil {
						log.Errorf("failed to remove outdated log directory: %s", d.getBasePath(dir.Name()))
						continue
					}
					log.Infof("removed outdated log directory: %s", d.getBasePath(dir.Name()))
				}
			}
		}
	}
}

func newFileLogDriver() Driver {
	// driver
	driver := &FileLogDriver{
		logFileName: "log.txt",
		mu:          sync.Mutex{},
	}

	// init
	driver.Init()

	return driver
}

var logDriver Driver
var logDriverOnce sync.Once

func GetFileLogDriver() Driver {
	logDriverOnce.Do(func() {
		logDriver = newFileLogDriver()
	})
	return logDriver
}