Merge pull request #81 from Neur0toxine/handler-connection-account-replace

replace handler, connection and account log fields instead of duplicating

core/logger/attrs.go

@@ -85,12 +85,18 @@ func Body(val any) zap.Field {
 			return zap.Any(BodyAttr, m)
 		}
 		return zap.String(BodyAttr, item)
-	case []byte:
+	case []byte, json.RawMessage:
+		var val []byte
+		if msg, ok := item.(json.RawMessage); ok {
+			val = msg
+		} else {
+			val = item.([]byte)
+		}
 		var m interface{}
-		if err := json.Unmarshal(item, &m); err == nil {
+		if err := json.Unmarshal(val, &m); err == nil {
 			return zap.Any(BodyAttr, m)
 		}
-		return zap.String(BodyAttr, string(item))
+		return zap.String(BodyAttr, string(val))
 	case io.Reader:
 		data, err := io.ReadAll(item)
 		if err != nil {

core/logger/attrs_test.go

@@ -2,6 +2,7 @@ package logger
 
 import (
 	"bytes"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"github.com/stretchr/testify/require"
@@ -125,6 +126,11 @@ func TestBody(t *testing.T) {
 			input:  []byte("test body"),
 			result: "test body",
 		},
+		{
+			name:   "json.RawMessage input",
+			input:  json.RawMessage("test body"),
+			result: "test body",
+		},
 		{
 			name:   "json byte slice input",
 			input:  []byte(`{"success":true}`),

core/logger/buffer_logger_test.go

@@ -85,16 +85,49 @@ func (l *bufferLogger) WithLazy(fields ...zapcore.Field) Logger {
 	}
 }
 
+// ForHandler returns a new logger that is associated with the given handler.
+// This will replace "handler" field if it was set before.
+// Note: chain calls like ForHandler().With().ForHandler() will DUPLICATE handler field!
 func (l *bufferLogger) ForHandler(handler any) Logger {
-	return l.WithLazy(zap.Any(HandlerAttr, handler))
+	if l.previous != previousFieldHandler {
+		result := l.With(zap.Any(HandlerAttr, handler))
+		result.(*bufferLogger).setPrevious(previousFieldHandler)
+		result.(*bufferLogger).parent = l.Logger
+		return result
+	}
+	result := l.clone(l.parentOrCurrent().With(zap.Any(HandlerAttr, handler)))
+	result.(*bufferLogger).setPrevious(previousFieldHandler)
+	return result
 }
 
+// ForConnection returns a new logger that is associated with the given connection.
+// This will replace "connection" field if it was set before.
+// Note: chain calls like ForConnection().With().ForConnection() will DUPLICATE connection field!
 func (l *bufferLogger) ForConnection(conn any) Logger {
-	return l.WithLazy(zap.Any(ConnectionAttr, conn))
+	if l.previous != previousFieldConnection {
+		result := l.With(zap.Any(ConnectionAttr, conn))
+		result.(*bufferLogger).setPrevious(previousFieldConnection)
+		result.(*bufferLogger).parent = l.Logger
+		return result
+	}
+	result := l.clone(l.parentOrCurrent().With(zap.Any(ConnectionAttr, conn)))
+	result.(*bufferLogger).setPrevious(previousFieldConnection)
+	return result
 }
 
+// ForAccount returns a new logger that is associated with the given account.
+// This will replace "account" field if it was set before.
+// Note: chain calls like ForAccount().With().ForAccount() will DUPLICATE account field!
 func (l *bufferLogger) ForAccount(acc any) Logger {
-	return l.WithLazy(zap.Any(AccountAttr, acc))
+	if l.previous != previousFieldAccount {
+		result := l.With(zap.Any(AccountAttr, acc))
+		result.(*bufferLogger).setPrevious(previousFieldAccount)
+		result.(*bufferLogger).parent = l.Logger
+		return result
+	}
+	result := l.clone(l.parentOrCurrent().With(zap.Any(AccountAttr, acc)))
+	result.(*bufferLogger).setPrevious(previousFieldAccount)
+	return result
 }
 
 // Read bytes from the logger buffer. io.Reader implementation.
@@ -117,6 +150,28 @@ func (l *bufferLogger) Reset() {
 	l.buf.Reset()
 }
 
+// clone creates a copy of the given logger.
+func (l *bufferLogger) clone(log *zap.Logger) Logger {
+	parent := l.parent
+	if parent == nil {
+		parent = l.Logger
+	}
+	return &bufferLogger{
+		Default: Default{
+			Logger: log,
+			parent: parent,
+		},
+	}
+}
+
+// parentOrCurrent returns parent logger if it exists or current logger otherwise.
+func (l *bufferLogger) parentOrCurrent() *zap.Logger {
+	if l.parent != nil {
+		return l.parent
+	}
+	return l.Logger
+}
+
 type lockableBuffer struct {
 	buf bytes.Buffer
 	rw  sync.RWMutex

core/logger/default.go

@@ -44,9 +44,19 @@ type Logger interface {
 	Sync() error
 }
 
+type previousField uint8
+
+const (
+	previousFieldHandler previousField = iota + 1
+	previousFieldConnection
+	previousFieldAccount
+)
+
 // Default is a default logger implementation.
 type Default struct {
 	*zap.Logger
+	parent   *zap.Logger
+	previous previousField
 }
 
 // NewDefault creates a new default logger with the given format and debug level.
@@ -67,23 +77,69 @@ func (l *Default) WithLazy(fields ...zap.Field) Logger {
 }
 
 // ForHandler returns a new logger that is associated with the given handler.
+// This will replace "handler" field if it was set before.
+// Note: chain calls like ForHandler().With().ForHandler() will DUPLICATE handler field!
 func (l *Default) ForHandler(handler any) Logger {
-	return l.WithLazy(zap.Any(HandlerAttr, handler))
+	if l.previous != previousFieldHandler {
+		result := l.With(zap.Any(HandlerAttr, handler))
+		result.(*Default).setPrevious(previousFieldHandler)
+		result.(*Default).parent = l.Logger
+		return result
+	}
+	result := l.clone(l.parentOrCurrent().With(zap.Any(HandlerAttr, handler)))
+	result.(*Default).setPrevious(previousFieldHandler)
+	return result
 }
 
 // ForConnection returns a new logger that is associated with the given connection.
+// This will replace "connection" field if it was set before.
+// Note: chain calls like ForConnection().With().ForConnection() will DUPLICATE connection field!
 func (l *Default) ForConnection(conn any) Logger {
-	return l.WithLazy(zap.Any(ConnectionAttr, conn))
+	if l.previous != previousFieldConnection {
+		result := l.With(zap.Any(ConnectionAttr, conn))
+		result.(*Default).setPrevious(previousFieldConnection)
+		result.(*Default).parent = l.Logger
+		return result
+	}
+	result := l.clone(l.parentOrCurrent().With(zap.Any(ConnectionAttr, conn)))
+	result.(*Default).setPrevious(previousFieldConnection)
+	return result
 }
 
 // ForAccount returns a new logger that is associated with the given account.
+// This will replace "account" field if it was set before.
+// Note: chain calls like ForAccount().With().ForAccount() will DUPLICATE account field!
 func (l *Default) ForAccount(acc any) Logger {
-	return l.WithLazy(zap.Any(AccountAttr, acc))
+	if l.previous != previousFieldAccount {
+		result := l.With(zap.Any(AccountAttr, acc))
+		result.(*Default).setPrevious(previousFieldAccount)
+		result.(*Default).parent = l.Logger
+		return result
+	}
+	result := l.clone(l.parentOrCurrent().With(zap.Any(AccountAttr, acc)))
+	result.(*Default).setPrevious(previousFieldAccount)
+	return result
 }
 
 // clone creates a copy of the given logger.
 func (l *Default) clone(log *zap.Logger) Logger {
-	return &Default{Logger: log}
+	parent := l.parent
+	if parent == nil {
+		parent = l.Logger
+	}
+	return &Default{Logger: log, parent: parent}
+}
+
+// parentOrCurrent returns parent logger if it exists or current logger otherwise.
+func (l *Default) parentOrCurrent() *zap.Logger {
+	if l.parent != nil {
+		return l.parent
+	}
+	return l.Logger
+}
+
+func (l *Default) setPrevious(prev previousField) {
+	l.previous = prev
 }
 
 // AnyZapFields converts an array of values to zap fields.

core/logger/default_test.go

@@ -61,6 +61,14 @@ func (s *TestDefaultSuite) TestForHandler() {
 	s.Assert().Equal("Handler", items[0].Handler)
 }
 
+func (s *TestDefaultSuite) TestForHandlerNoDuplicate() {
+	log := newBufferLogger()
+	log.ForHandler("handler1").ForHandler("handler2").Info("test")
+
+	s.Assert().Contains(log.String(), "handler2")
+	s.Assert().NotContains(log.String(), "handler1")
+}
+
 func (s *TestDefaultSuite) TestForConnection() {
 	log := newBufferLogger()
 	log.ForConnection("connection").Info("test")
@@ -71,6 +79,14 @@ func (s *TestDefaultSuite) TestForConnection() {
 	s.Assert().Equal("connection", items[0].Connection)
 }
 
+func (s *TestDefaultSuite) TestForConnectionNoDuplicate() {
+	log := newBufferLogger()
+	log.ForConnection("conn1").ForConnection("conn2").Info("test")
+
+	s.Assert().Contains(log.String(), "conn2")
+	s.Assert().NotContains(log.String(), "conn1")
+}
+
 func (s *TestDefaultSuite) TestForAccount() {
 	log := newBufferLogger()
 	log.ForAccount("account").Info("test")
@@ -81,6 +97,56 @@ func (s *TestDefaultSuite) TestForAccount() {
 	s.Assert().Equal("account", items[0].Account)
 }
 
+func (s *TestDefaultSuite) TestForAccountNoDuplicate() {
+	log := newBufferLogger()
+	log.ForAccount("acc1").ForAccount("acc2").Info("test")
+
+	s.Assert().Contains(log.String(), "acc2")
+	s.Assert().NotContains(log.String(), "acc1")
+}
+
+func (s *TestDefaultSuite) TestNoDuplicatesPersistRecords() {
+	log := newBufferLogger()
+	log.
+		ForHandler("handler1").
+		ForHandler("handler2").
+		ForConnection("conn1").
+		ForConnection("conn2").
+		ForAccount("acc1").
+		ForAccount("acc2").
+		Info("test")
+
+	s.Assert().Contains(log.String(), "handler2")
+	s.Assert().NotContains(log.String(), "handler1")
+	s.Assert().Contains(log.String(), "conn2")
+	s.Assert().NotContains(log.String(), "conn1")
+	s.Assert().Contains(log.String(), "acc2")
+	s.Assert().NotContains(log.String(), "acc1")
+}
+
+// TestPersistRecordsIncompatibleWith is not a unit test, but rather a demonstration how you shouldn't use For* methods.
+func (s *TestDefaultSuite) TestPersistRecordsIncompatibleWith() {
+	log := newBufferLogger()
+	log.
+		ForHandler("handler1").
+		With(zap.Int("f1", 1)).
+		ForHandler("handler2").
+		ForConnection("conn1").
+		With(zap.Int("f2", 2)).
+		ForConnection("conn2").
+		ForAccount("acc1").
+		With(zap.Int("f3", 3)).
+		ForAccount("acc2").
+		Info("test")
+
+	s.Assert().Contains(log.String(), "handler2")
+	s.Assert().Contains(log.String(), "handler1")
+	s.Assert().Contains(log.String(), "conn2")
+	s.Assert().Contains(log.String(), "conn1")
+	s.Assert().Contains(log.String(), "acc2")
+	s.Assert().Contains(log.String(), "acc1")
+}
+
 func TestAnyZapFields(t *testing.T) {
 	fields := AnyZapFields([]interface{}{zap.String("k0", "v0"), "ooga", "booga"})
 	require.Len(t, fields, 3)

core/logger/json_with_context_encoder.go

@@ -44,6 +44,10 @@ var _jsonWithContextPool = NewPool(func() *jsonWithContextEncoder {
 })
 
 func init() {
+	registerJSONWithContext()
+}
+
+func registerJSONWithContext() {
 	err := zap.RegisterEncoder("json-with-context", func(config zapcore.EncoderConfig) (zapcore.Encoder, error) {
 		return NewJSONWithContextEncoder(config), nil
 	})