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

db/clickhouse.go

@@ -22,6 +22,7 @@ type Pool struct {
 	open              func(*clickhouse.Options) (clickhouse.Conn, error)
 	mu                sync.RWMutex
 	conn              clickhouse.Conn
+	generation        uint64
 	healthy           bool
 	closed            bool
 }
@@ -82,6 +83,7 @@ func (p *Pool) connect(ctx context.Context) error {
 		_ = p.conn.Close()
 	}
 	p.conn = conn
+	p.generation++
 	p.healthy = true
 	p.mu.Unlock()
 	log.Printf("[db] connected and migrated")
@@ -192,11 +194,11 @@ func (p *Pool) watch(ctx context.Context) {
 			return
 		case <-t.C:
 			if p.Healthy() {
-				if conn := p.Get(); conn != nil {
+				if conn, generation := p.GetWithGeneration(); conn != nil {
 					pctx, cancel := context.WithTimeout(ctx, 3*time.Second)
 					if err := conn.Ping(pctx); err != nil {
 						log.Printf("[db] ping failed, marking unhealthy: %v", err)
-						p.MarkUnhealthy()
+						p.MarkUnhealthyGeneration(generation)
 					}
 					cancel()
 				}
@@ -223,6 +225,21 @@ func (p *Pool) MarkUnhealthy() {
 	p.mu.Unlock()
 }
 
+// MarkUnhealthyGeneration only changes health when the failed connection is still current.
+func (p *Pool) MarkUnhealthyGeneration(generation uint64) {
+	p.mu.Lock()
+	if p.generation == generation {
+		p.healthy = false
+	}
+	p.mu.Unlock()
+}
+
+func (p *Pool) Generation() uint64 {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return p.generation
+}
+
 // Get 返回当前连接，可能为 nil。
 func (p *Pool) Get() clickhouse.Conn {
 	p.mu.RLock()
@@ -230,6 +247,12 @@ func (p *Pool) Get() clickhouse.Conn {
 	return p.conn
 }
 
+func (p *Pool) GetWithGeneration() (clickhouse.Conn, uint64) {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return p.conn, p.generation
+}
+
 // Close starts closing the underlying connection and waits within ctx.
 func (p *Pool) Close(ctx context.Context) error {
 	p.mu.Lock()

db/clickhouse_test.go

@@ -214,6 +214,26 @@ func TestClosePreventsConcurrentConnectFromPublishing(t *testing.T) {
 	}
 }
 
+func TestOldGenerationCannotMarkReplacementUnhealthy(t *testing.T) {
+	pool := &Pool{conn: newFakeClickHouseConn(), healthy: true, generation: 1}
+	oldGeneration := pool.Generation()
+
+	pool.mu.Lock()
+	pool.conn = newFakeClickHouseConn()
+	pool.generation++
+	pool.healthy = true
+	pool.mu.Unlock()
+
+	pool.MarkUnhealthyGeneration(oldGeneration)
+	if !pool.Healthy() {
+		t.Fatal("old connection failure marked replacement connection unhealthy")
+	}
+	pool.MarkUnhealthyGeneration(pool.Generation())
+	if pool.Healthy() {
+		t.Fatal("current connection failure did not mark pool unhealthy")
+	}
+}
+
 func assertNonDestructiveMigration(t *testing.T, statements []string) {
 	t.Helper()
 	if len(statements) == 0 || !strings.Contains(strings.ToUpper(statements[0]), "CREATE TABLE IF NOT EXISTS") {