logfrog-go/handlers/rotatingFile.go

package handlers

import (
	"bufio"
	"compress/gzip"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"codeark.it/Bithero-Agency-Go/logfrog-go"
)

type RotationInterval int

const (
    RotateMinutely RotationInterval = iota
    RotateHourly
    RotateDaily
    RotateWeekly
)

const day = time.Hour * 24
const week = day * 7

const DefaultDailyTimeLayout = "2006-01-02"

// Returns the next timeframe for an interval
func (r RotationInterval) GetNextTimeFrame() time.Time {
    switch r {
    case RotateMinutely:
        return time.Now().Add(time.Minute).Truncate(time.Minute);
    case RotateHourly:
        return time.Now().Add(time.Hour).Truncate(time.Hour);
    case RotateDaily:
        return time.Now().Add(day).Truncate(day);
    case RotateWeekly:
        return time.Now().Add(week).Truncate(week);
    default:
        panic(fmt.Errorf("invalid rotation interval"))
    }
}

type RotatingFileHandler struct {
    *logfrog.BaseLogHandler

    basepath        string
    file            *os.File

    lognameFormat       string  // format for the current logfile
    oldLognameFormat    string  // format for the old logfiles

    rotateInterval      RotationInterval
    nextRotation        time.Time

    compress            bool
    compressionLevel    int

    mu                  sync.Mutex
}

// Creates a new RotatingFileHandler, that rotates daily, keeps at maximum 7 rotated files,
// compresses rotated logs with an level of gzip.DefaultCompression.
//
// The current logfile is named
//
//   logname + ".log"
//
// and rotated logs will be named with the date, i.e.:
//
//   logname + "_2006-01-02.log"
//
func NewDailyRotatingFileHandler(basepath, logname string) (*RotatingFileHandler, error) {
    return NewRotatingFileHandler(
        basepath,
        logname + ".log",
        logname + "_" + DefaultDailyTimeLayout + ".log",
        RotateDaily,
    )
}

// Creates an new RotatingFileHandler with the given attributes. Compresses per default with
// a level of gzip.DefaultCompression.
//
// The both format argumens (lognameFormat and oldLognameFormat) will be put through time.Format,
// so the same rules apply. I.e. the string "2006" will be replaced with the current year.
func NewRotatingFileHandler(
    basepath, lognameFormat, oldLognameFormat string, rotateInterval RotationInterval,
) (*RotatingFileHandler, error) {
    b := logfrog.NewBaseLogHandler()
    r := &RotatingFileHandler{
        BaseLogHandler: b,

        basepath: basepath,
        file: nil,
        lognameFormat: lognameFormat,
        oldLognameFormat: oldLognameFormat,

        rotateInterval: rotateInterval,
        nextRotation: (rotateInterval).GetNextTimeFrame(),

        compress: true, 
        compressionLevel: gzip.DefaultCompression,
    }
    r.LogHandler = r

    logfilename := r.getLogFilename()

    file, err := openOrCreateFile(logfilename)
    if err != nil {
        return nil, err
    }
    r.file = file

    return r, nil
}

// --------------------------------------------------------------------------------

// Enables compression of rotated logs with the specified compression level
func (r *RotatingFileHandler) EnableCompression(level int) *RotatingFileHandler {
    r.compress = true
    r.compressionLevel = level
    return r
}

// Disables compression of rotated logs
func (r *RotatingFileHandler) DisableCompression() *RotatingFileHandler {
    r.compress = false
    return r
}

// Sets the format of the filename for the current rotation
func (r *RotatingFileHandler) SetLognameFormat(lognameFormat string) *RotatingFileHandler {
    r.lognameFormat = lognameFormat
    return r
}

// Sets the format of the filename of rotated logfiles
func (r *RotatingFileHandler) SetOldLognameFormat(oldLognameFormat string) *RotatingFileHandler {
    r.oldLognameFormat = oldLognameFormat
    return r
}

// Returns the timestamp of the next rotation
func (r *RotatingFileHandler) GetNextRotation() time.Time {
    return r.nextRotation
}

// Returns the duration until the next rotation
func (r *RotatingFileHandler) GetUntilNextRotation() time.Duration {
    return time.Until(r.nextRotation)
}

// Returns the interval in which the log is rotated
func (r *RotatingFileHandler) GetRotationInterval() RotationInterval {
    return r.rotateInterval
}

// --------------------------------------------------------------------------------

func (r *RotatingFileHandler) OnClose() error {
    if r.file == nil {
        return nil
    }
    return r.file.Close()
}

// --------------------------------------------------------------------------------

func (r *RotatingFileHandler) getLogFilename() string {
    return filepath.Join(r.basepath, r.nextRotation.Format(r.lognameFormat))
}

func (r *RotatingFileHandler) getRotatedLogFilename() string {
    return filepath.Join(r.basepath, r.nextRotation.Format(r.oldLognameFormat))
}

func (rf *RotatingFileHandler) compressOldLog(rotatedfilename string) error {
    fmt.Printf("begin compressing of %s\n", rotatedfilename);

    fr, err := os.Open(rotatedfilename)
    if err != nil {
        fmt.Printf("[compressing %s] failed to open file: %v\n", rotatedfilename, err);
        return err
    }
    defer fr.Close()

    //fw, err := os.OpenFile(rotatedfilename + ".gz", os.O_WRONLY | os.O_CREATE, 0o666)
    fw, err := os.Create(rotatedfilename + ".gz")
    if err != nil {
        fmt.Printf("[compressing %s] failed to open output file: %v\n", rotatedfilename, err);
        return err
    }
    defer fw.Close()

    br := bufio.NewReader(fr)

    gzw, err := gzip.NewWriterLevel(fw, gzip.DefaultCompression)
    if err != nil {
        fmt.Printf("[compressing %s] failed to create gzip writer: %v\n", rotatedfilename, err);
        return err
    }
    defer gzw.Close()

    _, err = br.WriteTo(gzw)
    if err != nil {
        fmt.Printf("[compressing %s] failed WriteTo: %v\n", rotatedfilename, err);
        gzw.Close()
        fw.Close()
        os.Remove(rotatedfilename + ".gz")
        return err
    }

    fr.Close()
    err = os.Remove(rotatedfilename)
    fmt.Printf("[compressing %s] done with compressing: %v\n", rotatedfilename, err);
    return err
}

func (r *RotatingFileHandler) doRotation() error {
    // first create the filename of the logfile that will be rotated
    oldFilename := r.getLogFilename()

    // then create the filename for the old logfile with the current nextRotation
    rotatedfilename := r.getRotatedLogFilename()

    // update nextRotation to the next rotation
    r.nextRotation = r.rotateInterval.GetNextTimeFrame()

    // create the filename for the current logfile with the current nextRotation
    newFilename := r.getLogFilename()

    // close the old one and set the field additionally to nil
    r.file.Close()
    r.file = nil
    defer func(){
        // open the new logfile for writing
        file, err := openOrCreateFile(newFilename)
        if err != nil {
            panic(err)
        }
        r.file = file
    }()

    // rotate the file if necessary
    if oldFilename != rotatedfilename {
        os.Rename(oldFilename, rotatedfilename)
    }

    // if compression is enabled, compress the old logfile
    if r.compress {
        go r.compressOldLog(rotatedfilename)
    }
    return nil
}

func (r *RotatingFileHandler) Write(buf []byte) error {
    r.mu.Lock()
    defer r.mu.Unlock()

    if time.Now().After(r.nextRotation) {
        r.doRotation()
    }

    _, werr := r.file.Write(buf)
    return werr
}