diff --git a/logger/queue.go b/logger/queue.go index fbc4b4f..f831820 100644 --- a/logger/queue.go +++ b/logger/queue.go @@ -54,9 +54,14 @@ type FlushResult struct { Err error } +type queuedEntry struct { + entry *LogEntry + size int64 +} + // Queue is an asynchronous, bounded logger. Submit never waits for database work. type Queue struct { - ch chan *LogEntry + ch chan queuedEntry pool *db.Pool batchSize int batchInterval time.Duration @@ -101,7 +106,7 @@ func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval ti batchInterval = time.Second } return &Queue{ - ch: make(chan *LogEntry, queueSize), + ch: make(chan queuedEntry, queueSize), pool: pool, batchSize: batchSize, batchInterval: batchInterval, @@ -210,7 +215,7 @@ func (q *Queue) Submit(e *LogEntry) { return } select { - case q.ch <- e: + case q.ch <- queuedEntry{entry: e, size: size}: q.enq.Add(1) default: q.release(size) @@ -252,15 +257,15 @@ func (q *Queue) release(size int64) { } func (q *Queue) discardQueued() { - for e := range q.ch { - q.release(EstimatedBytes(e)) + for item := range q.ch { + q.release(item.size) q.failed.Add(1) } } func (q *Queue) run(workCtx context.Context) { defer q.wg.Done() - batch := make([]*LogEntry, 0, q.batchSize) + batch := make([]queuedEntry, 0, q.batchSize) ticker := time.NewTicker(q.batchInterval) defer ticker.Stop() @@ -268,9 +273,13 @@ func (q *Queue) run(workCtx context.Context) { if len(batch) == 0 { return } - q.flush(q.flushContext(workCtx), batch) - for _, e := range batch { - q.release(EstimatedBytes(e)) + entries := make([]*LogEntry, len(batch)) + for i, item := range batch { + entries[i] = item.entry + } + q.flush(q.flushContext(workCtx), entries) + for _, item := range batch { + q.release(item.size) } clear(batch) batch = batch[:0] @@ -278,12 +287,12 @@ func (q *Queue) run(workCtx context.Context) { for { select { - case e, ok := <-q.ch: + case item, ok := <-q.ch: if !ok { flush() return } - batch = append(batch, e) + batch = append(batch, item) if len(batch) >= q.batchSize { flush() } diff --git a/tests/logger/queue_test.go b/tests/logger/queue_test.go index 3c72c90..2dd6e43 100644 --- a/tests/logger/queue_test.go +++ b/tests/logger/queue_test.go @@ -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) { entry := &logger.LogEntry{RequestID: "queued"} q := logger.NewQueue(nil, 32, 32, 1, time.Hour, 32*logger.EstimatedBytes(entry))