fix: preserve logger byte reservations

This commit is contained in:
MiMoCode
2026-07-10 18:33:27 +08:00
parent 044bed77e6
commit 8eaab5ccfb
2 changed files with 46 additions and 11 deletions
+20 -11
View File
@@ -54,9 +54,14 @@ type FlushResult struct {
Err error Err error
} }
type queuedEntry struct {
entry *LogEntry
size int64
}
// Queue is an asynchronous, bounded logger. Submit never waits for database work. // Queue is an asynchronous, bounded logger. Submit never waits for database work.
type Queue struct { type Queue struct {
ch chan *LogEntry ch chan queuedEntry
pool *db.Pool pool *db.Pool
batchSize int batchSize int
batchInterval time.Duration batchInterval time.Duration
@@ -101,7 +106,7 @@ func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval ti
batchInterval = time.Second batchInterval = time.Second
} }
return &Queue{ return &Queue{
ch: make(chan *LogEntry, queueSize), ch: make(chan queuedEntry, queueSize),
pool: pool, pool: pool,
batchSize: batchSize, batchSize: batchSize,
batchInterval: batchInterval, batchInterval: batchInterval,
@@ -210,7 +215,7 @@ func (q *Queue) Submit(e *LogEntry) {
return return
} }
select { select {
case q.ch <- e: case q.ch <- queuedEntry{entry: e, size: size}:
q.enq.Add(1) q.enq.Add(1)
default: default:
q.release(size) q.release(size)
@@ -252,15 +257,15 @@ func (q *Queue) release(size int64) {
} }
func (q *Queue) discardQueued() { func (q *Queue) discardQueued() {
for e := range q.ch { for item := range q.ch {
q.release(EstimatedBytes(e)) q.release(item.size)
q.failed.Add(1) q.failed.Add(1)
} }
} }
func (q *Queue) run(workCtx context.Context) { func (q *Queue) run(workCtx context.Context) {
defer q.wg.Done() defer q.wg.Done()
batch := make([]*LogEntry, 0, q.batchSize) batch := make([]queuedEntry, 0, q.batchSize)
ticker := time.NewTicker(q.batchInterval) ticker := time.NewTicker(q.batchInterval)
defer ticker.Stop() defer ticker.Stop()
@@ -268,9 +273,13 @@ func (q *Queue) run(workCtx context.Context) {
if len(batch) == 0 { if len(batch) == 0 {
return return
} }
q.flush(q.flushContext(workCtx), batch) entries := make([]*LogEntry, len(batch))
for _, e := range batch { for i, item := range batch {
q.release(EstimatedBytes(e)) entries[i] = item.entry
}
q.flush(q.flushContext(workCtx), entries)
for _, item := range batch {
q.release(item.size)
} }
clear(batch) clear(batch)
batch = batch[:0] batch = batch[:0]
@@ -278,12 +287,12 @@ func (q *Queue) run(workCtx context.Context) {
for { for {
select { select {
case e, ok := <-q.ch: case item, ok := <-q.ch:
if !ok { if !ok {
flush() flush()
return return
} }
batch = append(batch, e) batch = append(batch, item)
if len(batch) >= q.batchSize { if len(batch) >= q.batchSize {
flush() flush()
} }
+26
View File
@@ -267,6 +267,32 @@ func TestQueueStopBeforeStartReleasesBudget(t *testing.T) {
} }
} }
func TestQueueReleasesSubmittedSizeAfterEntryMutation(t *testing.T) {
for _, tc := range []struct {
name string
mutate func(*logger.LogEntry)
}{
{name: "smaller", mutate: func(entry *logger.LogEntry) { entry.RequestBody = nil }},
{name: "larger", mutate: func(entry *logger.LogEntry) { entry.RequestBody = []byte("much larger body") }},
} {
t.Run(tc.name, func(t *testing.T) {
entry := &logger.LogEntry{RequestID: "queued", RequestBody: []byte("body")}
q := logger.NewQueue(nil, 2, 2, 1, time.Hour, logger.EstimatedBytes(entry))
q.Submit(entry)
tc.mutate(entry)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := q.Stop(ctx); err != nil {
t.Fatalf("Stop: %v", err)
}
if stats := q.Stats(); stats.Bytes != 0 || stats.Failed != 1 {
t.Fatalf("stats=%+v", stats)
}
})
}
}
func TestQueueCanceledStopEventuallyReleasesAllBudget(t *testing.T) { func TestQueueCanceledStopEventuallyReleasesAllBudget(t *testing.T) {
entry := &logger.LogEntry{RequestID: "queued"} entry := &logger.LogEntry{RequestID: "queued"}
q := logger.NewQueue(nil, 32, 32, 1, time.Hour, 32*logger.EstimatedBytes(entry)) q := logger.NewQueue(nil, 32, 32, 1, time.Hour, 32*logger.EstimatedBytes(entry))