Fix data race for memory logger #1487

This commit is contained in:
Alex X
2025-03-08 14:11:29 +03:00
parent fe2e372997
commit 830e476120
+25 -14
View File
@@ -4,12 +4,13 @@ import (
"io" "io"
"os" "os"
"strings" "strings"
"sync"
"github.com/mattn/go-isatty" "github.com/mattn/go-isatty"
"github.com/rs/zerolog" "github.com/rs/zerolog"
) )
var MemoryLog = newBuffer(16) var MemoryLog = newBuffer()
func GetLogger(module string) zerolog.Logger { func GetLogger(module string) zerolog.Logger {
if s, ok := modules[module]; ok { if s, ok := modules[module]; ok {
@@ -106,15 +107,19 @@ var modules = map[string]string{
"time": zerolog.TimeFormatUnixMs, "time": zerolog.TimeFormatUnixMs,
} }
const chunkSize = 1 << 16 const (
chunkCount = 16
chunkSize = 1 << 16
)
type circularBuffer struct { type circularBuffer struct {
chunks [][]byte chunks [][]byte
r, w int r, w int
mu sync.Mutex
} }
func newBuffer(chunks int) *circularBuffer { func newBuffer() *circularBuffer {
b := &circularBuffer{chunks: make([][]byte, 0, chunks)} b := &circularBuffer{chunks: make([][]byte, 0, chunkCount)}
// create first chunk // create first chunk
b.chunks = append(b.chunks, make([]byte, 0, chunkSize)) b.chunks = append(b.chunks, make([]byte, 0, chunkSize))
return b return b
@@ -123,16 +128,17 @@ func newBuffer(chunks int) *circularBuffer {
func (b *circularBuffer) Write(p []byte) (n int, err error) { func (b *circularBuffer) Write(p []byte) (n int, err error) {
n = len(p) n = len(p)
b.mu.Lock()
// check if chunk has size // check if chunk has size
if len(b.chunks[b.w])+n > chunkSize { if len(b.chunks[b.w])+n > chunkSize {
// increase write chunk index // increase write chunk index
if b.w++; b.w == cap(b.chunks) { if b.w++; b.w == chunkCount {
b.w = 0 b.w = 0
} }
// check overflow // check overflow
if b.r == b.w { if b.r == b.w {
// increase read chunk index // increase read chunk index
if b.r++; b.r == cap(b.chunks) { if b.r++; b.r == chunkCount {
b.r = 0 b.r = 0
} }
} }
@@ -147,29 +153,34 @@ func (b *circularBuffer) Write(p []byte) (n int, err error) {
} }
b.chunks[b.w] = append(b.chunks[b.w], p...) b.chunks[b.w] = append(b.chunks[b.w], p...)
b.mu.Unlock()
return return
} }
func (b *circularBuffer) WriteTo(w io.Writer) (n int64, err error) { func (b *circularBuffer) WriteTo(w io.Writer) (n int64, err error) {
for i := b.r; ; { buf := make([]byte, 0, chunkCount*chunkSize)
var nn int
if nn, err = w.Write(b.chunks[i]); err != nil {
return
}
n += int64(nn)
// use temp buffer inside mutex because w.Write can take some time
b.mu.Lock()
for i := b.r; ; {
buf = append(buf, b.chunks[i]...)
if i == b.w { if i == b.w {
break break
} }
if i++; i == cap(b.chunks) { if i++; i == chunkCount {
i = 0 i = 0
} }
} }
return b.mu.Unlock()
nn, err := w.Write(buf)
return int64(nn), err
} }
func (b *circularBuffer) Reset() { func (b *circularBuffer) Reset() {
b.mu.Lock()
b.chunks[0] = b.chunks[0][:0] b.chunks[0] = b.chunks[0][:0]
b.r = 0 b.r = 0
b.w = 0 b.w = 0
b.mu.Unlock()
} }