fix: enforce logger entry budget across batches

This commit is contained in:
MiMoCode
2026-07-10 18:30:46 +08:00
parent c4a985e9ca
commit 044bed77e6
2 changed files with 36 additions and 0 deletions
+15
View File
@@ -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() {
+21
View File
@@ -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))