From 74e0676617be0d0e2ecb028aa098dc350f77cef3 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Mon, 9 Jan 2023 09:39:30 +0800
Subject: [PATCH] feat: add config to truncate long log content (#2767)

---
 core/logx/config.go      |  1 +
 core/logx/logs.go        |  4 ++++
 core/logx/vars.go        |  7 +++++--
 core/logx/writer.go      |  9 +++++++++
 core/logx/writer_test.go | 36 ++++++++++++++++++++++++++++++++++--
 zrpc/client.go           |  1 +
 6 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/core/logx/config.go b/core/logx/config.go
index 7d1a4c1e..8f0398fd 100644
--- a/core/logx/config.go
+++ b/core/logx/config.go
@@ -8,6 +8,7 @@ type LogConf struct {
 	TimeFormat          string `json:",optional"`
 	Path                string `json:",default=logs"`
 	Level               string `json:",default=info,options=[debug,info,error,severe]"`
+	MaxContentLength    uint32 `json:",optional"`
 	Compress            bool   `json:",optional"`
 	Stat                bool   `json:",default=true"`
 	KeepDays            int    `json:",optional"`
diff --git a/core/logx/logs.go b/core/logx/logs.go
index 65047d49..734629a2 100644
--- a/core/logx/logs.go
+++ b/core/logx/logs.go
@@ -20,6 +20,8 @@ var (
 	timeFormat = "2006-01-02T15:04:05.000Z07:00"
 	logLevel   uint32
 	encoding   uint32 = jsonEncodingType
+	// maxContentLength is used to truncate the log content, 0 for not truncating.
+	maxContentLength uint32
 	// use uint32 for atomic operations
 	disableLog  uint32
 	disableStat uint32
@@ -238,6 +240,8 @@ func SetUp(c LogConf) (err error) {
 			timeFormat = c.TimeFormat
 		}
 
+		atomic.StoreUint32(&maxContentLength, c.MaxContentLength)
+
 		switch c.Encoding {
 		case plainEncoding:
 			atomic.StoreUint32(&encoding, plainEncodingType)
diff --git a/core/logx/vars.go b/core/logx/vars.go
index b95c3215..441bf973 100644
--- a/core/logx/vars.go
+++ b/core/logx/vars.go
@@ -16,13 +16,13 @@ const (
 const (
 	jsonEncodingType = iota
 	plainEncodingType
+)
 
+const (
 	plainEncoding    = "plain"
 	plainEncodingSep = '\t'
 	sizeRotationRule = "size"
-)
 
-const (
 	accessFilename = "access.log"
 	errorFilename  = "error.log"
 	severeFilename = "severe.log"
@@ -53,6 +53,7 @@ const (
 	spanKey      = "span"
 	timestampKey = "@timestamp"
 	traceKey     = "trace"
+	truncatedKey = "truncated"
 )
 
 var (
@@ -60,4 +61,6 @@ var (
 	ErrLogPathNotSet = errors.New("log path must be set")
 	// ErrLogServiceNameNotSet is an error that indicates that the service name is not set.
 	ErrLogServiceNameNotSet = errors.New("log service name must be set")
+
+	truncatedField = Field(truncatedKey, true)
 )
diff --git a/core/logx/writer.go b/core/logx/writer.go
index 2d4ade04..d9c2fb00 100644
--- a/core/logx/writer.go
+++ b/core/logx/writer.go
@@ -278,6 +278,15 @@ func combineGlobalFields(fields []LogField) []LogField {
 }
 
 func output(writer io.Writer, level string, val interface{}, fields ...LogField) {
+	// only truncate string content, don't know how to truncate the values of other types.
+	if v, ok := val.(string); ok {
+		maxLen := atomic.LoadUint32(&maxContentLength)
+		if maxLen > 0 && len(v) > int(maxLen) {
+			val = v[:maxLen]
+			fields = append(fields, truncatedField)
+		}
+	}
+
 	fields = combineGlobalFields(fields)
 
 	switch atomic.LoadUint32(&encoding) {
diff --git a/core/logx/writer_test.go b/core/logx/writer_test.go
index 9f9c3598..c2b476e4 100644
--- a/core/logx/writer_test.go
+++ b/core/logx/writer_test.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"log"
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -157,9 +158,40 @@ func TestWritePlainAny(t *testing.T) {
 
 }
 
+func TestLogWithLimitContentLength(t *testing.T) {
+	maxLen := atomic.LoadUint32(&maxContentLength)
+	atomic.StoreUint32(&maxContentLength, 10)
+
+	t.Cleanup(func() {
+		atomic.StoreUint32(&maxContentLength, maxLen)
+	})
+
+	t.Run("alert", func(t *testing.T) {
+		var buf bytes.Buffer
+		w := NewWriter(&buf)
+		w.Info("1234567890")
+		var v1 mockedEntry
+		if err := json.Unmarshal(buf.Bytes(), &v1); err != nil {
+			t.Fatal(err)
+		}
+		assert.Equal(t, "1234567890", v1.Content)
+		assert.False(t, v1.Truncated)
+
+		buf.Reset()
+		var v2 mockedEntry
+		w.Info("12345678901")
+		if err := json.Unmarshal(buf.Bytes(), &v2); err != nil {
+			t.Fatal(err)
+		}
+		assert.Equal(t, "1234567890", v2.Content)
+		assert.True(t, v2.Truncated)
+	})
+}
+
 type mockedEntry struct {
-	Level   string `json:"level"`
-	Content string `json:"content"`
+	Level     string `json:"level"`
+	Content   string `json:"content"`
+	Truncated bool   `json:"truncated"`
 }
 
 type easyToCloseWriter struct{}
diff --git a/zrpc/client.go b/zrpc/client.go
index 790a5a75..9966d038 100644
--- a/zrpc/client.go
+++ b/zrpc/client.go
@@ -89,6 +89,7 @@ func NewClientWithTarget(target string, opts ...ClientOption) (Client, error) {
 		Breaker:    true,
 		Timeout:    true,
 	}
+
 	return internal.NewClient(target, middlewares, opts...)
 }