From eb9ab48baaceb7d57d796febdf7138ec3fc6d8d9 Mon Sep 17 00:00:00 2001
From: hookokoko <648646891@qq.com>
Date: Tue, 7 Feb 2023 10:33:55 +0800
Subject: [PATCH] add non-block write log in asynchronous mode (#5150)

* add non-block write log in asynchronous mode

---------

Co-authored-by: chenhaokun <chenhaokun@itiger.com>
---
 CHANGELOG.md          |  1 +
 core/logs/log.go      | 25 ++++++++++-
 core/logs/log_test.go | 96 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 120 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c087e317..99a7b61a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,4 +1,5 @@
 # developing
+- [add non-block write log in asynchronous mode](https://github.com/beego/beego/pull/5150)
 - [Fix 5126: support bloom filter cache](https://github.com/beego/beego/pull/5126)
 - [Fix 5117: support write though cache](https://github.com/beego/beego/pull/5117)
 - [add read through for cache module](https://github.com/beego/beego/pull/5116)
diff --git a/core/logs/log.go b/core/logs/log.go
index 7d36d117..2dedc768 100644
--- a/core/logs/log.go
+++ b/core/logs/log.go
@@ -29,7 +29,6 @@
 //	log.Warn("warning")
 //	log.Debug("debug")
 //	log.Critical("critical")
-//
 package logs
 
 import (
@@ -115,6 +114,9 @@ type BeeLogger struct {
 	enableFuncCallDepth bool
 	enableFullFilePath  bool
 	asynchronous        bool
+	// Whether to discard logs when buffer is full and asynchronous is true
+	// No discard by default
+	logWithNonBlocking  bool
 	wg                  sync.WaitGroup
 	level               int
 	loggerFuncCallDepth int
@@ -175,6 +177,16 @@ func (bl *BeeLogger) Async(msgLen ...int64) *BeeLogger {
 	return bl
 }
 
+// AsyncNonBlockWrite Non-blocking write in asynchronous mode
+// Only works if asynchronous write logging is set
+func (bl *BeeLogger) AsyncNonBlockWrite() *BeeLogger {
+	if !bl.asynchronous {
+		return bl
+	}
+	bl.logWithNonBlocking = true
+	return bl
+}
+
 // SetLogger provides a given logger adapter into BeeLogger with config string.
 // config must in in JSON format like {"interval":360}}
 func (bl *BeeLogger) setLogger(adapterName string, configs ...string) error {
@@ -312,8 +324,17 @@ func (bl *BeeLogger) writeMsg(lm *LogMsg) error {
 		logM.FilePath = lm.FilePath
 		logM.LineNumber = lm.LineNumber
 		logM.Prefix = lm.Prefix
+
 		if bl.outputs != nil {
-			bl.msgChan <- lm
+			if bl.logWithNonBlocking {
+				select {
+				case bl.msgChan <- lm:
+				// discard log when channel is full
+				default:
+				}
+			} else {
+				bl.msgChan <- lm
+			}
 		} else {
 			logMsgPool.Put(lm)
 		}
diff --git a/core/logs/log_test.go b/core/logs/log_test.go
index c8a3a478..06f8b030 100644
--- a/core/logs/log_test.go
+++ b/core/logs/log_test.go
@@ -15,7 +15,11 @@
 package logs
 
 import (
+	"encoding/json"
+	"fmt"
+	"io"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -30,3 +34,95 @@ func TestBeeLoggerDelLogger(t *testing.T) {
 	SetPrefix("aaa")
 	Info("hello")
 }
+
+type mockLogger struct {
+	*logWriter
+	WriteCost time.Duration `json:"write_cost"` // Simulated log writing time consuming
+	writeCnt  int           // Count add 1 when writing log success, just for test result
+}
+
+func NewMockLogger() Logger {
+	return &mockLogger{
+		logWriter: &logWriter{writer: io.Discard},
+	}
+}
+
+func (m *mockLogger) Init(config string) error {
+	return json.Unmarshal([]byte(config), m)
+}
+
+func (m *mockLogger) WriteMsg(lm *LogMsg) error {
+	m.Lock()
+	msg := lm.Msg
+	msg += "\n"
+
+	time.Sleep(m.WriteCost)
+	if _, err := m.writer.Write([]byte(msg)); err != nil {
+		return err
+	}
+
+	m.writeCnt++
+	m.Unlock()
+	return nil
+}
+
+func (m *mockLogger) GetCnt() int {
+	return m.writeCnt
+}
+
+func (*mockLogger) Destroy()                    {}
+func (*mockLogger) Flush()                      {}
+func (*mockLogger) SetFormatter(_ LogFormatter) {}
+
+func TestBeeLogger_AsyncNonBlockWrite(t *testing.T) {
+	testCases := []struct {
+		name         string
+		before       func()
+		after        func()
+		msgLen       int64
+		writeCost    time.Duration
+		sendInterval time.Duration
+		writeCnt     int
+	}{
+		{
+			// Write log time is less than send log time, no blocking
+			name: "mock1",
+			after: func() {
+				_ = beeLogger.DelLogger("mock1")
+			},
+			msgLen:       5,
+			writeCnt:     10,
+			writeCost:    200 * time.Millisecond,
+			sendInterval: 300 * time.Millisecond,
+		},
+		{
+			// Write log time is less than send log time, discarded when blocking
+			name: "mock2",
+			after: func() {
+				_ = beeLogger.DelLogger("mock2")
+			},
+			writeCnt:     5,
+			msgLen:       5,
+			writeCost:    200 * time.Millisecond,
+			sendInterval: 10 * time.Millisecond,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			Register(tc.name, NewMockLogger)
+			err := beeLogger.SetLogger(tc.name, fmt.Sprintf(`{"write_cost": %d}`, tc.writeCost))
+			assert.Nil(t, err)
+
+			l := beeLogger.Async(tc.msgLen)
+			l.AsyncNonBlockWrite()
+
+			for i := 0; i < 10; i++ {
+				time.Sleep(tc.sendInterval)
+				l.Info(fmt.Sprintf("----%d----", i))
+			}
+			time.Sleep(1 * time.Second)
+			assert.Equal(t, tc.writeCnt, l.outputs[0].Logger.(*mockLogger).writeCnt)
+			tc.after()
+		})
+	}
+}