fix: restore reviewable migration evidence
This commit is contained in:
MiMoCode
+447
@@ -0,0 +1,447 @@
|
||||
package logger
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log"
|
||||
"math/rand"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
|
||||
|
||||
"git.misaka.ren/M1saka/token_thief/db"
|
||||
)
|
||||
|
||||
const (
|
||||
maxAttempts = 3
|
||||
BaseBackoff = 200 * time.Millisecond
|
||||
)
|
||||
|
||||
type Batch interface {
|
||||
Append(v ...any) error
|
||||
Send() error
|
||||
Abort() error
|
||||
}
|
||||
|
||||
type BatchPreparer interface {
|
||||
PrepareBatch(ctx context.Context, query string) (Batch, error)
|
||||
}
|
||||
|
||||
type clickHouseBatchPreparer struct {
|
||||
conn driver.Conn
|
||||
}
|
||||
|
||||
// The explicit adapter keeps the local Batch contract aligned with the real driver.
|
||||
func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string) (Batch, error) {
|
||||
return p.conn.PrepareBatch(ctx, query)
|
||||
}
|
||||
|
||||
type Stats struct {
|
||||
Enqueued uint64
|
||||
Dropped uint64
|
||||
Failed uint64
|
||||
Ambiguous uint64
|
||||
Bytes int64
|
||||
}
|
||||
|
||||
// FlushResult separates retry-safe Prepare failures from final or ambiguous failures.
|
||||
type FlushResult struct {
|
||||
Retry []*LogEntry
|
||||
Failed int
|
||||
Ambiguous int
|
||||
Err error
|
||||
}
|
||||
|
||||
// Queue is an asynchronous, bounded logger. Submit never waits for database work.
|
||||
type Queue struct {
|
||||
ch chan *LogEntry
|
||||
pool *db.Pool
|
||||
batchSize int
|
||||
batchInterval time.Duration
|
||||
workers int
|
||||
byteBudget int64
|
||||
|
||||
dropped atomic.Uint64
|
||||
failed atomic.Uint64
|
||||
ambiguous atomic.Uint64
|
||||
enq atomic.Uint64
|
||||
bytes atomic.Int64
|
||||
|
||||
mu sync.RWMutex
|
||||
started bool
|
||||
stopped bool
|
||||
shutdownCtx context.Context
|
||||
reporterCancel context.CancelFunc
|
||||
workCancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
done chan struct{}
|
||||
doneOnce sync.Once
|
||||
}
|
||||
|
||||
// NewQueue accepts an optional byte budget. A non-positive or omitted budget disables byte limiting.
|
||||
func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval time.Duration, byteBudget ...int64) *Queue {
|
||||
var budget int64
|
||||
if len(byteBudget) > 0 {
|
||||
budget = byteBudget[0]
|
||||
}
|
||||
if queueSize < 0 {
|
||||
queueSize = 0
|
||||
}
|
||||
if batchSize < 1 {
|
||||
batchSize = 1
|
||||
}
|
||||
if workers < 1 {
|
||||
workers = 1
|
||||
}
|
||||
if batchInterval <= 0 {
|
||||
batchInterval = time.Second
|
||||
}
|
||||
return &Queue{
|
||||
ch: make(chan *LogEntry, queueSize),
|
||||
pool: pool,
|
||||
batchSize: batchSize,
|
||||
batchInterval: batchInterval,
|
||||
workers: workers,
|
||||
byteBudget: budget,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// EstimatedBytes covers all variable-size strings and byte slices retained by an entry.
|
||||
func EstimatedBytes(e *LogEntry) int64 {
|
||||
if e == nil {
|
||||
return 0
|
||||
}
|
||||
return int64(len(e.RequestID) + len(e.Method) + len(e.Path) + len(e.Query) + len(e.ClientIP) + len(e.Error) +
|
||||
len(e.RequestHeaders) + len(e.RequestBody) + len(e.ResponseHeaders) + len(e.ResponseBody))
|
||||
}
|
||||
|
||||
func (q *Queue) Stats() Stats {
|
||||
return Stats{
|
||||
Enqueued: q.enq.Load(),
|
||||
Dropped: q.dropped.Load(),
|
||||
Failed: q.failed.Load(),
|
||||
Ambiguous: q.ambiguous.Load(),
|
||||
Bytes: q.bytes.Load(),
|
||||
}
|
||||
}
|
||||
|
||||
// Start starts workers once. The root context controls reporting only; Stop owns worker shutdown.
|
||||
func (q *Queue) Start(root context.Context) {
|
||||
q.mu.Lock()
|
||||
if q.started || q.stopped {
|
||||
q.mu.Unlock()
|
||||
return
|
||||
}
|
||||
q.started = true
|
||||
workCtx, workCancel := context.WithCancel(context.Background())
|
||||
reportCtx, reporterCancel := context.WithCancel(root)
|
||||
q.workCancel = workCancel
|
||||
q.reporterCancel = reporterCancel
|
||||
for i := 0; i < q.workers; i++ {
|
||||
q.wg.Add(1)
|
||||
go q.run(workCtx)
|
||||
}
|
||||
q.wg.Add(1)
|
||||
go q.reportLoop(reportCtx)
|
||||
q.mu.Unlock()
|
||||
}
|
||||
|
||||
// Stop closes submissions once and waits under the caller's single total deadline.
|
||||
// The variadic form permits legacy Stop() calls while new callers should pass a context.
|
||||
func (q *Queue) Stop(contexts ...context.Context) error {
|
||||
ctx := context.Background()
|
||||
if len(contexts) > 0 && contexts[0] != nil {
|
||||
ctx = contexts[0]
|
||||
}
|
||||
|
||||
q.mu.Lock()
|
||||
if !q.stopped {
|
||||
q.stopped = true
|
||||
q.shutdownCtx = ctx
|
||||
close(q.ch)
|
||||
if q.reporterCancel != nil {
|
||||
q.reporterCancel()
|
||||
}
|
||||
if q.started {
|
||||
go func() {
|
||||
q.wg.Wait()
|
||||
q.doneOnce.Do(func() { close(q.done) })
|
||||
}()
|
||||
} else {
|
||||
q.discardQueued()
|
||||
q.doneOnce.Do(func() { close(q.done) })
|
||||
}
|
||||
}
|
||||
done := q.done
|
||||
workCancel := q.workCancel
|
||||
q.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
if workCancel != nil {
|
||||
workCancel()
|
||||
}
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// Submit reserves memory and enqueues without waiting; full, over-budget, and stopped queues drop.
|
||||
func (q *Queue) Submit(e *LogEntry) {
|
||||
if e == nil {
|
||||
q.dropped.Add(1)
|
||||
return
|
||||
}
|
||||
size := EstimatedBytes(e)
|
||||
if !q.mu.TryRLock() {
|
||||
q.dropped.Add(1)
|
||||
return
|
||||
}
|
||||
if q.stopped || !q.reserve(size) {
|
||||
q.mu.RUnlock()
|
||||
q.dropped.Add(1)
|
||||
return
|
||||
}
|
||||
select {
|
||||
case q.ch <- e:
|
||||
q.enq.Add(1)
|
||||
default:
|
||||
q.release(size)
|
||||
q.dropped.Add(1)
|
||||
}
|
||||
q.mu.RUnlock()
|
||||
}
|
||||
|
||||
func (q *Queue) reserve(size int64) bool {
|
||||
if q.byteBudget <= 0 {
|
||||
q.bytes.Add(size)
|
||||
return true
|
||||
}
|
||||
for {
|
||||
used := q.bytes.Load()
|
||||
if size > q.byteBudget-used {
|
||||
return false
|
||||
}
|
||||
if q.bytes.CompareAndSwap(used, used+size) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (q *Queue) release(size int64) {
|
||||
q.bytes.Add(-size)
|
||||
}
|
||||
|
||||
func (q *Queue) discardQueued() {
|
||||
for e := range q.ch {
|
||||
q.release(EstimatedBytes(e))
|
||||
q.failed.Add(1)
|
||||
}
|
||||
}
|
||||
|
||||
func (q *Queue) run(workCtx context.Context) {
|
||||
defer q.wg.Done()
|
||||
batch := make([]*LogEntry, 0, q.batchSize)
|
||||
ticker := time.NewTicker(q.batchInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
flush := func() {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
q.flush(q.flushContext(workCtx), batch)
|
||||
for _, e := range batch {
|
||||
q.release(EstimatedBytes(e))
|
||||
}
|
||||
clear(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case e, ok := <-q.ch:
|
||||
if !ok {
|
||||
flush()
|
||||
return
|
||||
}
|
||||
batch = append(batch, e)
|
||||
if len(batch) >= q.batchSize {
|
||||
flush()
|
||||
}
|
||||
case <-ticker.C:
|
||||
flush()
|
||||
case <-workCtx.Done():
|
||||
flush()
|
||||
q.discardQueued()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (q *Queue) flushContext(workCtx context.Context) context.Context {
|
||||
q.mu.RLock()
|
||||
defer q.mu.RUnlock()
|
||||
if q.shutdownCtx != nil {
|
||||
return q.shutdownCtx
|
||||
}
|
||||
return workCtx
|
||||
}
|
||||
|
||||
func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
|
||||
if q.pool == nil || !q.pool.Healthy() {
|
||||
q.failed.Add(uint64(len(entries)))
|
||||
return
|
||||
}
|
||||
|
||||
retry := entries
|
||||
var lastErr error
|
||||
for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ {
|
||||
if err := ctx.Err(); err != nil {
|
||||
lastErr = err
|
||||
break
|
||||
}
|
||||
conn := q.pool.Get()
|
||||
if conn == nil {
|
||||
lastErr = errors.New("pool nil")
|
||||
break
|
||||
}
|
||||
result := Flush(ctx, clickHouseBatchPreparer{conn: conn}, retry)
|
||||
q.failed.Add(uint64(result.Failed))
|
||||
q.ambiguous.Add(uint64(result.Ambiguous))
|
||||
if result.Ambiguous > 0 {
|
||||
q.pool.MarkUnhealthy()
|
||||
}
|
||||
lastErr = result.Err
|
||||
retry = result.Retry
|
||||
if len(retry) == 0 {
|
||||
return
|
||||
}
|
||||
if attempt < maxAttempts && !waitBackoff(ctx, attempt) {
|
||||
lastErr = ctx.Err()
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if len(retry) > 0 {
|
||||
q.failed.Add(uint64(len(retry)))
|
||||
q.pool.MarkUnhealthy()
|
||||
log.Printf("[logger] giving up %d retry-safe rows after %d attempts: %v", len(retry), maxAttempts, lastErr)
|
||||
}
|
||||
}
|
||||
|
||||
func waitBackoff(ctx context.Context, attempt int) bool {
|
||||
wait := BaseBackoff
|
||||
for i := 1; i < attempt; i++ {
|
||||
wait *= 3
|
||||
}
|
||||
jitter := time.Duration((rand.Float64()*0.4 - 0.2) * float64(wait))
|
||||
timer := time.NewTimer(wait + jitter)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-timer.C:
|
||||
return true
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
const insertStatement = `INSERT INTO proxy_logs (
|
||||
request_id, method, path, query, client_ip,
|
||||
request_headers, request_body, request_truncated,
|
||||
status_code, response_headers, response_body, response_truncated,
|
||||
is_stream, latency_ms, started_at, finished_at, error
|
||||
)`
|
||||
|
||||
// Flush executes one batch. Only Prepare failures are retryable. Append failures are isolated by row;
|
||||
// Send failures are ambiguous and therefore never replayed.
|
||||
func Flush(ctx context.Context, conn BatchPreparer, entries []*LogEntry) FlushResult {
|
||||
if len(entries) == 0 {
|
||||
return FlushResult{}
|
||||
}
|
||||
result := flushOnce(ctx, conn, entries)
|
||||
if result.stage != flushAppend || len(entries) == 1 {
|
||||
return result.public(entries)
|
||||
}
|
||||
|
||||
combined := FlushResult{}
|
||||
for _, entry := range entries {
|
||||
single := flushOnce(ctx, conn, []*LogEntry{entry}).public([]*LogEntry{entry})
|
||||
combined.Retry = append(combined.Retry, single.Retry...)
|
||||
combined.Failed += single.Failed
|
||||
combined.Ambiguous += single.Ambiguous
|
||||
combined.Err = errors.Join(combined.Err, single.Err)
|
||||
}
|
||||
return combined
|
||||
}
|
||||
|
||||
type flushStage uint8
|
||||
|
||||
const (
|
||||
flushSuccess flushStage = iota
|
||||
flushPrepare
|
||||
flushAppend
|
||||
flushSend
|
||||
)
|
||||
|
||||
type flushAttempt struct {
|
||||
stage flushStage
|
||||
err error
|
||||
}
|
||||
|
||||
func (r flushAttempt) public(entries []*LogEntry) FlushResult {
|
||||
switch r.stage {
|
||||
case flushSuccess:
|
||||
return FlushResult{}
|
||||
case flushPrepare:
|
||||
return FlushResult{Retry: entries, Err: r.err}
|
||||
case flushSend:
|
||||
return FlushResult{Failed: len(entries), Ambiguous: len(entries), Err: r.err}
|
||||
default:
|
||||
return FlushResult{Failed: len(entries), Err: r.err}
|
||||
}
|
||||
}
|
||||
|
||||
func flushOnce(ctx context.Context, conn BatchPreparer, entries []*LogEntry) flushAttempt {
|
||||
batch, err := conn.PrepareBatch(ctx, insertStatement)
|
||||
if err != nil {
|
||||
return flushAttempt{stage: flushPrepare, err: err}
|
||||
}
|
||||
abort := func(cause error) error {
|
||||
return errors.Join(cause, batch.Abort())
|
||||
}
|
||||
for _, e := range entries {
|
||||
if err := batch.Append(
|
||||
e.RequestID, e.Method, e.Path, e.Query, e.ClientIP,
|
||||
string(e.RequestHeaders), string(e.RequestBody), e.RequestTruncated,
|
||||
int32(e.StatusCode), string(e.ResponseHeaders), string(e.ResponseBody), e.ResponseTruncated,
|
||||
e.IsStream, e.LatencyMS, e.StartedAt, e.FinishedAt, e.Error,
|
||||
); err != nil {
|
||||
return flushAttempt{stage: flushAppend, err: abort(err)}
|
||||
}
|
||||
}
|
||||
if err := batch.Send(); err != nil {
|
||||
return flushAttempt{stage: flushSend, err: abort(err)}
|
||||
}
|
||||
return flushAttempt{stage: flushSuccess}
|
||||
}
|
||||
|
||||
func (q *Queue) reportLoop(ctx context.Context) {
|
||||
defer q.wg.Done()
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
stats := q.Stats()
|
||||
healthy := q.pool != nil && q.pool.Healthy()
|
||||
log.Printf("[logger] metrics: enq=%d dropped_queue_full=%d dropped_db_fail=%d ambiguous_send=%d queue_len=%d queue_bytes=%d db_healthy=%v",
|
||||
stats.Enqueued, stats.Dropped, stats.Failed, stats.Ambiguous, len(q.ch), stats.Bytes, healthy)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user