From 044bed77e6755d69d3379335efab71fbb4efad5c Mon Sep 17 00:00:00 2001
From: MiMoCode <mimocode@localhost>
Date: Fri, 10 Jul 2026 18:30:46 +0800
Subject: [PATCH] fix: enforce logger entry budget across batches

---
 logger/queue.go            | 15 +++++++++++++++
 tests/logger/queue_test.go | 21 +++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/logger/queue.go b/logger/queue.go
index 0aeb581..fbc4b4f 100644
--- a/logger/queue.go
+++ b/logger/queue.go
@@ -61,12 +61,14 @@ type Queue struct {
 	batchSize     int
 	batchInterval time.Duration
 	workers       int
+	entryBudget   int64
 	byteBudget    int64
 
 	dropped   atomic.Uint64
 	failed    atomic.Uint64
 	ambiguous atomic.Uint64
 	enq       atomic.Uint64
+	entries   atomic.Int64
 	bytes     atomic.Int64
 
 	mu             sync.RWMutex
@@ -104,6 +106,7 @@ func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval ti
 		batchSize:     batchSize,
 		batchInterval: batchInterval,
 		workers:       workers,
+		entryBudget:   int64(queueSize),
 		byteBudget:    budget,
 		done:          make(chan struct{}),
 	}
@@ -217,6 +220,16 @@ func (q *Queue) Submit(e *LogEntry) {
 }
 
 func (q *Queue) reserve(size int64) bool {
+	for {
+		used := q.entries.Load()
+		if used >= q.entryBudget {
+			return false
+		}
+		if q.entries.CompareAndSwap(used, used+1) {
+			break
+		}
+	}
+
 	if q.byteBudget <= 0 {
 		q.bytes.Add(size)
 		return true
@@ -224,6 +237,7 @@ func (q *Queue) reserve(size int64) bool {
 	for {
 		used := q.bytes.Load()
 		if size > q.byteBudget-used {
+			q.entries.Add(-1)
 			return false
 		}
 		if q.bytes.CompareAndSwap(used, used+size) {
@@ -234,6 +248,7 @@ func (q *Queue) reserve(size int64) bool {
 
 func (q *Queue) release(size int64) {
 	q.bytes.Add(-size)
+	q.entries.Add(-1)
 }
 
 func (q *Queue) discardQueued() {
diff --git a/tests/logger/queue_test.go b/tests/logger/queue_test.go
index 6d77587..3c72c90 100644
--- a/tests/logger/queue_test.go
+++ b/tests/logger/queue_test.go
@@ -231,6 +231,27 @@ func TestQueueByteBudgetReleasedAfterFinalDiscard(t *testing.T) {
 	}
 }
 
+func TestQueueEntryBudgetIncludesWorkerBatch(t *testing.T) {
+	q := logger.NewQueue(nil, 1, 2, 1, time.Hour)
+	q.Submit(&logger.LogEntry{RequestID: "batched"})
+	q.Start(context.Background())
+
+	deadline := time.Now().Add(time.Second)
+	for q.Stats().Enqueued == 1 && time.Now().Before(deadline) {
+		q.Submit(&logger.LogEntry{RequestID: "queued"})
+		time.Sleep(time.Millisecond)
+	}
+	if stats := q.Stats(); stats.Enqueued != 1 {
+		t.Fatalf("stats=%+v; entry budget allowed channel plus worker batch", stats)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if err := q.Stop(ctx); err != nil {
+		t.Fatalf("Stop: %v", err)
+	}
+}
+
 func TestQueueStopBeforeStartReleasesBudget(t *testing.T) {
 	entry := &logger.LogEntry{RequestID: "queued"}
 	q := logger.NewQueue(nil, 2, 2, 1, time.Hour, logger.EstimatedBytes(entry))