test: prove logger shutdown deadline behavior

logger/queue.go

@@ -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)
 		}