From bd2033eb3519d8d4def5b4f7c2fd97d1e94c0b30 Mon Sep 17 00:00:00 2001 From: Kevin Wan Date: Sun, 7 Jul 2024 23:31:27 +0800 Subject: [PATCH] feat: support adding more writer, easy to write to console additionally (#4234) Signed-off-by: kevin --- core/logx/logs.go | 20 +++++ core/logx/logs_test.go | 15 ++++ core/logx/writer.go | 61 ++++++++++++++++ core/logx/writer_test.go | 153 +++++++++++++++++++++++++++++++++++++++ go.mod | 1 + 5 files changed, 250 insertions(+) diff --git a/core/logx/logs.go b/core/logx/logs.go index 9cf3a791..627bab9b 100644 --- a/core/logx/logs.go +++ b/core/logx/logs.go @@ -52,6 +52,26 @@ type ( } ) +// AddWriter adds a new writer. +// If there is already a writer, the new writer will be added to the writer chain. +// For example, to write logs to both file and console, if there is already a file writer, +// ```go +// logx.AddWriter(logx.NewWriter(os.Stdout)) +// ``` +func AddWriter(w Writer) { + ow := Reset() + if ow == nil { + SetWriter(w) + } else { + // no need to check if the existing writer is a comboWriter, + // because it is not common to add more than one writer. + // even more than one writer, the behavior is the same. + SetWriter(comboWriter{ + writers: []Writer{ow, w}, + }) + } +} + // Alert alerts v in alert level, and the message is written to error log. func Alert(v string) { getWriter().Alert(v) diff --git a/core/logx/logs_test.go b/core/logx/logs_test.go index 5d46b7b9..4c831a7c 100644 --- a/core/logx/logs_test.go +++ b/core/logx/logs_test.go @@ -679,6 +679,10 @@ func TestSetup(t *testing.T) { func TestDisable(t *testing.T) { Disable() + defer func() { + SetLevel(InfoLevel) + atomic.StoreUint32(&encoding, jsonEncodingType) + }() var opt logOptions WithKeepDays(1)(&opt) @@ -701,6 +705,17 @@ func TestDisableStat(t *testing.T) { assert.Equal(t, 0, w.builder.Len()) } +func TestAddWriter(t *testing.T) { + const message = "hello there" + w := new(mockWriter) + AddWriter(w) + w1 := new(mockWriter) + AddWriter(w1) + Error(message) + assert.Contains(t, w.String(), message) + assert.Contains(t, w1.String(), message) +} + func TestSetWriter(t *testing.T) { atomic.StoreUint32(&logLevel, 0) Reset() diff --git a/core/logx/writer.go b/core/logx/writer.go index 09ae34b7..520b063a 100644 --- a/core/logx/writer.go +++ b/core/logx/writer.go @@ -13,6 +13,7 @@ import ( fatihcolor "github.com/fatih/color" "github.com/zeromicro/go-zero/core/color" + "github.com/zeromicro/go-zero/core/errorx" ) type ( @@ -33,6 +34,10 @@ type ( lock sync.RWMutex } + comboWriter struct { + writers []Writer + } + concreteWriter struct { infoLog io.WriteCloser errorLog io.WriteCloser @@ -88,6 +93,62 @@ func (w *atomicWriter) Swap(v Writer) Writer { return old } +func (c comboWriter) Alert(v any) { + for _, w := range c.writers { + w.Alert(v) + } +} + +func (c comboWriter) Close() error { + var be errorx.BatchError + for _, w := range c.writers { + be.Add(w.Close()) + } + return be.Err() +} + +func (c comboWriter) Debug(v any, fields ...LogField) { + for _, w := range c.writers { + w.Debug(v, fields...) + } +} + +func (c comboWriter) Error(v any, fields ...LogField) { + for _, w := range c.writers { + w.Error(v, fields...) + } +} + +func (c comboWriter) Info(v any, fields ...LogField) { + for _, w := range c.writers { + w.Info(v, fields...) + } +} + +func (c comboWriter) Severe(v any) { + for _, w := range c.writers { + w.Severe(v) + } +} + +func (c comboWriter) Slow(v any, fields ...LogField) { + for _, w := range c.writers { + w.Slow(v, fields...) + } +} + +func (c comboWriter) Stack(v any) { + for _, w := range c.writers { + w.Stack(v) + } +} + +func (c comboWriter) Stat(v any, fields ...LogField) { + for _, w := range c.writers { + w.Stat(v, fields...) + } +} + func newConsoleWriter() Writer { outLog := newLogWriter(log.New(fatihcolor.Output, "", flags)) errLog := newLogWriter(log.New(fatihcolor.Error, "", flags)) diff --git a/core/logx/writer_test.go b/core/logx/writer_test.go index 6ac98bc1..7d223dd4 100644 --- a/core/logx/writer_test.go +++ b/core/logx/writer_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" ) func TestNewWriter(t *testing.T) { @@ -254,6 +255,117 @@ func TestLogWithLimitContentLength(t *testing.T) { }) } +func TestComboWriter(t *testing.T) { + var mockWriters []Writer + for i := 0; i < 3; i++ { + mockWriters = append(mockWriters, new(tracedWriter)) + } + + cw := comboWriter{ + writers: mockWriters, + } + + t.Run("Alert", func(t *testing.T) { + for _, mw := range cw.writers { + mw.(*tracedWriter).On("Alert", "test alert").Once() + } + cw.Alert("test alert") + for _, mw := range cw.writers { + mw.(*tracedWriter).AssertCalled(t, "Alert", "test alert") + } + }) + + t.Run("Close", func(t *testing.T) { + for i := range cw.writers { + if i == 1 { + cw.writers[i].(*tracedWriter).On("Close").Return(errors.New("error")).Once() + } else { + cw.writers[i].(*tracedWriter).On("Close").Return(nil).Once() + } + } + err := cw.Close() + assert.Error(t, err) + for _, mw := range cw.writers { + mw.(*tracedWriter).AssertCalled(t, "Close") + } + }) + + t.Run("Debug", func(t *testing.T) { + fields := []LogField{{Key: "key", Value: "value"}} + for _, mw := range cw.writers { + mw.(*tracedWriter).On("Debug", "test debug", fields).Once() + } + cw.Debug("test debug", fields...) + for _, mw := range cw.writers { + mw.(*tracedWriter).AssertCalled(t, "Debug", "test debug", fields) + } + }) + + t.Run("Error", func(t *testing.T) { + fields := []LogField{{Key: "key", Value: "value"}} + for _, mw := range cw.writers { + mw.(*tracedWriter).On("Error", "test error", fields).Once() + } + cw.Error("test error", fields...) + for _, mw := range cw.writers { + mw.(*tracedWriter).AssertCalled(t, "Error", "test error", fields) + } + }) + + t.Run("Info", func(t *testing.T) { + fields := []LogField{{Key: "key", Value: "value"}} + for _, mw := range cw.writers { + mw.(*tracedWriter).On("Info", "test info", fields).Once() + } + cw.Info("test info", fields...) + for _, mw := range cw.writers { + mw.(*tracedWriter).AssertCalled(t, "Info", "test info", fields) + } + }) + + t.Run("Severe", func(t *testing.T) { + for _, mw := range cw.writers { + mw.(*tracedWriter).On("Severe", "test severe").Once() + } + cw.Severe("test severe") + for _, mw := range cw.writers { + mw.(*tracedWriter).AssertCalled(t, "Severe", "test severe") + } + }) + + t.Run("Slow", func(t *testing.T) { + fields := []LogField{{Key: "key", Value: "value"}} + for _, mw := range cw.writers { + mw.(*tracedWriter).On("Slow", "test slow", fields).Once() + } + cw.Slow("test slow", fields...) + for _, mw := range cw.writers { + mw.(*tracedWriter).AssertCalled(t, "Slow", "test slow", fields) + } + }) + + t.Run("Stack", func(t *testing.T) { + for _, mw := range cw.writers { + mw.(*tracedWriter).On("Stack", "test stack").Once() + } + cw.Stack("test stack") + for _, mw := range cw.writers { + mw.(*tracedWriter).AssertCalled(t, "Stack", "test stack") + } + }) + + t.Run("Stat", func(t *testing.T) { + fields := []LogField{{Key: "key", Value: "value"}} + for _, mw := range cw.writers { + mw.(*tracedWriter).On("Stat", "test stat", fields).Once() + } + cw.Stat("test stat", fields...) + for _, mw := range cw.writers { + mw.(*tracedWriter).AssertCalled(t, "Stat", "test stat", fields) + } + }) +} + type mockedEntry struct { Level string `json:"level"` Content string `json:"content"` @@ -285,3 +397,44 @@ type hardToWriteWriter struct{} func (h hardToWriteWriter) Write(_ []byte) (_ int, _ error) { return 0, errors.New("write error") } + +type tracedWriter struct { + mock.Mock +} + +func (w *tracedWriter) Alert(v any) { + w.Called(v) +} + +func (w *tracedWriter) Close() error { + args := w.Called() + return args.Error(0) +} + +func (w *tracedWriter) Debug(v any, fields ...LogField) { + w.Called(v, fields) +} + +func (w *tracedWriter) Error(v any, fields ...LogField) { + w.Called(v, fields) +} + +func (w *tracedWriter) Info(v any, fields ...LogField) { + w.Called(v, fields) +} + +func (w *tracedWriter) Severe(v any) { + w.Called(v) +} + +func (w *tracedWriter) Slow(v any, fields ...LogField) { + w.Called(v, fields) +} + +func (w *tracedWriter) Stack(v any) { + w.Called(v) +} + +func (w *tracedWriter) Stat(v any, fields ...LogField) { + w.Called(v, fields) +} diff --git a/go.mod b/go.mod index 612e95fb..89194dac 100644 --- a/go.mod +++ b/go.mod @@ -96,6 +96,7 @@ require ( github.com/prometheus/common v0.45.0 // indirect github.com/prometheus/procfs v0.12.0 // indirect github.com/rivo/uniseg v0.2.0 // indirect + github.com/stretchr/objx v0.5.2 // indirect github.com/xdg-go/pbkdf2 v1.0.0 // indirect github.com/xdg-go/scram v1.1.2 // indirect github.com/xdg-go/stringprep v1.0.4 // indirect