Repository: zeromicro/go-zero
Pull request #1872 (2022-05-07 07:38:36)

feat: logx with color

MarkJoyMa: LGTM
re-dylan: zrpc, httpc can add color.

Patch:
diff --git a/.gitignore b/.gitignore
index 2f2cb2f2..e756a6be 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,6 +18,7 @@
 # for test purpose
 **/adhoc
 go.work
+go.work.sum
 
 # gitlab ci
 .cache
diff --git a/core/color/color.go b/core/color/color.go
new file mode 100644
index 00000000..6e892d29
--- /dev/null
+++ b/core/color/color.go
@@ -0,0 +1,73 @@
+package color
+
+import "github.com/fatih/color"
+
+const (
+	// NoColor is no color for both foreground and background.
+	NoColor Color = iota
+	// FgBlack is the foreground color black.
+	FgBlack
+	// FgRed is the foreground color red.
+	FgRed
+	// FgGreen is the foreground color green.
+	FgGreen
+	// FgYellow is the foreground color yellow.
+	FgYellow
+	// FgBlue is the foreground color blue.
+	FgBlue
+	// FgMagenta is the foreground color magenta.
+	FgMagenta
+	// FgCyan is the foreground color cyan.
+	FgCyan
+	// FgWhite is the foreground color white.
+	FgWhite
+
+	// BgBlack is the background color black.
+	BgBlack
+	// BgRed is the background color red.
+	BgRed
+	// BgGreen is the background color green.
+	BgGreen
+	// BgYellow is the background color yellow.
+	BgYellow
+	// BgBlue is the background color blue.
+	BgBlue
+	// BgMagenta is the background color magenta.
+	BgMagenta
+	// BgCyan is the background color cyan.
+	BgCyan
+	// BgWhite is the background color white.
+	BgWhite
+)
+
+var colors = map[Color][]color.Attribute{
+	FgBlack:   {color.FgBlack, color.Bold},
+	FgRed:     {color.FgRed, color.Bold},
+	FgGreen:   {color.FgGreen, color.Bold},
+	FgYellow:  {color.FgYellow, color.Bold},
+	FgBlue:    {color.FgBlue, color.Bold},
+	FgMagenta: {color.FgMagenta, color.Bold},
+	FgCyan:    {color.FgCyan, color.Bold},
+	FgWhite:   {color.FgWhite, color.Bold},
+	BgBlack:   {color.BgBlack, color.FgHiWhite, color.Bold},
+	BgRed:     {color.BgRed, color.FgHiWhite, color.Bold},
+	BgGreen:   {color.BgGreen, color.FgHiWhite, color.Bold},
+	BgYellow:  {color.BgHiYellow, color.FgHiBlack, color.Bold},
+	BgBlue:    {color.BgBlue, color.FgHiWhite, color.Bold},
+	BgMagenta: {color.BgMagenta, color.FgHiWhite, color.Bold},
+	BgCyan:    {color.BgCyan, color.FgHiWhite, color.Bold},
+	BgWhite:   {color.BgHiWhite, color.FgHiBlack, color.Bold},
+}
+
+type Color uint32
+
+// WithColor returns a string with the given color applied.
+func WithColor(text string, colour Color) string {
+	c := color.New(colors[colour]...)
+	return c.Sprint(text)
+}
+
+// WithColorPadding returns a string with the given color applied with leading and trailing spaces.
+func WithColorPadding(text string, colour Color) string {
+	return WithColor(" "+text+" ", colour)
+}
diff --git a/core/logx/color.go b/core/logx/color.go
new file mode 100644
index 00000000..769c74c2
--- /dev/null
+++ b/core/logx/color.go
@@ -0,0 +1,26 @@
+package logx
+
+import (
+	"sync/atomic"
+
+	"github.com/zeromicro/go-zero/core/color"
+)
+
+// WithColor is a helper function to add color to a string, only in plain encoding.
+func WithColor(text string, colour color.Color) string {
+	if atomic.LoadUint32(&encoding) == plainEncodingType {
+		return color.WithColor(text, colour)
+	}
+
+	return text
+}
+
+// WithColorPadding is a helper function to add color to a string with leading and trailing spaces,
+// only in plain encoding.
+func WithColorPadding(text string, colour color.Color) string {
+	if atomic.LoadUint32(&encoding) == plainEncodingType {
+		return color.WithColorPadding(text, colour)
+	}
+
+	return text
+}
diff --git a/core/logx/logs.go b/core/logx/logs.go
index b663c7df..f2569cf4 100644
--- a/core/logx/logs.go
+++ b/core/logx/logs.go
@@ -223,13 +223,13 @@ func SetUp(c LogConf) error {
 	}
 
 	switch c.Mode {
-	case consoleMode:
-		setupWithConsole()
-		return nil
+	case fileMode:
+		return setupWithFiles(c)
 	case volumeMode:
 		return setupWithVolume(c)
 	default:
-		return setupWithFiles(c)
+		setupWithConsole()
+		return nil
 	}
 }
 
diff --git a/core/logx/rotatelogger.go b/core/logx/rotatelogger.go
index bd31ce7e..07a03dd1 100644
--- a/core/logx/rotatelogger.go
+++ b/core/logx/rotatelogger.go
@@ -210,6 +210,12 @@ func (l *RotateLogger) maybeCompressFile(file string) {
 			ErrorStack(r)
 		}
 	}()
+
+	if _, err := os.Stat(file); err != nil {
+		// file not exists or other error, ignore compression
+		return
+	}
+
 	compressLogFile(file)
 }
 
@@ -292,7 +298,7 @@ func compressLogFile(file string) {
 	start := time.Now()
 	Infof("compressing log file: %s", file)
 	if err := gzipFile(file); err != nil {
-		ErrorStackf("compress error: %s", err)
+		Errorf("compress error: %s", err)
 	} else {
 		Infof("compressed log file: %s, took %s", file, time.Since(start))
 	}
diff --git a/core/logx/vars.go b/core/logx/vars.go
index 4229556c..e9f6772a 100644
--- a/core/logx/vars.go
+++ b/core/logx/vars.go
@@ -28,6 +28,7 @@ const (
 	statFilename   = "stat.log"
 
 	consoleMode = "console"
+	fileMode    = "file"
 	volumeMode  = "volume"
 
 	levelAlert  = "alert"
diff --git a/core/logx/writer.go b/core/logx/writer.go
index 75f8b4af..e494df60 100644
--- a/core/logx/writer.go
+++ b/core/logx/writer.go
@@ -10,6 +10,8 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+
+	"github.com/zeromicro/go-zero/core/color"
 )
 
 type (
@@ -239,6 +241,7 @@ func output(writer io.Writer, level string, val interface{}, fields ...LogField)
 
 	switch atomic.LoadUint32(&encoding) {
 	case plainEncodingType:
+		level = wrapLevelWithColor(level)
 		writePlainAny(writer, level, val, buildFields(fields...)...)
 	default:
 		entry := make(logEntryWithFields)
@@ -252,6 +255,30 @@ func output(writer io.Writer, level string, val interface{}, fields ...LogField)
 	}
 }
 
+func wrapLevelWithColor(level string) string {
+	var colour color.Color
+	switch level {
+	case levelAlert:
+		colour = color.FgRed
+	case levelError:
+		colour = color.FgRed
+	case levelFatal:
+		colour = color.FgRed
+	case levelInfo:
+		colour = color.FgBlue
+	case levelSlow:
+		colour = color.FgYellow
+	case levelStat:
+		colour = color.FgGreen
+	}
+
+	if colour == color.NoColor {
+		return level
+	}
+
+	return color.WithColorPadding(level, colour)
+}
+
 func writeJson(writer io.Writer, info interface{}) {
 	if content, err := json.Marshal(info); err != nil {
 		log.Println(err.Error())
diff --git a/go.mod b/go.mod
index 5312ad95..4a25311a 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/ClickHouse/clickhouse-go v1.5.1
 	github.com/DATA-DOG/go-sqlmock v1.5.0
 	github.com/alicebob/miniredis/v2 v2.17.0
+	github.com/fatih/color v1.10.0
 	github.com/globalsign/mgo v0.0.0-20181015135952-eeefdecb41b8
 	github.com/go-redis/redis/v8 v8.11.4
 	github.com/go-sql-driver/mysql v1.6.0
@@ -15,11 +16,11 @@ require (
 	github.com/justinas/alice v1.2.0
 	github.com/lib/pq v1.10.4
 	github.com/olekukonko/tablewriter v0.0.5
-	github.com/prometheus/client_golang v1.11.0
+	github.com/prometheus/client_golang v1.11.1
 	github.com/spaolacci/murmur3 v1.1.0
 	github.com/stretchr/testify v1.7.0
-	go.etcd.io/etcd/api/v3 v3.5.2
-	go.etcd.io/etcd/client/v3 v3.5.2
+	go.etcd.io/etcd/api/v3 v3.5.4
+	go.etcd.io/etcd/client/v3 v3.5.4
 	go.mongodb.org/mongo-driver v1.9.0
 	go.opentelemetry.io/otel v1.3.0
 	go.opentelemetry.io/otel/exporters/jaeger v1.3.0
@@ -42,7 +43,6 @@ require (
 )
 
 require (
-	github.com/fatih/color v1.10.0 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/google/gofuzz v1.2.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
diff --git a/go.sum b/go.sum
index 415c98eb..d8f46d3e 100644
--- a/go.sum
+++ b/go.sum
@@ -353,8 +353,9 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
 github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
 github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo=
 github.com/prometheus/client_golang v1.7.1/go.mod h1:PY5Wy2awLA44sXw4AOSfFBetzPP4j5+D6mVACh+pe2M=
-github.com/prometheus/client_golang v1.11.0 h1:HNkLOAEQMIDv/K+04rukrLx6ch7msSRwf3/SASFAGtQ=
 github.com/prometheus/client_golang v1.11.0/go.mod h1:Z6t4BnS23TR94PD6BsDNk8yVqroYurpAkEiz0P2BEV0=
+github.com/prometheus/client_golang v1.11.1 h1:+4eQaD7vAZ6DsfsxB15hbE0odUjGI5ARs9yskGu1v4s=
+github.com/prometheus/client_golang v1.11.1/go.mod h1:Z6t4BnS23TR94PD6BsDNk8yVqroYurpAkEiz0P2BEV0=
 github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
 github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
@@ -418,12 +419,12 @@ github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9dec
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 github.com/yuin/gopher-lua v0.0.0-20200816102855-ee81675732da h1:NimzV1aGyq29m5ukMK0AMWEhFaL/lrEOaephfuoiARg=
 github.com/yuin/gopher-lua v0.0.0-20200816102855-ee81675732da/go.mod h1:E1AXubJBdNmFERAOucpDIxNzeGfLzg0mYh+UfMWdChA=
-go.etcd.io/etcd/api/v3 v3.5.2 h1:tXok5yLlKyuQ/SXSjtqHc4uzNaMqZi2XsoSPr/LlJXI=
-go.etcd.io/etcd/api/v3 v3.5.2/go.mod h1:5GB2vv4A4AOn3yk7MftYGHkUfGtDHnEraIjym4dYz5A=
-go.etcd.io/etcd/client/pkg/v3 v3.5.2 h1:4hzqQ6hIb3blLyQ8usCU4h3NghkqcsohEQ3o3VetYxE=
-go.etcd.io/etcd/client/pkg/v3 v3.5.2/go.mod h1:IJHfcCEKxYu1Os13ZdwCwIUTUVGYTSAM3YSwc9/Ac1g=
-go.etcd.io/etcd/client/v3 v3.5.2 h1:WdnejrUtQC4nCxK0/dLTMqKOB+U5TP/2Ya0BJL+1otA=
-go.etcd.io/etcd/client/v3 v3.5.2/go.mod h1:kOOaWFFgHygyT0WlSmL8TJiXmMysO/nNUlEsSsN6W4o=
+go.etcd.io/etcd/api/v3 v3.5.4 h1:OHVyt3TopwtUQ2GKdd5wu3PmmipR4FTwCqoEjSyRdIc=
+go.etcd.io/etcd/api/v3 v3.5.4/go.mod h1:5GB2vv4A4AOn3yk7MftYGHkUfGtDHnEraIjym4dYz5A=
+go.etcd.io/etcd/client/pkg/v3 v3.5.4 h1:lrneYvz923dvC14R54XcA7FXoZ3mlGZAgmwhfm7HqOg=
+go.etcd.io/etcd/client/pkg/v3 v3.5.4/go.mod h1:IJHfcCEKxYu1Os13ZdwCwIUTUVGYTSAM3YSwc9/Ac1g=
+go.etcd.io/etcd/client/v3 v3.5.4 h1:p83BUL3tAYS0OT/r0qglgc3M1JjhM0diV8DSWAhVXv4=
+go.etcd.io/etcd/client/v3 v3.5.4/go.mod h1:ZaRkVgBZC+L+dLCjTcF1hRXpgZXQPOvnA/Ak/gq3kiY=
 go.mongodb.org/mongo-driver v1.9.0 h1:f3aLGJvQmBl8d9S40IL+jEyBC6hfLPbJjv9t5hEM9ck=
 go.mongodb.org/mongo-driver v1.9.0/go.mod h1:0sQWfOeY63QTntERDJJ/0SuKK0T1uVSgKCuAROlKEPY=
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
diff --git a/rest/handler/loghandler.go b/rest/handler/loghandler.go
index f9b0a872..8d9ae44f 100644
--- a/rest/handler/loghandler.go
+++ b/rest/handler/loghandler.go
@@ -11,9 +11,11 @@ import (
 	"net"
 	"net/http"
 	"net/http/httputil"
+	"strconv"
 	"strings"
 	"time"
 
+	"github.com/zeromicro/go-zero/core/color"
 	"github.com/zeromicro/go-zero/core/iox"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/syncx"
@@ -157,15 +159,21 @@ func dumpRequest(r *http.Request) string {
 	return string(reqContent)
 }
 
+func isOkResponse(code int) bool {
+	// not server error
+	return code < http.StatusInternalServerError
+}
+
 func logBrief(r *http.Request, code int, timer *utils.ElapsedTimer, logs *internal.LogCollector) {
 	var buf bytes.Buffer
 	duration := timer.Duration()
 	logger := logx.WithContext(r.Context()).WithDuration(duration)
-	buf.WriteString(fmt.Sprintf("[HTTP] %s - %d - %s - %s - %s",
-		r.Method, code, r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent()))
+	buf.WriteString(fmt.Sprintf("[HTTP] %s - %s %s - %s - %s",
+		wrapStatusCode(code), wrapMethod(r.Method), r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent()))
 	if duration > slowThreshold.Load() {
-		logger.Slowf("[HTTP] %s - %d - %s - %s - %s - slowcall(%s)",
-			r.Method, code, r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent(), timex.ReprOfDuration(duration))
+		logger.Slowf("[HTTP] %s - %s - %s %s - %s - slowcall(%s)",
+			wrapStatusCode(code), wrapMethod(r.Method), r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent(),
+			fmt.Sprintf("slowcall(%s)", timex.ReprOfDuration(duration)))
 	}
 
 	ok := isOkResponse(code)
@@ -201,8 +209,8 @@ func logDetails(r *http.Request, response *detailLoggedResponseWriter, timer *ut
 	buf.WriteString(fmt.Sprintf("[HTTP] %s - %d - %s - %s\n=> %s\n",
 		r.Method, code, r.RemoteAddr, timex.ReprOfDuration(duration), dumpRequest(r)))
 	if duration > defaultSlowThreshold {
-		logger.Slowf("[HTTP] %s - %d - %s - slowcall(%s)\n=> %s\n",
-			r.Method, code, r.RemoteAddr, timex.ReprOfDuration(duration), dumpRequest(r))
+		logger.Slowf("[HTTP] %s - %d - %s - slowcall(%s)\n=> %s\n", r.Method, code, r.RemoteAddr,
+			fmt.Sprintf("slowcall(%s)", timex.ReprOfDuration(duration)), dumpRequest(r))
 	}
 
 	body := logs.Flush()
@@ -222,7 +230,44 @@ func logDetails(r *http.Request, response *detailLoggedResponseWriter, timer *ut
 	}
 }
 
-func isOkResponse(code int) bool {
-	// not server error
-	return code < http.StatusInternalServerError
+func wrapMethod(method string) string {
+	var colour color.Color
+	switch method {
+	case http.MethodGet:
+		colour = color.BgBlue
+	case http.MethodPost:
+		colour = color.BgCyan
+	case http.MethodPut:
+		colour = color.BgYellow
+	case http.MethodDelete:
+		colour = color.BgRed
+	case http.MethodPatch:
+		colour = color.BgGreen
+	case http.MethodHead:
+		colour = color.BgMagenta
+	case http.MethodOptions:
+		colour = color.BgWhite
+	}
+
+	if colour == color.NoColor {
+		return method
+	}
+
+	return logx.WithColorPadding(method, colour)
+}
+
+func wrapStatusCode(code int) string {
+	var colour color.Color
+	switch {
+	case code >= http.StatusOK && code < http.StatusMultipleChoices:
+		colour = color.BgGreen
+	case code >= http.StatusMultipleChoices && code < http.StatusBadRequest:
+		colour = color.BgBlue
+	case code >= http.StatusBadRequest && code < http.StatusInternalServerError:
+		colour = color.BgMagenta
+	default:
+		colour = color.BgYellow
+	}
+
+	return logx.WithColorPadding(strconv.Itoa(code), colour)
 }


Test patch:
diff --git a/core/color/color_test.go b/core/color/color_test.go
new file mode 100644
index 00000000..74ee2a12
--- /dev/null
+++ b/core/color/color_test.go
@@ -0,0 +1,17 @@
+package color
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWithColor(t *testing.T) {
+	output := WithColor("Hello", BgRed)
+	assert.Equal(t, "Hello", output)
+}
+
+func TestWithColorPadding(t *testing.T) {
+	output := WithColorPadding("Hello", BgRed)
+	assert.Equal(t, " Hello ", output)
+}
diff --git a/core/logx/color_test.go b/core/logx/color_test.go
new file mode 100644
index 00000000..b3a9205a
--- /dev/null
+++ b/core/logx/color_test.go
@@ -0,0 +1,33 @@
+package logx
+
+import (
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/color"
+)
+
+func TestWithColor(t *testing.T) {
+	old := atomic.SwapUint32(&encoding, plainEncodingType)
+	defer atomic.StoreUint32(&encoding, old)
+
+	output := WithColor("hello", color.BgBlue)
+	assert.Equal(t, "hello", output)
+
+	atomic.StoreUint32(&encoding, jsonEncodingType)
+	output = WithColor("hello", color.BgBlue)
+	assert.Equal(t, "hello", output)
+}
+
+func TestWithColorPadding(t *testing.T) {
+	old := atomic.SwapUint32(&encoding, plainEncodingType)
+	defer atomic.StoreUint32(&encoding, old)
+
+	output := WithColorPadding("hello", color.BgBlue)
+	assert.Equal(t, " hello ", output)
+
+	atomic.StoreUint32(&encoding, jsonEncodingType)
+	output = WithColorPadding("hello", color.BgBlue)
+	assert.Equal(t, "hello", output)
+}
diff --git a/rest/server_test.go b/rest/server_test.go
index 4d4d21be..7f0c667e 100644
--- a/rest/server_test.go
+++ b/rest/server_test.go
@@ -4,6 +4,7 @@ import (
 	"crypto/tls"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -11,11 +12,16 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/conf"
+	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/rest/httpx"
 	"github.com/zeromicro/go-zero/rest/router"
 )
 
 func TestNewServer(t *testing.T) {
+	writer := logx.Reset()
+	defer logx.SetWriter(writer)
+	logx.SetWriter(logx.NewWriter(ioutil.Discard))
+
 	const configYaml = `
 Name: foo
 Port: 54321
@@ -31,7 +37,6 @@ Port: 54321
 		{
 			c:    RestConf{},
 			opts: []RunOption{WithRouter(mockedRouter{}), WithCors()},
-			fail: true,
 		},
 		{
 			c:    cnf,