From b7442ab51407665a1a043db8eb1d5ab4936129d2 Mon Sep 17 00:00:00 2001
From: MiMoCode <mimocode@localhost>
Date: Fri, 10 Jul 2026 19:48:06 +0800
Subject: [PATCH] fix: prevent reconnect after pool close

---
 db/clickhouse.go      | 17 ++++++++++++
 db/clickhouse_test.go | 62 ++++++++++++++++++++++++++++++++++++-------
 2 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/db/clickhouse.go b/db/clickhouse.go
index 69778f9..a16d5c3 100644
--- a/db/clickhouse.go
+++ b/db/clickhouse.go
@@ -2,6 +2,7 @@ package db
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"net"
@@ -22,8 +23,11 @@ type Pool struct {
 	mu                sync.RWMutex
 	conn              clickhouse.Conn
 	healthy           bool
+	closed            bool
 }
 
+var errPoolClosed = errors.New("db pool is closed")
+
 // NewPool 创建 Pool。即使首次连接失败也返回非 nil 实例，后台会持续重试。
 func NewPool(ctx context.Context, dsn string, reconnectInterval time.Duration) *Pool {
 	p := &Pool{dsn: dsn, reconnectInterval: reconnectInterval}
@@ -35,6 +39,13 @@ func NewPool(ctx context.Context, dsn string, reconnectInterval time.Duration) *
 }
 
 func (p *Pool) connect(ctx context.Context) error {
+	p.mu.RLock()
+	closed := p.closed
+	p.mu.RUnlock()
+	if closed {
+		return errPoolClosed
+	}
+
 	cctx, cancel := context.WithTimeout(ctx, 5*time.Second)
 	defer cancel()
 	opts, err := ClickHouseOptions(p.dsn)
@@ -62,6 +73,11 @@ func (p *Pool) connect(ctx context.Context) error {
 		return err
 	}
 	p.mu.Lock()
+	if p.closed {
+		p.mu.Unlock()
+		_ = conn.Close()
+		return errPoolClosed
+	}
 	if p.conn != nil {
 		_ = p.conn.Close()
 	}
@@ -217,6 +233,7 @@ func (p *Pool) Get() clickhouse.Conn {
 // Close starts closing the underlying connection and waits within ctx.
 func (p *Pool) Close(ctx context.Context) error {
 	p.mu.Lock()
+	p.closed = true
 	conn := p.conn
 	p.conn = nil
 	p.healthy = false
diff --git a/db/clickhouse_test.go b/db/clickhouse_test.go
index 5f23d73..89496ca 100644
--- a/db/clickhouse_test.go
+++ b/db/clickhouse_test.go
@@ -183,6 +183,37 @@ func TestInitialConnectWithIncompatibleSchemaStaysUnhealthy(t *testing.T) {
 	assertNonDestructiveMigration(t, candidate.statements)
 }
 
+func TestClosePreventsConcurrentConnectFromPublishing(t *testing.T) {
+	old := newFakeClickHouseConn()
+	pingStarted := make(chan struct{})
+	pingRelease := make(chan struct{})
+	candidate := newFakeClickHouseConn().withBlockedPing(pingStarted, pingRelease)
+	pool := &Pool{dsn: "localhost:9000", conn: old, healthy: true}
+	pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil }
+
+	connectDone := make(chan error, 1)
+	go func() { connectDone <- pool.connect(context.Background()) }()
+	<-pingStarted
+
+	if err := pool.Close(context.Background()); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	close(pingRelease)
+	if err := <-connectDone; err == nil {
+		t.Fatal("connect succeeded after pool was closed")
+	}
+
+	if pool.Get() != nil || pool.Healthy() {
+		t.Fatalf("closed pool connection = %T, healthy = %v; want nil, false", pool.Get(), pool.Healthy())
+	}
+	if old.closeCalls != 1 {
+		t.Fatalf("old connection close calls = %d, want 1", old.closeCalls)
+	}
+	if candidate.closeCalls != 1 {
+		t.Fatalf("candidate close calls = %d, want 1", candidate.closeCalls)
+	}
+}
+
 func assertNonDestructiveMigration(t *testing.T, statements []string) {
 	t.Helper()
 	if len(statements) == 0 || !strings.Contains(strings.ToUpper(statements[0]), "CREATE TABLE IF NOT EXISTS") {
@@ -197,12 +228,14 @@ func assertNonDestructiveMigration(t *testing.T, statements []string) {
 }
 
 type fakeClickHouseConn struct {
-	rows       *fakeRows
-	row        driver.Row
-	execErr    error
-	queryErr   error
-	statements []string
-	closeCalls int
+	rows        *fakeRows
+	row         driver.Row
+	execErr     error
+	queryErr    error
+	statements  []string
+	closeCalls  int
+	pingStarted chan struct{}
+	pingRelease chan struct{}
 }
 
 func newFakeClickHouseConn() *fakeClickHouseConn {
@@ -229,6 +262,11 @@ func (c *fakeClickHouseConn) withColumnType(index int, typ string) *fakeClickHou
 	c.rows.values[index][1] = typ
 	return c
 }
+func (c *fakeClickHouseConn) withBlockedPing(started, release chan struct{}) *fakeClickHouseConn {
+	c.pingStarted = started
+	c.pingRelease = release
+	return c
+}
 
 func (c *fakeClickHouseConn) Contributors() []string                            { return nil }
 func (c *fakeClickHouseConn) ServerVersion() (*driver.ServerVersion, error)     { return nil, nil }
@@ -249,9 +287,15 @@ func (c *fakeClickHouseConn) Exec(_ context.Context, query string, _ ...any) err
 	return c.execErr
 }
 func (c *fakeClickHouseConn) AsyncInsert(context.Context, string, bool, ...any) error { return nil }
-func (c *fakeClickHouseConn) Ping(context.Context) error                              { return nil }
-func (c *fakeClickHouseConn) Stats() driver.Stats                                     { return driver.Stats{} }
-func (c *fakeClickHouseConn) Close() error                                            { c.closeCalls++; return nil }
+func (c *fakeClickHouseConn) Ping(context.Context) error {
+	if c.pingStarted != nil {
+		close(c.pingStarted)
+		<-c.pingRelease
+	}
+	return nil
+}
+func (c *fakeClickHouseConn) Stats() driver.Stats { return driver.Stats{} }
+func (c *fakeClickHouseConn) Close() error        { c.closeCalls++; return nil }
 
 type fakeRows struct {
 	values    [][]string