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 batchSize int
batchInterval time.Duration batchInterval time.Duration
workers int workers int
entryBudget int64
byteBudget int64 byteBudget int64
dropped atomic.Uint64 dropped atomic.Uint64
failed atomic.Uint64 failed atomic.Uint64
ambiguous atomic.Uint64 ambiguous atomic.Uint64
enq atomic.Uint64 enq atomic.Uint64
entries atomic.Int64
bytes atomic.Int64 bytes atomic.Int64
mu sync.RWMutex mu sync.RWMutex
@@ -104,6 +106,7 @@ func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval ti
batchSize: batchSize, batchSize: batchSize,
batchInterval: batchInterval, batchInterval: batchInterval,
workers: workers, workers: workers,
entryBudget: int64(queueSize),
byteBudget: budget, byteBudget: budget,
done: make(chan struct{}), done: make(chan struct{}),
} }
@@ -217,6 +220,16 @@ func (q *Queue) Submit(e *LogEntry) {
} }
func (q *Queue) reserve(size int64) bool { 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 { if q.byteBudget <= 0 {
q.bytes.Add(size) q.bytes.Add(size)
return true return true
@@ -224,6 +237,7 @@ func (q *Queue) reserve(size int64) bool {
for { for {
used := q.bytes.Load() used := q.bytes.Load()
if size > q.byteBudget-used { if size > q.byteBudget-used {
q.entries.Add(-1)
return false return false
} }
if q.bytes.CompareAndSwap(used, used+size) { if q.bytes.CompareAndSwap(used, used+size) {
@@ -234,6 +248,7 @@ func (q *Queue) reserve(size int64) bool {
func (q *Queue) release(size int64) { func (q *Queue) release(size int64) {
q.bytes.Add(-size) q.bytes.Add(-size)
q.entries.Add(-1)
} }
func (q *Queue) discardQueued() { 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) { func TestQueueStopBeforeStartReleasesBudget(t *testing.T) {
entry := &logger.LogEntry{RequestID: "queued"} entry := &logger.LogEntry{RequestID: "queued"}
q := logger.NewQueue(nil, 2, 2, 1, time.Hour, logger.EstimatedBytes(entry)) q := logger.NewQueue(nil, 2, 2, 1, time.Hour, logger.EstimatedBytes(entry))