github.com/v2pro/plz@v0.0.0-20221028024117-e5f9aec5b631/countlog/output/rotation/rotation.go (about)

     1  package rotation
     2  
     3  import (
     4  	"os"
     5  	"time"
     6  	"path"
     7  	"github.com/v2pro/plz/concurrent"
     8  	"context"
     9  	"unsafe"
    10  	"sync/atomic"
    11  	"github.com/v2pro/plz/countlog/spi"
    12  	"math/rand"
    13  	"github.com/v2pro/plz/countlog/loglog"
    14  )
    15  
    16  // normal => triggered => opened new => normal
    17  const statusNormal = 0
    18  const statusTriggered = 1
    19  const statusArchived = 2
    20  
    21  type Writer struct {
    22  	cfg *Config
    23  	// file is owned by the write goroutine
    24  	file *os.File
    25  	// newFile and status shared between write and rotate goroutine
    26  	newFile     unsafe.Pointer
    27  	status      int32
    28  	stat        interface{}
    29  	executor    *concurrent.UnboundedExecutor
    30  	archiveChan chan struct{}
    31  }
    32  
    33  type NamingStrategy interface {
    34  	ListFiles() ([]Archive, error)
    35  	NextFile() (string, error)
    36  }
    37  
    38  type ArchiveStrategy interface {
    39  	Archive(path string) ([]Archive, error)
    40  }
    41  
    42  type Compressor interface {
    43  }
    44  
    45  type ArchiveByCompression struct {
    46  	RawArchive ArchiveStrategy
    47  	Retention  RetainStrategy
    48  	Naming     NamingStrategy
    49  	Compressor Compressor
    50  }
    51  
    52  type Archive struct {
    53  	Path       string
    54  	ArchivedAt time.Time
    55  	Size       int64
    56  }
    57  
    58  type RetainStrategy interface {
    59  	PurgeSet(archives []Archive) []Archive
    60  }
    61  
    62  type TriggerStrategy interface {
    63  	UpdateStat(stat interface{}, file *os.File, buf []byte) (interface{}, bool, error)
    64  	TimeToTrigger() time.Duration
    65  }
    66  
    67  type PurgeStrategy interface {
    68  	Purge(purgeSet []Archive) error
    69  }
    70  
    71  type Config struct {
    72  	WritePath       string
    73  	FileMode        os.FileMode
    74  	DirectoryMode   os.FileMode
    75  	TriggerStrategy TriggerStrategy
    76  	ArchiveStrategy ArchiveStrategy
    77  	RetainStrategy  RetainStrategy
    78  	PurgeStrategy   PurgeStrategy
    79  }
    80  
    81  func NewWriter(cfg Config) (*Writer, error) {
    82  	if cfg.FileMode == 0 {
    83  		cfg.FileMode = 0644
    84  	}
    85  	if cfg.DirectoryMode == 0 {
    86  		cfg.DirectoryMode = 0755
    87  	}
    88  	executor := concurrent.NewUnboundedExecutor()
    89  	writer := &Writer{executor: executor, cfg: &cfg}
    90  	err := writer.reopen()
    91  	if err != nil {
    92  		return nil, err
    93  	}
    94  	executor.Go(writer.rotateInBackground)
    95  	return writer, nil
    96  }
    97  
    98  func (writer *Writer) reopen() error {
    99  	cfg := writer.cfg
   100  	file, err := os.OpenFile(cfg.WritePath, os.O_WRONLY|os.O_APPEND, cfg.FileMode)
   101  	if err != nil {
   102  		if !os.IsNotExist(err) {
   103  			return err
   104  		}
   105  		os.MkdirAll(path.Dir(cfg.WritePath), cfg.DirectoryMode)
   106  		file, err = os.OpenFile(cfg.WritePath,
   107  			os.O_CREATE|os.O_WRONLY|os.O_TRUNC, cfg.FileMode)
   108  		if err != nil {
   109  			return err
   110  		}
   111  	}
   112  	writer.file = file
   113  	return nil
   114  }
   115  
   116  func (writer *Writer) Write(buf []byte) (int, error) {
   117  	if atomic.LoadInt32(&writer.status) == statusArchived {
   118  		err := writer.file.Close()
   119  		if err != nil {
   120  			loglog.Error(err)
   121  		}
   122  		writer.reopen()
   123  	}
   124  	file := writer.file
   125  	triggerStrategy := writer.cfg.TriggerStrategy
   126  	n, err := file.Write(buf)
   127  	if atomic.LoadInt32(&writer.status) == statusNormal {
   128  		var triggered bool
   129  		var err error
   130  		writer.stat, triggered, err = triggerStrategy.UpdateStat(writer.stat, file, buf[:n])
   131  		if err != nil {
   132  			loglog.Error(err)
   133  			return n, err
   134  		}
   135  		if triggered {
   136  			atomic.StoreInt32(&writer.status, statusTriggered)
   137  		}
   138  	}
   139  	return n, err
   140  }
   141  
   142  func (writer *Writer) Close() error {
   143  	writer.executor.StopAndWaitForever()
   144  	return writer.file.Close()
   145  }
   146  
   147  func (writer *Writer) rotateInBackground(ctx context.Context) {
   148  	triggerStrategy := writer.cfg.TriggerStrategy
   149  	archiveStrategy := writer.cfg.ArchiveStrategy
   150  	retainStrategy := writer.cfg.RetainStrategy
   151  	purgeStrategy := writer.cfg.PurgeStrategy
   152  	var timer <-chan time.Time
   153  	for {
   154  		duration := triggerStrategy.TimeToTrigger()
   155  		if duration > 0 {
   156  			duration += time.Duration(rand.Int63n(int64(duration)))
   157  			timer = time.NewTimer(duration).C
   158  		}
   159  		select {
   160  		case <-ctx.Done():
   161  			return
   162  		case <-writer.archiveChan:
   163  		case <-timer:
   164  		}
   165  		archives, err := archiveStrategy.Archive(writer.cfg.WritePath)
   166  		if err != nil {
   167  			loglog.Error(err)
   168  			// retry after one minute
   169  			timer = time.NewTimer(time.Minute).C
   170  			continue
   171  		}
   172  		atomic.StoreInt32(&writer.status, statusArchived)
   173  		purgeSet := retainStrategy.PurgeSet(archives)
   174  		err = purgeStrategy.Purge(purgeSet)
   175  		if err != nil {
   176  			loglog.Error(err)
   177  		}
   178  	}
   179  }