log/log.go

   1 package log
   2
   3 import (
   4         "fmt"
   5         "io/ioutil"
   6         "os"
   7         "path/filepath"
   8         "strings"
   9         "sync"
  10         "time"
  11
  12         rotatelogs "github.com/lestrrat-go/file-rotatelogs"
  13         "github.com/sirupsen/logrus"
  14
  15         "github.com/vapor/config"
  16 )
  17
  18 const (
  19         rotationTime int64 = 86400
  20         maxAge       int64 = 604800
  21 )
  22
  23 var defaultFormatter = &logrus.TextFormatter{DisableColors: true}
  24
  25 func InitLogFile(config *config.Config) error {
  26         logPath := config.LogDir()
  27         if err := clearLockFiles(logPath); err != nil {
  28                 return err
  29         }
  30
  31         hook := newBtmHook(logPath)
  32         logrus.AddHook(hook)
  33         logrus.SetOutput(ioutil.Discard) //控制台不输出
  34         fmt.Printf("all logs are output in the %s directory\n", logPath)
  35         return nil
  36 }
  37
  38 type BtmHook struct {
  39         logPath string
  40         lock    *sync.Mutex
  41 }
  42
  43 func newBtmHook(logPath string) *BtmHook {
  44         hook := &BtmHook{lock: new(sync.Mutex)}
  45         hook.logPath = logPath
  46         return hook
  47 }
  48
  49 // Write a log line to an io.Writer.
  50 func (hook *BtmHook) ioWrite(entry *logrus.Entry) error {
  51         module := "general"
  52         if data, ok := entry.Data["module"]; ok {
  53                 module = data.(string)
  54         }
  55
  56         logPath := filepath.Join(hook.logPath, module)
  57         writer, err := rotatelogs.New(
  58                 logPath+".%Y%m%d",
  59                 rotatelogs.WithMaxAge(time.Duration(maxAge)*time.Second),
  60                 rotatelogs.WithRotationTime(time.Duration(rotationTime)*time.Second),
  61         )
  62         if err != nil {
  63                 return err
  64         }
  65
  66         msg, err := defaultFormatter.Format(entry)
  67         if err != nil {
  68                 return err
  69         }
  70
  71         if _, err = writer.Write(msg); err != nil {
  72                 return err
  73         }
  74
  75         return writer.Close()
  76 }
  77
  78 func clearLockFiles(logPath string) error {
  79         files, err := ioutil.ReadDir(logPath)
  80         if os.IsNotExist(err) {
  81                 return nil
  82         } else if err != nil {
  83                 return err
  84         }
  85
  86         for _, file := range files {
  87                 if ok := strings.HasSuffix(file.Name(), "_lock"); ok {
  88                         if err := os.Remove(filepath.Join(logPath, file.Name())); err != nil {
  89                                 return err
  90                         }
  91                 }
  92         }
  93         return nil
  94 }
  95
  96 func (hook *BtmHook) Fire(entry *logrus.Entry) error {
  97         hook.lock.Lock()
  98         defer hook.lock.Unlock()
  99         return hook.ioWrite(entry)
 100 }
 101
 102 // Levels returns configured log levels.
 103 func (hook *BtmHook) Levels() []logrus.Level {
 104         return logrus.AllLevels
 105 }