logrus_hook.go 2.1 KB
package log

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/beego/beego/v2/core/logs"
	"io"
	"sync/atomic"
)

var errOutOfMaxSize = errors.New("msg  size is out of limit ")

type BaseHook struct {
	syncWrite   bool
	msgChan     chan []byte
	maxSize     int32
	currentSize int32
	closeChan   chan struct{}
	w           io.Writer
}

func (b *BaseHook) Write(p []byte) (n int, err error) {
	if b.syncWrite {
		return b.w.Write(p)
	}
	if b.currentSize >= b.maxSize {
		fmt.Println(errOutOfMaxSize.Error(), b.currentSize)
		return 0, errOutOfMaxSize
	}
	b.msgChan <- p
	atomic.AddInt32(&b.currentSize, 1)

	return len(p), nil
}

// NewBaseHook
// syncWriteFlag 同步写标识 true:同步写 false:异步写
func NewBaseHook(syncWriteFlag bool, internalIo io.WriteCloser) (*BaseHook, error) {
	writer := &BaseHook{
		syncWrite: syncWriteFlag,
		maxSize:   10000,
		msgChan:   make(chan []byte, 10000),
		closeChan: make(chan struct{}),
		w:         internalIo,
	}
	go writer.ConsumeMsg()
	return writer, nil
}

func (b *BaseHook) ConsumeMsg() {
	for {
		select {
		case <-b.closeChan:
			return
		case m, ok := <-b.msgChan:
			if ok {
				atomic.AddInt32(&b.currentSize, -1)
				if _, err := b.w.Write(m); err != nil {
					fmt.Println(err)
				}
			}
		}
	}
}

func (b *BaseHook) Close() {
	close(b.msgChan)
	b.closeChan <- struct{}{}
	//b.wc.Close()
}

type LoggerConfig struct {
	Level    int    `json:"level,omitempty"`
	Filename string `json:"filename,omitempty"`
	MaxSize  int    `json:"maxsize,omitempty"`
	//MaxBackups int    `json:"max_backups,omitempty"`
	//MaxAge     int    `json:"max_age,omitempty"`
	//Compress   bool   `json:"compress,omitempty"`
}

type internalLog struct {
	*logs.BeeLogger
}

func NewBeegoLogWriter(conf LoggerConfig) io.Writer {
	logger := logs.GetBeeLogger()
	logger.SetLevel(conf.Level)
	logger.EnableFuncCallDepth(true)
	logger.SetLogFuncCallDepth(2)
	confByte, _ := json.Marshal(conf)
	err := logger.SetLogger(logs.AdapterFile, string(confByte))
	if err != nil {
		fmt.Println(err.Error())
	}
	return &internalLog{logger}
}

func (l *internalLog) Write(p []byte) (n int, err error) {
	l.Debug(string(p))
	return len(p), nil
}