refactor: improved memlogger and remove html log formatting

internal/logging/html.go

@@ -1,159 +0,0 @@
-package logging
-
-import (
-	"errors"
-	"fmt"
-	"time"
-
-	"github.com/rs/zerolog"
-	"github.com/yusing/go-proxy/internal/common"
-)
-
-var levelHTMLFormats = [][]byte{
-	[]byte(` <span class="log-trace">TRC</span> `),
-	[]byte(` <span class="log-debug">DBG</span> `),
-	[]byte(` <span class="log-info">INF</span> `),
-	[]byte(` <span class="log-warn">WRN</span> `),
-	[]byte(` <span class="log-error">ERR</span> `),
-	[]byte(` <span class="log-fatal">FTL</span> `),
-	[]byte(` <span class="log-panic">PAN</span> `),
-}
-
-var colorToClass = map[string]string{
-	"1":  "log-bold",
-	"3":  "log-italic",
-	"4":  "log-underline",
-	"30": "log-black",
-	"31": "log-red",
-	"32": "log-green",
-	"33": "log-yellow",
-	"34": "log-blue",
-	"35": "log-magenta",
-	"36": "log-cyan",
-	"37": "log-white",
-	"90": "log-bright-black",
-	"91": "log-red",
-	"92": "log-bright-green",
-	"93": "log-bright-yellow",
-	"94": "log-bright-blue",
-	"95": "log-bright-magenta",
-	"96": "log-bright-cyan",
-	"97": "log-bright-white",
-}
-
-// FormatMessageToHTMLBytes converts text with ANSI color codes to HTML with class names.
-// ANSI codes are mapped to classes via a static map, and reset codes ([0m) close all spans.
-// Time complexity is O(n) with minimal allocations.
-func FormatMessageToHTMLBytes(msg string, buf []byte) ([]byte, error) {
-	buf = append(buf, "<span class=\"log-message\">"...)
-	var stack []string
-	lastPos := 0
-
-	for i := 0; i < len(msg); {
-		if msg[i] == '\x1b' && i+1 < len(msg) && msg[i+1] == '[' {
-			if lastPos < i {
-				escapeAndAppend(msg[lastPos:i], &buf)
-			}
-			i += 2 // Skip \x1b[
-
-			start := i
-			for ; i < len(msg) && msg[i] != 'm'; i++ {
-				if !isANSICodeChar(msg[i]) {
-					return nil, fmt.Errorf("invalid ANSI char: %c", msg[i])
-				}
-			}
-
-			if i >= len(msg) {
-				return nil, errors.New("unterminated ANSI sequence")
-			}
-
-			codeStr := msg[start:i]
-			i++ // Skip 'm'
-			lastPos = i
-
-			startPart := 0
-			for j := 0; j <= len(codeStr); j++ {
-				if j == len(codeStr) || codeStr[j] == ';' {
-					part := codeStr[startPart:j]
-					if part == "" {
-						return nil, errors.New("empty code part")
-					}
-
-					if part == "0" {
-						for range stack {
-							buf = append(buf, "</span>"...)
-						}
-						stack = stack[:0]
-					} else {
-						className, ok := colorToClass[part]
-						if !ok {
-							return nil, fmt.Errorf("invalid ANSI code: %s", part)
-						}
-						stack = append(stack, className)
-						buf = append(buf, `<span class="`...)
-						buf = append(buf, className...)
-						buf = append(buf, `">`...)
-					}
-					startPart = j + 1
-				}
-			}
-		} else {
-			i++
-		}
-	}
-
-	if lastPos < len(msg) {
-		escapeAndAppend(msg[lastPos:], &buf)
-	}
-
-	for range stack {
-		buf = append(buf, "</span>"...)
-	}
-
-	buf = append(buf, "</span>"...)
-	return buf, nil
-}
-
-func isANSICodeChar(c byte) bool {
-	return (c >= '0' && c <= '9') || c == ';'
-}
-
-func escapeAndAppend(s string, buf *[]byte) {
-	for i, r := range s {
-		switch r {
-		case '•':
-			*buf = append(*buf, "&middot;"...)
-		case '&':
-			*buf = append(*buf, "&amp;"...)
-		case '<':
-			*buf = append(*buf, "&lt;"...)
-		case '>':
-			*buf = append(*buf, "&gt;"...)
-		case '\t':
-			*buf = append(*buf, "&#9;"...)
-		case '\n':
-			*buf = append(*buf, "<br>"...)
-			*buf = append(*buf, prefixHTML...)
-		default:
-			*buf = append(*buf, s[i])
-		}
-	}
-}
-
-func timeNowHTML() []byte {
-	if !common.IsTest {
-		return []byte(time.Now().Format(timeFmt))
-	}
-	return []byte(time.Date(2024, 1, 1, 1, 1, 1, 1, time.UTC).Format(timeFmt))
-}
-
-func FormatLogEntryHTML(level zerolog.Level, message string, buf []byte) []byte {
-	buf = append(buf, []byte(`<pre class="log-entry">`)...)
-	buf = append(buf, timeNowHTML()...)
-	if level < zerolog.NoLevel {
-		buf = append(buf, levelHTMLFormats[level+1]...)
-	}
-	buf, _ = FormatMessageToHTMLBytes(message, buf)
-	buf = append(buf, []byte("</pre>")...)
-	return buf
-}

internal/logging/html_test.go

@@ -1,30 +0,0 @@
-package logging
-
-import (
-	"testing"
-
-	"github.com/rs/zerolog"
-	. "github.com/yusing/go-proxy/internal/utils/testing"
-)
-
-func TestFormatHTML(t *testing.T) {
-	buf := make([]byte, 0, 100)
-	buf = FormatLogEntryHTML(zerolog.InfoLevel, "This is a test.\nThis is a new line.", buf)
-	ExpectEqual(t, string(buf), `<pre class="log-entry">01-01 01:01 <span class="log-info">INF</span> <span class="log-message">This is a test.<br>`+prefix+`This is a new line.</span></pre>`)
-}
-
-func TestFormatHTMLANSI(t *testing.T) {
-	buf := make([]byte, 0, 100)
-	buf = FormatLogEntryHTML(zerolog.InfoLevel, "This is \x1b[91m\x1b[1ma test.\x1b[0mOK!.", buf)
-	ExpectEqual(t, string(buf), `<pre class="log-entry">01-01 01:01 <span class="log-info">INF</span> <span class="log-message">This is <span class="log-red"><span class="log-bold">a test.</span></span>OK!.</span></pre>`)
-	buf = buf[:0]
-	buf = FormatLogEntryHTML(zerolog.InfoLevel, "This is \x1b[91ma \x1b[1mtest.\x1b[0mOK!.", buf)
-	ExpectEqual(t, string(buf), `<pre class="log-entry">01-01 01:01 <span class="log-info">INF</span> <span class="log-message">This is <span class="log-red">a <span class="log-bold">test.</span></span>OK!.</span></pre>`)
-}
-
-func BenchmarkFormatLogEntryHTML(b *testing.B) {
-	buf := make([]byte, 0, 250)
-	for range b.N {
-		FormatLogEntryHTML(zerolog.InfoLevel, "This is \x1b[91ma \x1b[1mtest.\x1b[0mOK!.", buf)
-	}
-}

internal/logging/logging.go

@@ -3,6 +3,7 @@ package logging
 
 import (
 	"io"
+	"log"
 	"strings"
 
 	"github.com/rs/zerolog"
@@ -11,11 +12,10 @@ import (
 )
 
 var (
-	logger     zerolog.Logger
+	logger  zerolog.Logger
-	timeFmt    string
+	timeFmt string
-	level      zerolog.Level
+	level   zerolog.Level
-	prefix     string
+	prefix  string
-	prefixHTML []byte
 )
 
 func init() {
@@ -32,12 +32,6 @@ func init() {
 	}
 	prefixLength := len(timeFmt) + 5 // level takes 3 + 2 spaces
 	prefix = strings.Repeat(" ", prefixLength)
-	// prefixHTML = []byte(strings.Repeat(" ", prefixLength))
-	prefixHTML = []byte(prefix)
-
-	if zerolog.TraceLevel != -1 && zerolog.NoLevel != 6 {
-		panic("zerolog implementation changed")
-	}
 }
 
 func fmtMessage(msg string) string {
@@ -62,6 +56,9 @@ func InitLogger(out io.Writer) {
 	logger = zerolog.New(
 		writer,
 	).Level(level).With().Timestamp().Logger()
+	log.SetOutput(writer)
+	log.SetPrefix("")
+	log.SetFlags(0)
 }
 
 func DiscardLogger() { zerolog.SetGlobalLevel(zerolog.Disabled) }

internal/logging/memlogger/mem_logger.go

@@ -1,20 +1,19 @@
-package v1
+package memlogger
 
 import (
 	"bytes"
 	"context"
 	"io"
 	"net/http"
+	"os"
 	"sync"
 	"time"
 
 	"github.com/coder/websocket"
 	"github.com/rs/zerolog"
-	"github.com/yusing/go-proxy/internal/api/v1/utils"
-	"github.com/yusing/go-proxy/internal/common"
-	config "github.com/yusing/go-proxy/internal/config/types"
 	"github.com/yusing/go-proxy/internal/logging"
-	"github.com/yusing/go-proxy/internal/task"
+	"github.com/yusing/go-proxy/internal/net/gphttp"
+	"github.com/yusing/go-proxy/internal/net/gphttp/gpwebsocket"
 	F "github.com/yusing/go-proxy/internal/utils/functional"
 )
 
@@ -27,15 +26,10 @@ type memLogger struct {
 	sync.RWMutex
 	notifyLock sync.RWMutex
 	connChans  F.Map[chan *logEntryRange, struct{}]
-
+	listeners  F.Map[chan []byte, struct{}]
-	bufPool sync.Pool // used in hook mode
 }
 
-type MemLogger interface {
+type MemLogger io.Writer
-	io.Writer
-	// TODO: hook does not pass in fields, looking for a workaround to do server side log rendering
-	zerolog.Hook
-}
 
 type buffer struct {
 	data []byte
@@ -45,54 +39,35 @@ const (
 	maxMemLogSize         = 16 * 1024
 	truncateSize          = maxMemLogSize / 2
 	initialWriteChunkSize = 4 * 1024
-	hookModeBufSize       = 256
 )
 
 var memLoggerInstance = &memLogger{
 	connChans: F.NewMapOf[chan *logEntryRange, struct{}](),
-	bufPool: sync.Pool{
+	listeners: F.NewMapOf[chan []byte, struct{}](),
-		New: func() any {
-			return &buffer{
-				data: make([]byte, 0, hookModeBufSize),
-			}
-		},
-	},
 }
 
 func init() {
-	if !common.EnableLogStreaming {
-		return
-	}
 	memLoggerInstance.Grow(maxMemLogSize)
-
+	w := zerolog.MultiLevelWriter(os.Stderr, memLoggerInstance)
-	if common.DebugMemLogger {
+	logging.InitLogger(w)
-		ticker := time.NewTicker(1 * time.Second)
-
-		go func() {
-			defer ticker.Stop()
-
-			for {
-				select {
-				case <-task.RootContextCanceled():
-					return
-				case <-ticker.C:
-					logging.Info().Msgf("mem logger size: %d, active conns: %d",
-						memLoggerInstance.Len(),
-						memLoggerInstance.connChans.Size())
-				}
-			}
-		}()
-	}
-}
-
-func LogsWS() func(config config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
-	return memLoggerInstance.ServeHTTP
 }
 
 func GetMemLogger() MemLogger {
 	return memLoggerInstance
 }
 
+func Handler() http.Handler {
+	return memLoggerInstance
+}
+
+func HandlerFunc() http.HandlerFunc {
+	return memLoggerInstance.ServeHTTP
+}
+
+func Events() (<-chan []byte, func()) {
+	return memLoggerInstance.events()
+}
+
 func (m *memLogger) truncateIfNeeded(n int) {
 	m.RLock()
 	needTruncate := m.Len()+n > maxMemLogSize
@@ -111,22 +86,35 @@ func (m *memLogger) truncateIfNeeded(n int) {
 }
 
 func (m *memLogger) notifyWS(pos, n int) {
-	if m.connChans.Size() > 0 {
+	if m.connChans.Size() == 0 && m.listeners.Size() == 0 {
-		timeout := time.NewTimer(1 * time.Second)
+		return
-		defer timeout.Stop()
+	}
 
-		m.notifyLock.RLock()
+	timeout := time.NewTimer(3 * time.Second)
-		defer m.notifyLock.RUnlock()
+	defer timeout.Stop()
-		m.connChans.Range(func(ch chan *logEntryRange, _ struct{}) bool {
+
+	m.notifyLock.RLock()
+	defer m.notifyLock.RUnlock()
+
+	m.connChans.Range(func(ch chan *logEntryRange, _ struct{}) bool {
+		select {
+		case ch <- &logEntryRange{pos, pos + n}:
+			return true
+		case <-timeout.C:
+			return false
+		}
+	})
+
+	if m.listeners.Size() > 0 {
+		msg := m.Buffer.Bytes()[pos : pos+n]
+		m.listeners.Range(func(ch chan []byte, _ struct{}) bool {
 			select {
-			case ch <- &logEntryRange{pos, pos + n}:
-				return true
 			case <-timeout.C:
-				logging.Warn().Msg("mem logger: timeout logging to channel")
 				return false
+			case ch <- msg:
+				return true
 			}
 		})
-		return
 	}
 }
 
@@ -138,29 +126,6 @@ func (m *memLogger) writeBuf(b []byte) (pos int, err error) {
 	return
 }
 
-// Run implements zerolog.Hook.
-func (m *memLogger) Run(e *zerolog.Event, level zerolog.Level, message string) {
-	bufStruct := m.bufPool.Get().(*buffer)
-	buf := bufStruct.data
-	defer func() {
-		bufStruct.data = bufStruct.data[:0]
-		m.bufPool.Put(bufStruct)
-	}()
-
-	buf = logging.FormatLogEntryHTML(level, message, buf)
-	n := len(buf)
-
-	m.truncateIfNeeded(n)
-
-	pos, err := m.writeBuf(buf)
-	if err != nil {
-		// not logging the error here, it will cause Run to be called again = infinite loop
-		return
-	}
-
-	m.notifyWS(pos, n)
-}
-
 // Write implements io.Writer.
 func (m *memLogger) Write(p []byte) (n int, err error) {
 	n = len(p)
@@ -176,17 +141,16 @@ func (m *memLogger) Write(p []byte) (n int, err error) {
 	return
 }
 
-func (m *memLogger) ServeHTTP(config config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
+func (m *memLogger) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	conn, err := utils.InitiateWS(config, w, r)
+	conn, err := gpwebsocket.Initiate(w, r)
 	if err != nil {
-		utils.HandleErr(w, r, err)
+		gphttp.ServerError(w, r, err)
 		return
 	}
 
 	logCh := make(chan *logEntryRange)
 	m.connChans.Store(logCh, struct{}{})
 
-	/* trunk-ignore(golangci-lint/errcheck) */
 	defer func() {
 		_ = conn.CloseNow()
 
@@ -197,13 +161,27 @@ func (m *memLogger) ServeHTTP(config config.ConfigInstance, w http.ResponseWrite
 	}()
 
 	if err := m.wsInitial(r.Context(), conn); err != nil {
-		utils.HandleErr(w, r, err)
+		gphttp.ServerError(w, r, err)
 		return
 	}
 
 	m.wsStreamLog(r.Context(), conn, logCh)
 }
 
+func (m *memLogger) events() (logs <-chan []byte, cancel func()) {
+	ch := make(chan []byte)
+	m.notifyLock.Lock()
+	defer m.notifyLock.Unlock()
+	m.listeners.Store(ch, struct{}{})
+
+	return ch, func() {
+		m.notifyLock.Lock()
+		defer m.notifyLock.Unlock()
+		m.listeners.Delete(ch)
+		close(ch)
+	}
+}
+
 func (m *memLogger) writeBytes(ctx context.Context, conn *websocket.Conn, b []byte) error {
 	return conn.Write(ctx, websocket.MessageText, b)
 }