fix: enforce logger entry budget across batches
This commit is contained in:
@@ -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() {
|
||||
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user