修复日志持久化可靠性与数据库安全默认值

2026-07-13 11:27:34 +08:00
parent db325ec413
commit 6f37165256
10 changed files with 612 additions and 75 deletions
@@ -25,6 +25,16 @@ type Pool struct {
 	generation        uint64
 	healthy           bool
 	closed            bool
+	leases            map[uint64]*connectionLease
+	leaseChanged      chan struct{}
+	closing           int
+}
+
+type connectionLease struct {
+	conn    clickhouse.Conn
+	active  int
+	retired bool
+	closed  bool
 }

 var errPoolClosed = errors.New("db pool is closed")
@@ -79,13 +89,16 @@ func (p *Pool) connect(ctx context.Context) error {
 		_ = conn.Close()
 		return errPoolClosed
 	}
-	if p.conn != nil {
-		_ = p.conn.Close()
-	}
+	old := p.retireCurrentLocked()
 	p.conn = conn
 	p.generation++
+	p.ensureLeaseLocked(p.generation, conn)
 	p.healthy = true
 	p.mu.Unlock()
+	if old != nil {
+		_ = old.Close()
+		p.finishClosing()
+	}
 	log.Printf("[db] connected and migrated")
 	return nil
 }
@@ -194,13 +207,14 @@ func (p *Pool) watch(ctx context.Context) {
 			return
 		case <-t.C:
 			if p.Healthy() {
-				if conn, generation := p.GetWithGeneration(); conn != nil {
+				if conn, generation, release := p.Acquire(); 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.MarkUnhealthyGeneration(generation)
 					}
 					cancel()
+					release()
 				}
 				continue
 			}
@@ -253,24 +267,126 @@ func (p *Pool) GetWithGeneration() (clickhouse.Conn, uint64) {
 	return p.conn, p.generation
 }

+// Acquire pins the current connection until release is called.
+func (p *Pool) Acquire() (clickhouse.Conn, uint64, func()) {
+	p.mu.Lock()
+	if p.closed || p.conn == nil {
+		p.mu.Unlock()
+		return nil, p.generation, nil
+	}
+	generation := p.generation
+	lease := p.ensureLeaseLocked(generation, p.conn)
+	lease.active++
+	p.mu.Unlock()
+
+	var once sync.Once
+	return lease.conn, generation, func() {
+		once.Do(func() { p.release(generation) })
+	}
+}
+
+func (p *Pool) ensureLeaseLocked(generation uint64, conn clickhouse.Conn) *connectionLease {
+	if p.leases == nil {
+		p.leases = make(map[uint64]*connectionLease)
+	}
+	lease := p.leases[generation]
+	if lease == nil {
+		lease = &connectionLease{conn: conn}
+		p.leases[generation] = lease
+	}
+	return lease
+}
+
+func (p *Pool) notifyLeaseChangedLocked() {
+	if p.leaseChanged != nil {
+		close(p.leaseChanged)
+	}
+	p.leaseChanged = make(chan struct{})
+}
+
+func (p *Pool) retireCurrentLocked() clickhouse.Conn {
+	if p.conn == nil {
+		return nil
+	}
+	lease := p.ensureLeaseLocked(p.generation, p.conn)
+	lease.retired = true
+	if lease.active != 0 || lease.closed {
+		return nil
+	}
+	lease.closed = true
+	delete(p.leases, p.generation)
+	p.closing++
+	return lease.conn
+}
+
+func (p *Pool) finishClosing() {
+	p.mu.Lock()
+	p.closing--
+	p.notifyLeaseChangedLocked()
+	p.mu.Unlock()
+}
+
+func (p *Pool) release(generation uint64) {
+	p.mu.Lock()
+	lease := p.leases[generation]
+	if lease == nil {
+		p.mu.Unlock()
+		return
+	}
+	lease.active--
+	var conn clickhouse.Conn
+	if lease.active == 0 && lease.retired && !lease.closed {
+		lease.closed = true
+		conn = lease.conn
+		p.closing++
+	}
+	p.mu.Unlock()
+	if conn != nil {
+		_ = conn.Close()
+		p.mu.Lock()
+		delete(p.leases, generation)
+		p.closing--
+		p.notifyLeaseChangedLocked()
+		p.mu.Unlock()
+	}
+}
+
 // 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
+	conn := p.retireCurrentLocked()
 	p.conn = nil
 	p.healthy = false
 	p.mu.Unlock()
-	if conn == nil {
-		return nil
+	var closeErr error
+	if conn != nil {
+		done := make(chan error, 1)
+		go func() {
+			done <- conn.Close()
+			p.finishClosing()
+		}()
+		select {
+		case closeErr = <-done:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
 	}
-
-	done := make(chan error, 1)
-	go func() { done <- conn.Close() }()
-	select {
-	case err := <-done:
-		return err
-	case <-ctx.Done():
-		return ctx.Err()
+	for {
+		p.mu.Lock()
+		if len(p.leases) == 0 && p.closing == 0 {
+			p.mu.Unlock()
+			return closeErr
+		}
+		if p.leaseChanged == nil {
+			p.leaseChanged = make(chan struct{})
+		}
+		changed := p.leaseChanged
+		p.mu.Unlock()
+		select {
+		case <-changed:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
 	}
 }
@@ -6,6 +6,7 @@ import (
 	"reflect"
 	"strings"
 	"testing"
+	"time"

 	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
@@ -133,6 +134,29 @@ func TestValidateProxyLogsSchema(t *testing.T) {
 	}
 }

+func TestMissingProxyLogColumnsReturnsAdditiveDDL(t *testing.T) {
+	columns := append([]schemaColumn(nil), proxyLogsColumns...)
+	columns = append(columns[:7], columns[8:]...)
+
+	statements, err := missingProxyLogColumns(columns)
+	if err != nil {
+		t.Fatalf("missingProxyLogColumns: %v", err)
+	}
+	want := []string{"ALTER TABLE proxy_logs ADD COLUMN IF NOT EXISTS request_truncated Bool DEFAULT false"}
+	if !reflect.DeepEqual(statements, want) {
+		t.Fatalf("statements=%q want %q", statements, want)
+	}
+}
+
+func TestValidateProxyLogsSchemaAllowsExtraColumns(t *testing.T) {
+	columns := append([]schemaColumn(nil), proxyLogsColumns...)
+	columns = append(columns, schemaColumn{"deployment", "String"})
+	table := schemaTable{"MergeTree", "toYYYYMM(started_at)", "started_at, request_id"}
+	if err := validateProxyLogsSchema(columns, table); err != nil {
+		t.Fatalf("validateProxyLogsSchema: %v", err)
+	}
+}
+
 func TestConnectRejectsUnhealthySchemaWithoutReplacingHealthyConnection(t *testing.T) {
 	tests := []struct {
 		name      string
@@ -183,6 +207,24 @@ func TestInitialConnectWithIncompatibleSchemaStaysUnhealthy(t *testing.T) {
 	assertNonDestructiveMigration(t, candidate.statements)
 }

+func TestMigrationDoesNotAlterIncompatibleTable(t *testing.T) {
+	values := schemaValues()
+	values = append(values[:7], values[8:]...)
+	candidate := newFakeClickHouseConn().withRows(&fakeRows{values: values, scanErrAt: -1})
+	candidate.row = fakeRow{values: []string{"ReplacingMergeTree", "toYYYYMM(started_at)", "started_at, request_id"}}
+	pool := &Pool{dsn: "localhost:9000"}
+	pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil }
+
+	if err := pool.connect(context.Background()); err == nil {
+		t.Fatal("connect succeeded with incompatible table")
+	}
+	for _, statement := range candidate.statements {
+		if strings.Contains(strings.ToUpper(statement), "ALTER TABLE") {
+			t.Fatalf("migration altered incompatible table: %q", statement)
+		}
+	}
+}
+
 func TestClosePreventsConcurrentConnectFromPublishing(t *testing.T) {
 	old := newFakeClickHouseConn()
 	pingStarted := make(chan struct{})
@@ -234,6 +276,85 @@ func TestOldGenerationCannotMarkReplacementUnhealthy(t *testing.T) {
 	}
 }

+func TestConnectionLeaseDefersReplacementCloseUntilRelease(t *testing.T) {
+	old := newFakeClickHouseConn()
+	candidate := newFakeClickHouseConn()
+	pool := &Pool{dsn: "localhost:9000", conn: old, healthy: true, generation: 1}
+	pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil }
+
+	conn, generation, release := pool.Acquire()
+	if conn != old || generation != 1 || release == nil {
+		t.Fatalf("Acquire=(%T,%d,%v), want old,1,release", conn, generation, release != nil)
+	}
+	if err := pool.connect(context.Background()); err != nil {
+		t.Fatalf("connect: %v", err)
+	}
+	if old.closeCalls != 0 {
+		t.Fatalf("leased old connection closed %d times", old.closeCalls)
+	}
+	release()
+	if old.closeCalls != 1 {
+		t.Fatalf("old connection close calls=%d want 1 after release", old.closeCalls)
+	}
+	release()
+	if old.closeCalls != 1 {
+		t.Fatalf("release was not idempotent: close calls=%d", old.closeCalls)
+	}
+}
+
+func TestCloseWaitsForConnectionLease(t *testing.T) {
+	conn := newFakeClickHouseConn()
+	pool := &Pool{conn: conn, healthy: true, generation: 1}
+	_, _, release := pool.Acquire()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if err := pool.Close(ctx); !errors.Is(err, context.Canceled) {
+		t.Fatalf("Close error=%v want canceled", err)
+	}
+	if conn.closeCalls != 0 {
+		t.Fatalf("leased connection closed %d times", conn.closeCalls)
+	}
+	if acquired, _, _ := pool.Acquire(); acquired != nil {
+		t.Fatal("Acquire succeeded after Close")
+	}
+	release()
+	if conn.closeCalls != 1 {
+		t.Fatalf("connection close calls=%d want 1 after release", conn.closeCalls)
+	}
+}
+
+func TestCloseWaitsForRetiredConnectionLease(t *testing.T) {
+	old := newFakeClickHouseConn()
+	current := newFakeClickHouseConn()
+	pool := &Pool{conn: old, healthy: true, generation: 1}
+	_, _, release := pool.Acquire()
+	pool.mu.Lock()
+	pool.retireCurrentLocked()
+	pool.conn = current
+	pool.generation++
+	pool.ensureLeaseLocked(pool.generation, current)
+	pool.mu.Unlock()
+
+	closeDone := make(chan error, 1)
+	go func() { closeDone <- pool.Close(context.Background()) }()
+	select {
+	case err := <-closeDone:
+		t.Fatalf("Close returned before retired lease release: %v", err)
+	case <-time.After(25 * time.Millisecond):
+	}
+	if old.closeCalls != 0 {
+		t.Fatalf("retired leased connection closed %d times", old.closeCalls)
+	}
+	release()
+	if err := <-closeDone; err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	if old.closeCalls != 1 {
+		t.Fatalf("retired connection close calls=%d want 1", old.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") {
@@ -38,43 +38,81 @@ func migrate(ctx context.Context, conn clickhouse.Conn) error {
 		return fmt.Errorf("create proxy_logs: %w", err)
 	}

-	rows, err := conn.Query(ctx, `
-SELECT name, type
-FROM system.columns
-WHERE database = currentDatabase() AND table = 'proxy_logs'
-ORDER BY position`)
+	columns, err := queryProxyLogColumns(ctx, conn)
 	if err != nil {
-		return fmt.Errorf("query proxy_logs columns: %w", err)
+		return err
 	}
-	var columns []schemaColumn
-	for rows.Next() {
-		var column schemaColumn
-		if err := rows.Scan(&column.name, &column.typ); err != nil {
-			rows.Close()
-			return fmt.Errorf("scan proxy_logs columns: %w", err)
+	table, err := queryProxyLogsTable(ctx, conn)
+	if err != nil {
+		return err
+	}
+	if err := validateProxyLogsTable(table); err != nil {
+		return fmt.Errorf("incompatible proxy_logs schema: %w", err)
+	}
+	statements, err := missingProxyLogColumns(columns)
+	if err != nil {
+		return fmt.Errorf("plan proxy_logs migration: %w", err)
+	}
+	for _, statement := range statements {
+		if err := conn.Exec(ctx, statement); err != nil {
+			return fmt.Errorf("alter proxy_logs: %w", err)
 		}
-		columns = append(columns, column)
 	}
-	if err := rows.Err(); err != nil {
-		rows.Close()
-		return fmt.Errorf("read proxy_logs columns: %w", err)
+	if len(statements) > 0 {
+		columns, err = queryProxyLogColumns(ctx, conn)
+		if err != nil {
+			return err
+		}
 	}
-	rows.Close()

+	table, err = queryProxyLogsTable(ctx, conn)
+	if err != nil {
+		return err
+	}
+	if err := validateProxyLogsSchema(columns, table); err != nil {
+		return fmt.Errorf("incompatible proxy_logs schema: %w", err)
+	}
+	return nil
+}
+
+func queryProxyLogsTable(ctx context.Context, conn clickhouse.Conn) (schemaTable, error) {
 	var table schemaTable
-	err = conn.QueryRow(ctx, `
+	err := conn.QueryRow(ctx, `
 SELECT engine, partition_key, sorting_key
 FROM system.tables
 WHERE database = currentDatabase() AND name = 'proxy_logs'`).Scan(
 		&table.engine, &table.partitionKey, &table.sortingKey,
 	)
 	if err != nil {
-		return fmt.Errorf("query proxy_logs table: %w", err)
+		return schemaTable{}, fmt.Errorf("query proxy_logs table: %w", err)
 	}
-	if err := validateProxyLogsSchema(columns, table); err != nil {
-		return fmt.Errorf("incompatible proxy_logs schema: %w", err)
+	return table, nil
+}
+
+func queryProxyLogColumns(ctx context.Context, conn clickhouse.Conn) ([]schemaColumn, error) {
+	rows, err := conn.Query(ctx, `
+SELECT name, type
+FROM system.columns
+WHERE database = currentDatabase() AND table = 'proxy_logs'
+ORDER BY position`)
+	if err != nil {
+		return nil, fmt.Errorf("query proxy_logs columns: %w", err)
 	}
-	return nil
+	var columns []schemaColumn
+	for rows.Next() {
+		var column schemaColumn
+		if err := rows.Scan(&column.name, &column.typ); err != nil {
+			rows.Close()
+			return nil, fmt.Errorf("scan proxy_logs columns: %w", err)
+		}
+		columns = append(columns, column)
+	}
+	if err := rows.Err(); err != nil {
+		rows.Close()
+		return nil, fmt.Errorf("read proxy_logs columns: %w", err)
+	}
+	rows.Close()
+	return columns, nil
 }

 type schemaColumn struct {
@@ -108,15 +146,58 @@ var proxyLogsColumns = []schemaColumn{
 	{"error", "String"},
 }

-func validateProxyLogsSchema(columns []schemaColumn, table schemaTable) error {
-	if len(columns) != len(proxyLogsColumns) {
-		return fmt.Errorf("got %d columns, want %d", len(columns), len(proxyLogsColumns))
+var proxyLogColumnDDL = map[string]string{
+	"request_id":         "request_id String",
+	"method":             "method String",
+	"path":               "path String",
+	"query":              "query String",
+	"client_ip":          "client_ip String",
+	"request_headers":    "request_headers String",
+	"request_body":       "request_body String",
+	"request_truncated":  "request_truncated Bool DEFAULT false",
+	"status_code":        "status_code Int32",
+	"response_headers":   "response_headers String",
+	"response_body":      "response_body String",
+	"response_truncated": "response_truncated Bool DEFAULT false",
+	"is_stream":          "is_stream Bool DEFAULT false",
+	"latency_ms":         "latency_ms Int64",
+	"started_at":         "started_at DateTime64(3)",
+	"finished_at":        "finished_at DateTime64(3)",
+	"error":              "error String",
+}
+
+func missingProxyLogColumns(columns []schemaColumn) ([]string, error) {
+	existing := make(map[string]string, len(columns))
+	for _, column := range columns {
+		existing[column.name] = column.typ
 	}
-	for i, want := range proxyLogsColumns {
-		if columns[i] != want {
-			return fmt.Errorf("column %d is %s %s, want %s %s", i+1, columns[i].name, columns[i].typ, want.name, want.typ)
+	var statements []string
+	for _, required := range proxyLogsColumns {
+		if typ, ok := existing[required.name]; ok {
+			if typ != required.typ {
+				return nil, fmt.Errorf("column %s has type %s, want %s", required.name, typ, required.typ)
+			}
+			continue
+		}
+		statements = append(statements, "ALTER TABLE proxy_logs ADD COLUMN IF NOT EXISTS "+proxyLogColumnDDL[required.name])
+	}
+	return statements, nil
+}
+
+func validateProxyLogsSchema(columns []schemaColumn, table schemaTable) error {
+	existing := make(map[string]string, len(columns))
+	for _, column := range columns {
+		existing[column.name] = column.typ
+	}
+	for _, want := range proxyLogsColumns {
+		if typ, ok := existing[want.name]; !ok || typ != want.typ {
+			return fmt.Errorf("column %s is %s, want %s", want.name, typ, want.typ)
 		}
 	}
+	return validateProxyLogsTable(table)
+}
+
+func validateProxyLogsTable(table schemaTable) error {
 	if table.engine != "MergeTree" {
 		return fmt.Errorf("engine is %q, want MergeTree", table.engine)
 	}