修复代理与日志链路的可靠性问题

logger/queue.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 
@@ -46,11 +47,16 @@ func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string)
 }
 
 type poolBackend struct {
-	pool *db.Pool
+	pool       *db.Pool
+	conn       driver.Conn
+	generation uint64
 }
 
 func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, error) {
-	conn := b.pool.Get()
+	conn := b.conn
+	if conn == nil {
+		conn, _ = b.pool.GetWithGeneration()
+	}
 	if conn == nil {
 		return nil, errors.New("pool nil")
 	}
@@ -58,7 +64,14 @@ func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, err
 }
 
 func (b poolBackend) Healthy() bool  { return b.pool.Healthy() }
-func (b poolBackend) MarkUnhealthy() { b.pool.MarkUnhealthy() }
+func (b poolBackend) MarkUnhealthy() { b.pool.MarkUnhealthyGeneration(b.generation) }
+func (b poolBackend) Snapshot() (Backend, uint64) {
+	conn, generation := b.pool.GetWithGeneration()
+	return poolBackend{pool: b.pool, conn: conn, generation: generation}, generation
+}
+func (b poolBackend) MarkUnhealthyGeneration(generation uint64) {
+	b.pool.MarkUnhealthyGeneration(generation)
+}
 
 type Stats struct {
 	Enqueued  uint64
@@ -153,8 +166,8 @@ func EstimatedBytes(e *LogEntry) int64 {
 	if e == nil {
 		return 0
 	}
-	return int64(len(e.RequestID) + len(e.Method) + len(e.Path) + len(e.Query) + len(e.ClientIP) + len(e.Error) +
-		len(e.RequestHeaders) + len(e.RequestBody) + len(e.ResponseHeaders) + len(e.ResponseBody))
+	return int64(unsafe.Sizeof(*e)+unsafe.Sizeof(queuedEntry{})) + int64(len(e.RequestID)+len(e.Method)+len(e.Path)+len(e.Query)+len(e.ClientIP)+len(e.Error)+
+		len(e.RequestHeaders)+len(e.RequestBody)+len(e.ResponseHeaders)+len(e.ResponseBody))
 }
 
 func (q *Queue) Stats() Stats {
@@ -195,7 +208,6 @@ func (q *Queue) Stop(contexts ...context.Context) error {
 	if len(contexts) > 0 && contexts[0] != nil {
 		ctx = contexts[0]
 	}
-
 	q.mu.Lock()
 	shutdownOwner := false
 	if !q.stopped {
@@ -226,7 +238,12 @@ func (q *Queue) Stop(contexts ...context.Context) error {
 	case <-ctx.Done():
 		if shutdownOwner && workCancel != nil {
 			workCancel()
-			<-done
+			timer := time.NewTimer(100 * time.Millisecond)
+			defer timer.Stop()
+			select {
+			case <-done:
+			case <-timer.C:
+			}
 		}
 		return ctx.Err()
 	}
@@ -238,8 +255,7 @@ func (q *Queue) Submit(e *LogEntry) {
 		q.dropped.Add(1)
 		return
 	}
-	entry := cloneLogEntry(e)
-	size := EstimatedBytes(entry)
+	size := EstimatedBytes(e)
 	if !q.mu.TryRLock() {
 		q.dropped.Add(1)
 		return
@@ -249,6 +265,7 @@ func (q *Queue) Submit(e *LogEntry) {
 		q.dropped.Add(1)
 		return
 	}
+	entry := cloneLogEntry(e)
 	select {
 	case q.ch <- queuedEntry{entry: entry, size: size}:
 		q.enq.Add(1)
@@ -367,16 +384,19 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
 
 	retry := entries
 	var lastErr error
+	var failedGeneration uint64
 	for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ {
 		if err := ctx.Err(); err != nil {
 			lastErr = err
 			break
 		}
-		result := Flush(ctx, q.backend, retry)
+		attemptBackend, generation := backendSnapshot(q.backend)
+		failedGeneration = generation
+		result := Flush(ctx, attemptBackend, retry)
 		q.failed.Add(uint64(result.Failed))
 		q.ambiguous.Add(uint64(result.Ambiguous))
 		if result.Ambiguous > 0 {
-			q.backend.MarkUnhealthy()
+			markBackendUnhealthy(q.backend, failedGeneration)
 		}
 		lastErr = result.Err
 		retry = result.Retry
@@ -391,11 +411,31 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
 
 	if len(retry) > 0 {
 		q.failed.Add(uint64(len(retry)))
-		q.backend.MarkUnhealthy()
+		markBackendUnhealthy(q.backend, failedGeneration)
 		log.Printf("[logger] giving up %d retry-safe rows after %d attempts: %v", len(retry), maxAttempts, lastErr)
 	}
 }
 
+type generationBackend interface {
+	Snapshot() (Backend, uint64)
+	MarkUnhealthyGeneration(uint64)
+}
+
+func backendSnapshot(backend Backend) (Backend, uint64) {
+	if versioned, ok := backend.(generationBackend); ok {
+		return versioned.Snapshot()
+	}
+	return backend, 0
+}
+
+func markBackendUnhealthy(backend Backend, generation uint64) {
+	if versioned, ok := backend.(generationBackend); ok {
+		versioned.MarkUnhealthyGeneration(generation)
+		return
+	}
+	backend.MarkUnhealthy()
+}
+
 func waitBackoff(ctx context.Context, attempt int) bool {
 	wait := BaseBackoff
 	for i := 1; i < attempt; i++ {