m1saka/token_thief
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test: prove logger shutdown deadline behavior

Browse Source
This commit is contained in:
MiMoCode
2026-07-10 18:36:40 +08:00
parent 8eaab5ccfb
commit 3292eb38fc
2 changed files with 170 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
logger/queue.go
+38 -12
View File
@@ -29,6 +29,13 @@ type BatchPreparer interface {
PrepareBatch(ctx context.Context, query string) (Batch, error)
}
// Backend is the database contract required by the asynchronous queue.
type Backend interface {
BatchPreparer
Healthy() bool
MarkUnhealthy()
}
type clickHouseBatchPreparer struct {
conn driver.Conn
}
@@ -38,6 +45,21 @@ func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string)
return p.conn.PrepareBatch(ctx, query)
}
type poolBackend struct {
pool *db.Pool
}
func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, error) {
conn := b.pool.Get()
if conn == nil {
return nil, errors.New("pool nil")
}
return clickHouseBatchPreparer{conn: conn}.PrepareBatch(ctx, query)
}
func (b poolBackend) Healthy() bool { return b.pool.Healthy() }
func (b poolBackend) MarkUnhealthy() { b.pool.MarkUnhealthy() }
type Stats struct {
Enqueued uint64
Dropped uint64
@@ -62,7 +84,7 @@ type queuedEntry struct {
// Queue is an asynchronous, bounded logger. Submit never waits for database work.
type Queue struct {
ch chan queuedEntry
pool *db.Pool
backend Backend
batchSize int
batchInterval time.Duration
workers int
@@ -89,6 +111,15 @@ type Queue struct {
// NewQueue accepts an optional byte budget. A non-positive or omitted budget disables byte limiting.
func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval time.Duration, byteBudget ...int64) *Queue {
var backend Backend
if pool != nil {
backend = poolBackend{pool: pool}
}
return NewQueueWithBackend(backend, queueSize, batchSize, workers, batchInterval, byteBudget...)
}
// NewQueueWithBackend builds a queue against the database operations used by workers.
func NewQueueWithBackend(backend Backend, queueSize, batchSize, workers int, batchInterval time.Duration, byteBudget ...int64) *Queue {
var budget int64
if len(byteBudget) > 0 {
budget = byteBudget[0]
@@ -107,7 +138,7 @@ func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval ti
}
return &Queue{
ch: make(chan queuedEntry, queueSize),
pool: pool,
backend: backend,
batchSize: batchSize,
batchInterval: batchInterval,
workers: workers,
@@ -316,7 +347,7 @@ func (q *Queue) flushContext(workCtx context.Context) context.Context {
}
func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
if q.pool == nil || !q.pool.Healthy() {
if q.backend == nil || !q.backend.Healthy() {
q.failed.Add(uint64(len(entries)))
return
}
@@ -328,16 +359,11 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
lastErr = err
break
}
conn := q.pool.Get()
if conn == nil {
lastErr = errors.New("pool nil")
break
}
result := Flush(ctx, clickHouseBatchPreparer{conn: conn}, retry)
result := Flush(ctx, q.backend, retry)
q.failed.Add(uint64(result.Failed))
q.ambiguous.Add(uint64(result.Ambiguous))
if result.Ambiguous > 0 {
q.pool.MarkUnhealthy()
q.backend.MarkUnhealthy()
}
lastErr = result.Err
retry = result.Retry
@@ -352,7 +378,7 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
if len(retry) > 0 {
q.failed.Add(uint64(len(retry)))
q.pool.MarkUnhealthy()
q.backend.MarkUnhealthy()
log.Printf("[logger] giving up %d retry-safe rows after %d attempts: %v", len(retry), maxAttempts, lastErr)
}
}
@@ -463,7 +489,7 @@ func (q *Queue) reportLoop(ctx context.Context) {
return
case <-ticker.C:
stats := q.Stats()
healthy := q.pool != nil && q.pool.Healthy()
healthy := q.backend != nil && q.backend.Healthy()
log.Printf("[logger] metrics: enq=%d dropped_queue_full=%d dropped_db_fail=%d ambiguous_send=%d queue_len=%d queue_bytes=%d db_healthy=%v",
stats.Enqueued, stats.Dropped, stats.Failed, stats.Ambiguous, len(q.ch), stats.Bytes, healthy)
}
tests/logger/queue_test.go
+132
View File
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
@@ -43,6 +44,26 @@ type fakePool struct {
scripts []batchScript
}
type queueBackend struct {
prepare func(context.Context, string) (logger.Batch, error)
healthy atomic.Bool
calls atomic.Int32
}
func newQueueBackend(prepare func(context.Context, string) (logger.Batch, error)) *queueBackend {
b := &queueBackend{prepare: prepare}
b.healthy.Store(true)
return b
}
func (b *queueBackend) PrepareBatch(ctx context.Context, query string) (logger.Batch, error) {
b.calls.Add(1)
return b.prepare(ctx, query)
}
func (b *queueBackend) Healthy() bool { return b.healthy.Load() }
func (b *queueBackend) MarkUnhealthy() { b.healthy.Store(false) }
func (p *fakePool) PrepareBatch(_ context.Context, query string) (logger.Batch, error) {
idx := p.calls
p.calls++
@@ -360,6 +381,117 @@ func TestStopIsIndependentFromStartContext(t *testing.T) {
}
}
func TestQueueStopDeadlineCancelsBlockedPrepare(t *testing.T) {
started := make(chan struct{})
backend := newQueueBackend(func(ctx context.Context, _ string) (logger.Batch, error) {
close(started)
<-ctx.Done()
return nil, ctx.Err()
})
q := logger.NewQueueWithBackend(backend, 1, 1, 1, time.Hour)
q.Start(context.Background())
q.Submit(&logger.LogEntry{RequestID: "blocked-prepare"})
<-started
ctx, cancel := context.WithTimeout(context.Background(), 40*time.Millisecond)
defer cancel()
startedAt := time.Now()
if err := q.Stop(ctx); !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("Stop error=%v want deadline exceeded", err)
}
if elapsed := time.Since(startedAt); elapsed > 250*time.Millisecond {
t.Fatalf("Stop exceeded total deadline: %v", elapsed)
}
}
func TestQueueStopDeadlineCancelsBlockedSend(t *testing.T) {
started := make(chan struct{})
backend := newQueueBackend(func(ctx context.Context, _ string) (logger.Batch, error) {
return &blockingSendBatch{ctx: ctx, started: started}, nil
})
q := logger.NewQueueWithBackend(backend, 1, 1, 1, time.Hour)
q.Start(context.Background())
q.Submit(&logger.LogEntry{RequestID: "blocked-send"})
<-started
ctx, cancel := context.WithTimeout(context.Background(), 40*time.Millisecond)
defer cancel()
if err := q.Stop(ctx); !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("Stop error=%v want deadline exceeded", err)
}
deadline := time.Now().Add(250 * time.Millisecond)
for q.Stats().Bytes != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if stats := q.Stats(); stats.Bytes != 0 || stats.Ambiguous != 1 || stats.Failed != 1 {
t.Fatalf("stats after canceled Send=%+v", stats)
}
if backend.Healthy() {
t.Fatal("Send failure did not mark backend unhealthy")
}
}
func TestQueueStopDeadlineCancelsPrepareRetryBackoff(t *testing.T) {
backend := newQueueBackend(func(context.Context, string) (logger.Batch, error) {
return nil, errors.New("prepare failed")
})
q := logger.NewQueueWithBackend(backend, 1, 1, 1, time.Hour)
q.Start(context.Background())
q.Submit(&logger.LogEntry{RequestID: "backoff"})
deadline := time.Now().Add(time.Second)
for backend.calls.Load() == 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
ctx, cancel := context.WithTimeout(context.Background(), 40*time.Millisecond)
defer cancel()
if err := q.Stop(ctx); !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("Stop error=%v want deadline exceeded", err)
}
time.Sleep(20 * time.Millisecond)
if calls := backend.calls.Load(); calls != 1 {
t.Fatalf("PrepareBatch calls=%d want 1 before deadline cancels backoff", calls)
}
}
func TestQueueExhaustsPrepareRetriesAndMarksBackendUnhealthy(t *testing.T) {
backend := newQueueBackend(func(context.Context, string) (logger.Batch, error) {
return nil, errors.New("prepare failed")
})
q := logger.NewQueueWithBackend(backend, 2, 2, 1, time.Hour)
q.Start(context.Background())
q.Submit(&logger.LogEntry{RequestID: "a"})
q.Submit(&logger.LogEntry{RequestID: "b"})
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
if err := q.Stop(ctx); err != nil {
t.Fatalf("Stop: %v", err)
}
if calls := backend.calls.Load(); calls != 3 {
t.Fatalf("PrepareBatch calls=%d want 3", calls)
}
if stats := q.Stats(); stats.Failed != 2 || stats.Ambiguous != 0 || stats.Bytes != 0 {
t.Fatalf("stats=%+v", stats)
}
if backend.Healthy() {
t.Fatal("exhausted Prepare retries did not mark backend unhealthy")
}
}
type blockingSendBatch struct {
ctx context.Context
started chan struct{}
}
func (*blockingSendBatch) Append(...any) error { return nil }
func (b *blockingSendBatch) Send() error {
close(b.started)
<-b.ctx.Done()
return b.ctx.Err()
}
func (*blockingSendBatch) Abort() error { return nil }
func contains(s, substr string) bool {
for i := 0; i+len(substr) <= len(s); i++ {
if s[i:i+len(substr)] == substr {