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