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

logger/queue.go

@@ -50,24 +50,21 @@ type poolBackend struct {
 	pool       *db.Pool
 	conn       driver.Conn
 	generation uint64
+	release    func()
 }
 
 func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, error) {
-	conn := b.conn
-	if conn == nil {
-		conn, _ = b.pool.GetWithGeneration()
-	}
-	if conn == nil {
+	if b.conn == nil {
 		return nil, errors.New("pool nil")
 	}
-	return clickHouseBatchPreparer{conn: conn}.PrepareBatch(ctx, query)
+	return clickHouseBatchPreparer{conn: b.conn}.PrepareBatch(ctx, query)
 }
 
 func (b poolBackend) Healthy() bool  { return b.pool.Healthy() }
 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
+	conn, generation, release := b.pool.Acquire()
+	return poolBackend{pool: b.pool, conn: conn, generation: generation, release: release}, generation
 }
 func (b poolBackend) MarkUnhealthyGeneration(generation uint64) {
 	b.pool.MarkUnhealthyGeneration(generation)
@@ -330,43 +327,77 @@ func (q *Queue) run(workCtx context.Context) {
 	ticker := time.NewTicker(q.batchInterval)
 	defer ticker.Stop()
 
-	flush := func() {
+	flush := func(final bool) bool {
 		if len(batch) == 0 {
-			return
+			return true
 		}
-		entries := make([]*LogEntry, len(batch))
-		for i, item := range batch {
-			entries[i] = item.entry
+		if q.backend == nil || !q.backend.Healthy() {
+			if !final {
+				return false
+			}
+			q.failed.Add(uint64(len(batch)))
+		} else {
+			q.flush(q.flushContext(workCtx), entriesFromBatch(batch))
 		}
-		q.flush(q.flushContext(workCtx), entries)
 		for _, item := range batch {
 			q.release(item.size)
 		}
 		clear(batch)
 		batch = batch[:0]
+		return true
 	}
 
 	for {
+		if len(batch) >= q.batchSize && (q.backend == nil || !q.backend.Healthy()) {
+			if q.isStopped() {
+				flush(true)
+				q.discardQueued()
+				return
+			}
+			select {
+			case <-ticker.C:
+				flush(false)
+			case <-workCtx.Done():
+				flush(true)
+				q.discardQueued()
+				return
+			}
+			continue
+		}
 		select {
 		case item, ok := <-q.ch:
 			if !ok {
-				flush()
+				flush(true)
 				return
 			}
 			batch = append(batch, item)
 			if len(batch) >= q.batchSize {
-				flush()
+				flush(false)
 			}
 		case <-ticker.C:
-			flush()
+			flush(false)
 		case <-workCtx.Done():
-			flush()
+			flush(true)
 			q.discardQueued()
 			return
 		}
 	}
 }
 
+func (q *Queue) isStopped() bool {
+	q.mu.RLock()
+	defer q.mu.RUnlock()
+	return q.stopped
+}
+
+func entriesFromBatch(batch []queuedEntry) []*LogEntry {
+	entries := make([]*LogEntry, len(batch))
+	for i, item := range batch {
+		entries[i] = item.entry
+	}
+	return entries
+}
+
 func (q *Queue) flushContext(workCtx context.Context) context.Context {
 	q.mu.RLock()
 	defer q.mu.RUnlock()
@@ -377,11 +408,6 @@ func (q *Queue) flushContext(workCtx context.Context) context.Context {
 }
 
 func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
-	if q.backend == nil || !q.backend.Healthy() {
-		q.failed.Add(uint64(len(entries)))
-		return
-	}
-
 	retry := entries
 	var lastErr error
 	var failedGeneration uint64
@@ -390,9 +416,12 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
 			lastErr = err
 			break
 		}
-		attemptBackend, generation := backendSnapshot(q.backend)
+		attemptBackend, generation, release := backendSnapshot(q.backend)
 		failedGeneration = generation
 		result := Flush(ctx, attemptBackend, retry)
+		if release != nil {
+			release()
+		}
 		q.failed.Add(uint64(result.Failed))
 		q.ambiguous.Add(uint64(result.Ambiguous))
 		if result.Ambiguous > 0 {
@@ -421,11 +450,21 @@ type generationBackend interface {
 	MarkUnhealthyGeneration(uint64)
 }
 
-func backendSnapshot(backend Backend) (Backend, uint64) {
+func backendSnapshot(backend Backend) (Backend, uint64, func()) {
 	if versioned, ok := backend.(generationBackend); ok {
-		return versioned.Snapshot()
+		snapshot, generation := versioned.Snapshot()
+		if leased, ok := snapshot.(interface{ Release() }); ok {
+			return snapshot, generation, leased.Release
+		}
+		return snapshot, generation, nil
+	}
+	return backend, 0, nil
+}
+
+func (b poolBackend) Release() {
+	if b.release != nil {
+		b.release()
 	}
-	return backend, 0
 }
 
 func markBackendUnhealthy(backend Backend, generation uint64) {