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

db/clickhouse.go

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