youtube_collection/youtube_collection_server/Utils/LogUtil.go

package utils

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path"
	"time"

	"github.com/gin-gonic/gin"
	rotatelogs "github.com/lestrrat-go/file-rotatelogs"
	"github.com/logrusorgru/aurora"
	"github.com/rifflock/lfshook"
	"github.com/sirupsen/logrus"
)

var Logger *logrus.Logger
var isInit bool

func CheckInitLog() bool {
	return isInit
}
func init() {
	isInit = false
}

const (
	maxAgeHour   = 168
	rotationHour = 24
)

type SaveFormatter struct {
}

func (m *SaveFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	var b *bytes.Buffer
	if entry.Buffer != nil {
		b = entry.Buffer
	} else {
		b = &bytes.Buffer{}
	}
	timestamp := entry.Time.Format("2006-01-02 15:04:05")
	var newLog string
	newLog = fmt.Sprintf("[%s] [%s:%d] [%s] %s \n", timestamp, path.Base(entry.Caller.File), entry.Caller.Line, entry.Level, entry.Message)
	b.WriteString(newLog)
	return b.Bytes(), nil
}

type ConsoleFormatter struct {
}

func (m *ConsoleFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	var b *bytes.Buffer
	if entry.Buffer != nil {
		b = entry.Buffer
	} else {
		b = &bytes.Buffer{}
	}
	timestamp := entry.Time.Format("2006-01-02 15:04:05")
	var newLog string
	var levelString aurora.Value
	switch entry.Level.String() {
	case "info":
		levelString = aurora.Green(entry.Level)
	case "warning":
		levelString = aurora.Yellow(entry.Level)
	case "debug":
		levelString = aurora.Gray(16-1, entry.Level)
	case "error":
		levelString = aurora.Red(entry.Level)
	case "fatal":
		levelString = aurora.Red(entry.Level)
	case "panic":
		levelString = aurora.Red(entry.Level)
	}
	newLog = fmt.Sprintf("[%s] [%s:%d] [%s] %s \n", timestamp, path.Base(entry.Caller.File), entry.Caller.Line, levelString, entry.Message)
	b.WriteString(newLog)
	return b.Bytes(), nil
}

func InitLogger(env string, logPath string) *logrus.Logger {
	logFilePath := ""
	errFilePath := ""
	if dir, err := os.Getwd(); err == nil {
		logFilePath = dir + "/logs/" + logPath + "/all.log"
		errFilePath = dir + "/logs/" + logPath + "/error.log"
	}

	accessWriter, err := rotatelogs.New(
		logFilePath+".%Y-%m-%d",
		rotatelogs.WithLinkName(logFilePath),
		rotatelogs.WithRotationTime(time.Hour*rotationHour),
		rotatelogs.WithMaxAge(time.Hour*maxAgeHour),
	)
	if err != nil {
		panic(err)
	}
	errorWriter, err := rotatelogs.New(
		errFilePath+".%Y-%m-%d",
		rotatelogs.WithLinkName(errFilePath),
		rotatelogs.WithRotationTime(time.Hour*rotationHour),
		rotatelogs.WithMaxAge(time.Hour*maxAgeHour),
	)
	if err != nil {
		panic(err)
	}
	allWriter := io.MultiWriter(accessWriter, errorWriter)
	//实例化
	logger := logrus.New()
	logger.Out = os.Stdout
	//设置日志格式
	logger.SetFormatter(&ConsoleFormatter{})
	logger.SetReportCaller(true)
	//设置输出
	if env != "dev" {
		logger.AddHook(lfshook.NewHook(
			lfshook.WriterMap{
				logrus.DebugLevel: accessWriter,
				logrus.InfoLevel:  accessWriter,
				logrus.ErrorLevel: allWriter,
				logrus.PanicLevel: allWriter,
			},
			&SaveFormatter{},
		))
	}
	//设置日志级别
	logger.SetLevel(logrus.DebugLevel)
	Logger = logger
	isInit = true
	return logger
}

func InitReqLogger(env string, logPath string) *logrus.Logger {
	logFilePath := ""
	errFilePath := ""
	if dir, err := os.Getwd(); err == nil {
		logFilePath = dir + "/logs/" + logPath + "/req.log"
		errFilePath = dir + "/logs/" + logPath + "/req-error.log"
	}

	accessWriter, err := rotatelogs.New(
		logFilePath+".%Y-%m-%d",
		rotatelogs.WithLinkName(logFilePath),
		rotatelogs.WithRotationTime(time.Hour*rotationHour),
		rotatelogs.WithMaxAge(time.Hour*maxAgeHour),
	)
	if err != nil {
		panic(err)
	}
	errorWriter, err := rotatelogs.New(
		errFilePath+".%Y-%m-%d",
		rotatelogs.WithLinkName(errFilePath),
		rotatelogs.WithRotationTime(time.Hour*rotationHour),
		rotatelogs.WithMaxAge(time.Hour*maxAgeHour),
	)
	if err != nil {
		panic(err)
	}
	allWriter := io.MultiWriter(accessWriter, errorWriter)
	//实例化
	logger := logrus.New()
	logger.Out = os.Stdout
	//设置日志格式
	logger.SetFormatter(&ConsoleFormatter{})
	logger.SetReportCaller(true)
	//设置输出
	if env != "dev" {
		logger.AddHook(lfshook.NewHook(
			lfshook.WriterMap{
				logrus.DebugLevel: accessWriter,
				logrus.InfoLevel:  accessWriter,
				logrus.ErrorLevel: allWriter,
				logrus.PanicLevel: allWriter,
			},
			&SaveFormatter{},
		))
	}
	//设置日志级别
	logger.SetLevel(logrus.DebugLevel)
	Logger = logger
	isInit = true
	return logger
}

func LoggerToFile(env string, logPath string) gin.HandlerFunc {
	logger := InitReqLogger(env, logPath)
	return func(c *gin.Context) {
		// 开始时间
		startTime := time.Now()
		// 处理请求
		c.Next()
		// 结束时间
		endTime := time.Now()
		// 执行时间
		latencyTime := endTime.Sub(startTime)
		// 请求方式
		reqMethod := c.Request.Method
		// 请求路由
		reqUri := c.Request.RequestURI
		// 状态码
		statusCode := c.Writer.Status()
		// 请求IP
		clientIP := c.ClientIP()
		//日志格式
		logger.Infof("|%3d|%13v|%15s|%s|%s|",
			statusCode,
			latencyTime,
			clientIP,
			reqMethod,
			reqUri,
		)
	}
}