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